# SPDX-License-Identifier: Apache-2.0
"""Immutable, typed stage checkpoints, displayed as native Markdown ports.

The Markdown is a view of a server-owned manifest, not editable configuration.
Only its opaque /artifact/<id> capability is consumed. Browser-supplied paths,
metadata, coordinate claims and HTML are never trusted as artifact contents.
"""
from __future__ import annotations

import hashlib
import json
import math
import os
import re
import threading
import uuid
from pathlib import Path
from workflow_diagnostics import StageTrace
from workflow_execution import cancellation_check, ensure_gpu_idle, ExecutionBusy

ROOT = Path(os.environ.get("APP_DATA_ROOT", "/tmp/viggle-mine-workflow")).resolve()
PUBLIC = ROOT / "outputs"
STATE = ROOT / "state"
for directory in (PUBLIC, STATE):
    directory.mkdir(parents=True, exist_ok=True)

# Also serializes V22's module globals, its progress bridge and CPU model state.
# One process/worker per Space, matching the canonical launch model.
EXECUTION_LOCK = threading.RLock()
POINTS_LOCK = threading.RLock()
ID = re.compile(r"^[0-9a-f]{32}$")
REF = re.compile(r"\[Inspect checkpoint\]\(/artifact/([0-9a-f]{32})\)")
SPACES = {
    "prepared": "normalized_source",
    "tracked": "normalized_source",
    "guarded": "normalized_source",
    "camera": "stabilized",
    "pass_a": "stabilized_generated",
    "recovery": "stabilized",
    "pass_b": "normalized_source",
    "standard": "normalized_source",
}


def atomic_json(path, value):
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp = path.with_name(path.name + "." + uuid.uuid4().hex + ".tmp")
    tmp.write_text(json.dumps(value, indent=2, allow_nan=False), encoding="utf-8")
    tmp.replace(path)


def digest(value):
    return hashlib.sha256(json.dumps(value, sort_keys=True, separators=(",", ":"),
                                     allow_nan=False).encode()).hexdigest()


