#!/usr/bin/env python3
#
#    JKS - Measurement database system
#    Copyright (C) 2026  Christoph Lehner (christoph.lehner@ur.de, https://github.com/lehner/jks)
#
#    This program is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 2 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License along
#    with this program; if not, write to the Free Software Foundation, Inc.,
#    51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA.
#
import jks, sys, os, asyncio
import jks.flow.core as flow

usage = """%s command flow.sh [...]

  status flow.sh                     state of every node
  run flow.sh [-j N] [-f] [id ...]   compute the nodes (default: all) and what they need;
                                     -j: parallel steps, -f: recompute the given nodes
  export flow.sh id out.jks          write the database of a node
  log flow.sh id                     output of the step that computed a node
  source flow.sh id path             add a source database
  list flow.sh id word ...           add a list of words for loops ($(jks_list_values id))
  plot flow.sh id input out.pdf cmd ...  add a jks_plot2 figure of node input (out.pdf: - keeps
                                     it in the cache); cmd as for jks_plot2, e.g. c1:C:title
  add flow.sh id parent script args  add a step; parent - for scripts that write a new
                                     database; @ is the node's database, @id another node's
  block flow.sh id parent|- [-i id] [-f file] script.sh|-
                                     add a bash block: it changes $DB (a copy of parent), reads
                                     $IN_<id> of each -i node and the -f files; script.sh|-: stdin
  rm flow.sh id                      remove a node nobody reads
  rebase flow.sh id parent           let node id (and all after it) start from another node
  fmt flow.sh                        rewrite the file in canonical form
  du flow.sh                         disk usage of the work directory and of every node
  import flow.sh mk                  a new flow from an mk driver (a bash script of jks_* calls);
                                     what cannot be mapped is reported and kept as comments
  gc flow.sh [-n]                    delete stored results no node needs (-n: only report)

The flow file is a bash script; "bash flow.sh" replays it without the cache.  Loops
(jks_begin/for/jks_apply/done) are written into the file directly.""" % sys.argv[0]

if len(sys.argv) < 3 or sys.argv[1] not in ("status", "run", "export", "log", "source", "list", "plot", "add", "rm", "rebase", "fmt", "du", "gc", "import", "block"):
    print(usage)
    sys.exit(0)

cmd, fn, args = sys.argv[1], sys.argv[2], sys.argv[3:]


def fail(msg):
    print("ERROR: %s" % msg)
    sys.exit(1)


def size(n):
    for u in ["B", "KB", "MB", "GB"]:
        if n < 1024 or u == "GB":
            return "%.0f %s" % (n, u) if u == "B" else "%.1f %s" % (n, u)
        n /= 1024.0


