# SPDX-License-Identifier: Apache-2.0
"""Resume a long CAD build from the last section whose code and inputs haven't changed.

A build function is split into sections at its top-level `tick("name")` calls. After each section the runner
pickles the section's state (build123d shapes pickle exactly, as BREP) under a key: the hash of every input
(the module's code outside the build function, the extra files named by the caller, the spec, the library
versions), of every earlier section's code, and of this section's code. Change the code of one section and the
next run loads the checkpoint before it and runs from there; change an input and it runs from the start.

    run(main, module, inputs=[...], fresh=False, keep_last=1)

`keep_last` sections at the end (export, layout.json) always run: they write the files. Checkpoints live in
~/.cache/goldfish-cad/<module>/, one chain at a time (the previous chain is deleted when a new one is written).
A release build passes fresh=True and ignores the cache.
"""
from __future__ import annotations

import ast
import copy
import hashlib
import inspect
import pickle
import shutil
import textwrap
import time
import types
import zlib
from pathlib import Path

CACHE = Path.home() / ".cache" / "goldfish-cad"


def _sections(fn):
    """The body of fn as [(label, code)], split after each top-level tick("label") call."""
    src = textwrap.dedent(inspect.getsource(fn))
    body = ast.parse(src).body[0].body
    lines = src.splitlines()
    out, start, label = [], body[0].lineno, None
    for st in body:
        call = st.value if isinstance(st, ast.Expr) and isinstance(st.value, ast.Call) else None
        if call and getattr(call.func, "id", None) == "tick":
            out.append((call.args[0].value, textwrap.dedent("\n".join(lines[start - 1:st.end_lineno]))))
            start = st.end_lineno + 1
    out.append(("end", textwrap.dedent("\n".join(lines[start - 1:]))))
    return out


def _defs_only(code):
    """The function definitions in a section: re-run when a section is skipped, so later sections can call them."""
    tree = ast.parse(code)
    keep = [n for n in tree.body if isinstance(n, (ast.FunctionDef, ast.Import, ast.ImportFrom))
            or (isinstance(n, ast.Assign) and isinstance(n.value, ast.Lambda))]
    return ast.unparse(ast.Module(body=keep, type_ignores=[]))


def _bare(v):
    """Shapes rebuilt from their geometry alone: a boolean's topology history (and the parent links of a solid
    taken out of a compound) don't pickle and aren't needed. Recurses into lists, tuples and dicts."""
    if hasattr(v, "wrapped") and hasattr(v, "_history"):
        try:
            return type(v)(v.wrapped)
        except Exception:
            w = copy.copy(v)
            w._history = None
            return w
    if isinstance(v, (list, tuple)):
        return type(v)(_bare(x) for x in v)
    if isinstance(v, dict):
        return {k: _bare(x) for k, x in v.items()}
    return v


def _picklable(ns, base):
    keep = {}
    for k, v in ns.items():
        if k.startswith("__") or k in base and base[k] is v:
            continue
        if isinstance(v, (types.FunctionType, types.ModuleType, type)) or callable(v) and not hasattr(v, "wrapped"):
            continue
        v = _bare(v)
        try:
            pickle.dumps(v)
        except Exception as e:
            print(f"    (not checkpointed: {k}, {type(v).__name__}: {e})")
            continue
        keep[k] = v
    return keep


def run(fn, module, inputs=(), fresh=False, keep_last=1):
    import build123d
    import OCP
    mod_src = inspect.getsource(module)
    fn_src = inspect.getsource(fn)
    h = hashlib.sha256()
    h.update(mod_src.replace(fn_src, "").encode())
    for p in inputs:
        h.update(Path(p).read_bytes())
    h.update(f"{build123d.__version__} {getattr(OCP, '__version__', '')}".encode())
    secs = _sections(fn)
    keys, k = [], h.hexdigest()
    for label, code in secs:
        k = hashlib.sha256((k + code).encode()).hexdigest()
        keys.append(k)
    cdir = CACHE / Path(module.__file__).stem
    cdir.mkdir(parents=True, exist_ok=True)
    cacheable = len(secs) - keep_last

    start = 0
    ns = dict(vars(module))
    base = dict(ns)
    if not fresh:
        for i in range(cacheable - 1, -1, -1):
            f = cdir / f"{i:02d}-{keys[i][:16]}.pkl"
            if f.exists():
                for j in range(i + 1):                      # functions defined in the skipped sections
                    exec(compile(_defs_only(secs[j][1]), f"<{secs[j][0]}>", "exec"), ns)
                ns.update(pickle.loads(zlib.decompress(f.read_bytes())))
                start = i + 1
                print(f"  resumed after '{secs[i][0]}' (sections 1–{i + 1} of {len(secs)} unchanged)", flush=True)
                break
    if "_t" in ns:
        ns["_t"] = [time.time()]
    for i in range(start, len(secs)):
        label, code = secs[i]
        exec(compile(code, f"<{label}>", "exec"), ns)
        if i < cacheable:
            for old in cdir.glob(f"{i:02d}-*.pkl"):         # one chain at a time
                old.unlink()
            (cdir / f"{i:02d}-{keys[i][:16]}.pkl").write_bytes(zlib.compress(pickle.dumps(_picklable(ns, base)), 1))


def clear(module_name=None):
    shutil.rmtree(CACHE / module_name if module_name else CACHE, ignore_errors=True)