def file_digest(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for block in iter(lambda: f.read(1024 * 1024), b""):
            h.update(block)
    return h.hexdigest()


def checkpoint_id(value):
    if not isinstance(value, str):
        raise ValueError("Use the connected checkpoint produced by the previous stage.")
    if ID.fullmatch(value):
        return value
    matches = REF.findall(value)
    if len(matches) != 1:
        raise ValueError("Checkpoint missing. Run the previous stage first.")
    return matches[0]


def load(value, expected=None):
    ident = checkpoint_id(value)
    path = STATE / "artifacts" / f"{ident}.json"
    if not path.is_file():
        raise ValueError("This checkpoint has expired after a Space restart. Prepare the source again.")
    artifact = json.loads(path.read_text())
    if artifact.get("schema") != "viggle.workflow.checkpoint.v1" or artifact.get("id") != ident:
        raise ValueError("Invalid checkpoint manifest.")
    kind = artifact.get("kind")
    if kind not in SPACES or artifact.get("coordinate_space") != SPACES[kind]:
        raise ValueError("Checkpoint coordinate-space contract is invalid.")
    if expected and kind not in ((expected,) if isinstance(expected, str) else expected):
        raise ValueError(f"Expected {expected}, received {kind}. Do not mix coordinate spaces.")
    for file in artifact.get("files", []):
        p = Path(file).resolve()
        if not p.is_relative_to(PUBLIC) or not p.is_file():
            raise ValueError("Checkpoint media has expired. Run its producing stage again.")
    return artifact


def create(kind, payload, files=(), parents=()):
    ident = uuid.uuid4().hex
    value = {"schema": "viggle.workflow.checkpoint.v1", "id": ident,
             "kind": kind, "coordinate_space": SPACES[kind],
             "parents": list(parents), "files": list(map(str, files)), **payload}
    atomic_json(STATE / "artifacts" / f"{ident}.json", value)
    return load(ident, kind)


def card(artifact, body):
    return f"{body}\n\n[Inspect checkpoint](/artifact/{artifact['id']})"


def cached(stage, inputs, compute, *, slot=None):
    """Replay completed work even after canvas refresh/force-run/API traversal.

    A generation take is immutable. Reusing its slot with changed inputs fails
    instead of silently spending another GPU allocation during Recovery.
    A failed GPU attempt is NOT cached; retrying the same take is permitted.
    """
    fingerprint = digest([stage, inputs])
    key = digest([stage, slot]) if slot is not None else fingerprint
    path = STATE / "cache" / f"{key}.json"
    check_cancelled = cancellation_check()
    check_cancelled()
    # Do not queue indefinitely behind a cancelled or stalled SDK call. V22
    # shares mutable CPU model/progress state across otherwise separate nodes.
    if not EXECUTION_LOCK.acquire(blocking=False):
        raise ExecutionBusy("Another V22 stage is running. Wait for it to finish before running this stage.")
    try:
        with StageTrace(stage) as trace:
            if path.is_file():
                entry = json.loads(path.read_text())
                if entry["fingerprint"] != fingerprint:
                    raise ValueError("This take already has different inputs. Increase Take to explicitly generate a new result; Recovery never regenerates an existing Pass A take.")
                for ident in entry["checkpoints"]:
                    load(ident)
                trace.set_phase("cache_hit")
                return tuple(entry["result"])
            ensure_gpu_idle()
            check_cancelled()
            trace.set_phase("calling_backend")
            result, checkpoints = compute()
            check_cancelled()
            trace.set_phase("saving_checkpoint")
            atomic_json(path, {"fingerprint": fingerprint, "result": result,
                               "checkpoints": checkpoints})
            return result
    finally:
        EXECUTION_LOCK.release()


def number(value, name, low, high, integer=False):
    x = float(value)
    if not math.isfinite(x) or not low <= x <= high or (integer and int(x) != x):
        raise ValueError(f"{name} must be {'an integer ' if integer else ''}between {low} and {high}.")
    return int(x) if integer else x


def boolean(value, name):
    if not isinstance(value, bool):
        raise ValueError(f"{name} requires a checkbox boolean, not text.")
    return value


def validate_points(value, prepared):
    meta = prepared["meta"]
    if not isinstance(value, dict) or value.get("coordinate_space") != "normalized_source":
        raise ValueError("Picker points must use normalized_source pixel coordinates.")
    if value.get("preparation_id") != prepared["id"]:
        raise ValueError("Points belong to another prepared source.")
    if (value.get("w"), value.get("h")) != (meta["w"], meta["h"]):
        raise ValueError("Point dimensions do not match this prepared source.")
    anchors = value.get("anchors")
    # The data contract is a list. Only this V22 adapter enforces three anchors.
    if not isinstance(anchors, list) or len(anchors) != len(meta["frames"]) or len(anchors) != 3:
        raise ValueError("V22 currently requires three prepared anchors.")
    clean = []
    for i, (a, frame, time) in enumerate(zip(anchors, meta["frames"], meta["times"])):
        if not isinstance(a, dict):
            raise ValueError(f"Anchor {i + 1} must be an object with frame, time and points.")
        actual_time = float(a.get("time", -1))
        if a.get("frame") != frame or not math.isfinite(actual_time) or abs(actual_time - time) > 1e-6:
            raise ValueError(f"Anchor {i + 1} timestamp/frame does not match preparation.")
        points = a.get("points")
        if not isinstance(points, list) or not 3 <= len(points) <= 12:
            raise ValueError(f"Anchor {i + 1}: select 3–12 visible points on the same subject.")
        pts = []
        for p in points:
            if not isinstance(p, (list, tuple)) or len(p) != 2:
                raise ValueError("Each point must be [x, y].")
            x, y = map(float, p)
            if not (math.isfinite(x) and math.isfinite(y) and 0 <= x < meta["w"] and 0 <= y < meta["h"]):
                raise ValueError("Point outside the normalized source image.")
            pts.append([x, y])
        clean.append({"frame": frame, "time": time, "points": pts})
    return {"coordinate_space": "normalized_source", "preparation_id": prepared["id"],
            "w": meta["w"], "h": meta["h"], "anchors": clean}


def save_points(preparation_id, value):
    prepared = load(preparation_id, "prepared")
    clean = validate_points(value, prepared)
    clean["revision"] = digest(clean)
    with POINTS_LOCK:
        atomic_json(STATE / "points" / f"{prepared['id']}.json", clean)
    return clean


def points(prepared):
    path = STATE / "points" / f"{prepared['id']}.json"
    with POINTS_LOCK:
        if not path.is_file():
            raise ValueError("Open Point Picker in the Source checkpoint, select 3–12 points on each anchor, and Save points. Then run Track Subject.")
        value = json.loads(path.read_text())
    clean = validate_points(value, prepared)
    if value.get("revision") != digest(clean):
        raise ValueError("Saved point revision is invalid.")
    return value


def assert_current_points(artifact):
    current = points(load(artifact["preparation_id"], "prepared"))
    if current["revision"] != artifact["points_revision"]:
        raise ValueError("Points changed in the picker. Run Track Subject again before building a new camera or Pass A.")