try:
    if cmd == "import":
        from jks.flow import mk
        if len(args) != 1:
            fail("usage: jks_flow import flow.sh mk")
        if os.path.exists(fn):
            fail("%s exists already" % fn)
        fl, report, files, lists = mk.import_mk(args[0], fn)
        fl.save()
        for line, text in report:
            print("%s:%d: %s" % (args[0], line, text))
        steps = sum(1 for n in fl.nodes.values() if n.kind != "source")
        print("%s: %d nodes (%d sources%s)" % (fn, len(fl.nodes), len(fl.nodes) - steps,
                                            ", lists " + " ".join(lists) if lists else ""))
        for p, i in sorted(files.items()):
            print("  %s is node %s" % (p, i))
        print("compute with: jks_flow run %s" % fn)
        sys.exit(0)

    if cmd in ("source", "list", "plot", "add", "block") and not os.path.exists(fn):
        fl = flow.flow(fn)
    else:
        fl = flow.flow.load(fn)
    en = flow.engine(fl)

    if cmd == "status":
        for r in en.status():
            n = fl.nodes[r["id"]]
            if n.kind == "source":
                what = "source %s" % n.source
            elif n.kind == "list":
                what = "list " + " ".join(n.values)
            elif n.kind == "plot":
                what = "plot %s -> %s (%d commands)" % (n.plot["script"], n.plot["out"], len(n.plot["cmds"]))
            elif n.kind == "step":
                env, words = n.cmd.words()
                what = " ".join(env + words)
            elif n.kind == "block":
                first = next((l.strip() for l in n.block["script"].split("\n")
                              if l.strip() and not l.strip().startswith("#")), "")
                what = "block%s: %s" % (" (inputs %s)" % " ".join(n.block["inputs"]) if n.block["inputs"] else "", first)
            else:
                what = "loop over %s: %s (%s)" % (", ".join(n.loop_vars()), ", ".join(c.name for c in n.body), r["detail"])
            print("%-16s %-8s %-22s %s" % (r["id"], r["status"], r["reason"], what[:100]))

    elif cmd == "run":
        jobs, force, ids = 1, False, []
        i = 0
        while i < len(args):
            if args[i] == "-j":
                jobs = int(args[i + 1])
                i += 2
                continue
            if args[i] == "-f":
                force = True
            else:
                ids.append(args[i])
            i += 1
        for i in ids:
            if i not in fl.nodes:
                fail("no node %s" % i)
        verbose = os.environ.get("JKS_FLOW_VERBOSE")
        res = asyncio.run(en.run(ids or None, jobs=jobs, force=force,
                                 log=lambda i, line: print("  %s| %s" % (i, line)) if verbose else None))
        bad = 0
        for i in fl.nodes:
            if i in res:
                r = res[i]
                if r.get("delivered"):
                    print("%-16s copied the stored figure to %s" % (i, fl.nodes[i].plot["out"]))
                elif r["ok"]:
                    extra = ", %s over %d iterations, %d script run%s" % (
                        r["mode"], r["iterations"], r["computed"], "" if r["computed"] == 1 else "s") if "mode" in r else ""
                    print("%-16s done in %.2fs, stored %s (%s)%s" % (i, r["seconds"], r["stored"], size(r["bytes"]), extra))
                else:
                    bad += 1
                    print("%-16s FAILED: %s" % (i, r["error"]))
                    for line in r.get("log", "").splitlines()[-5:]:
                        print("  %s| %s" % (i, line))
        if not res:
            print("nothing to do")
        if bad:
            sys.exit(1)

    elif cmd == "export":
        if len(args) != 2:
            fail("export flow.sh id out.jks")
        en.export(args[0], args[1])

    elif cmd == "log":
        if len(args) != 1 or args[0] not in fl.nodes:
            fail("log flow.sh id")
        k = en.keys()[args[0]]
        st = en._read_json("state.json")
        f = st.get("failed", {}).get(args[0], {})
        m = en.meta(k) if k else None
        if m is None and f.get("key") == k:
            print("# failed: %s" % f["error"])
            print(f.get("log", ""))
        else:
            m = m or en.meta(st.get("built", {}).get(args[0], "") or "")
            if m is None:
                fail("%s has no stored result" % args[0])
            print("# %s  (%s)" % (m["command"], m["created"]))
            print(m.get("log", ""))

    elif cmd == "source":
        if len(args) != 2:
            fail("source flow.sh id path")
        fl.add(flow.node(args[0], "source", source=args[1]))
        fl.save()

    elif cmd == "plot":
        if len(args) < 4:
            fail("plot flow.sh id input out.pdf cmd ...")
        fl.add(flow.node(args[0], "plot", args[1], plot={"script": "jks_plot2", "out": args[2], "cmds": args[3:]}))
        fl.save()

    elif cmd == "list":
        if len(args) < 1:
            fail("list flow.sh id word ...")
        fl.add(flow.node(args[0], "list", values=args[1:]))
        fl.save()

    elif cmd == "add":
        if len(args) < 4:
            fail("add flow.sh id parent script args")
        env = {}
        while args and "=" in args[0] and args[0].split("=")[0].isidentifier():
            k, v = args.pop(0).split("=", 1)
            env[k] = v
        fl.add(flow.node(args[0], "step", None if args[1] == "-" else args[1],
                         flow.command(os.path.basename(args[2]), args[3:], env)))
        fl.save()

    elif cmd == "block":
        if len(args) < 3:
            fail("block flow.sh id parent|- [-i id] [-f file] script.sh|-")
        id, parent, rest, inputs, files = args[0], args[1], args[2:], [], []
        while len(rest) > 1 and rest[0] in ("-i", "-f"):
            (inputs if rest[0] == "-i" else files).append(rest[1])
            rest = rest[2:]
        if len(rest) != 1:
            fail("block flow.sh id parent|- [-i id] [-f file] script.sh|-")
        script = sys.stdin.read() if rest[0] == "-" else open(rest[0]).read()
        fl.add(flow.node(id, "block", None if parent == "-" else parent, meta={"files": files} if files else None,
                         block={"script": script if script.endswith("\n") else script + "\n", "inputs": inputs}))
        fl.save()

    elif cmd == "rm":
        if len(args) != 1:
            fail("rm flow.sh id")
        fl.remove(args[0])
        fl.save()

    elif cmd == "rebase":
        if len(args) != 2:
            fail("rebase flow.sh id parent")
        fl.rebase(args[0], args[1])
        fl.save()
        print("%s now starts from %s; recompute with: jks_flow run %s" % (args[0], args[1], fn))

    elif cmd == "fmt":
        fl.save()

    elif cmd == "du":
        u = en.usage()
        what = {"delta": "delta", "full": "full copy", "plot": "figure", "snapshot": "snapshot of the source"}
        for id, x in u["nodes"].items():
            c, b = x.get("current"), x.get("built")
            line = "%-16s %9s" % (id, size(c["bytes"] + c["iterations"]) if c else "-")
            if c:
                line += "  %s" % what.get(c["stored"], c["stored"])
                if c["stored"] == "delta" and c["changed"]:
                    line += " of %d tags" % c["changed"]
                if c["iterations"]:
                    line += ", %s in iterations" % size(c["iterations"])
            if b:
                line += "  (previous result %s)" % size(b["bytes"] + b["iterations"])
            print(line)
        print("")
        print("%s in %s: %s" % (size(u["total"]), u["work"], ", ".join(
            "%s %s" % (size(u["parts"][k]), k) for k in ("store", "tmp", "files", "other") if u["parts"][k])))
        print("jks_flow gc would free %s (%d stored results)" % (size(u["free"]), u["free_count"]))

    elif cmd == "gc":
        dry = "-n" in args
        n, count = en.gc(dry=dry)
        print("%s %s (%d stored results)" % ("would free" if dry else "freed", size(n), count))

except flow.FlowError as e:
    fail(e)
