"""Cloud-lane ledger: pure logic (spec docs/specs cloud-lane-design §4.4). IO: wf_cloud.py."""
from __future__ import annotations

import base64
import binascii
import dataclasses
import datetime as dt
import gzip
import hashlib
import json
import re

from . import tasks as TK, usage as U

OVERHEAD = 1.15          # title generation, setup: what the session can't see
LOST_AFTER = dt.timedelta(hours=24)
ENDED = ("done", "awaiting", "handback", "lost")


class CloudError(Exception):
    pass


def new() -> dict:
    return {"budget": 240.0, "spent": 0.0, "reserve_per_task": 4.0, "max_parallel": 3, "entries": []}


def loads(text: str) -> dict:
    if not text.strip():
        return new()
    try:
        led = json.loads(text)
        led["budget"], led["spent"], led["entries"]
    except (ValueError, KeyError, TypeError) as e:
        raise CloudError(f"cloud ledger unreadable: {e}")
    return {**new(), **led}


def dumps(led: dict) -> str:
    return json.dumps(led, indent=1) + "\n"


def running(led: dict) -> int:
    return sum(e["state"] == "running" for e in led["entries"])


def balance(led: dict) -> float:
    return led["budget"] - led["spent"] - led["reserve_per_task"] * running(led)


def refusal(led: dict) -> str | None:
    """One line why `send` must refuse (exit 3), or None."""
    bal, res = balance(led), led["reserve_per_task"]
    if bal < res:
        return f"cloud busy: balance ${bal:.2f} < reserve ${res:.2f}"
    if running(led) >= led["max_parallel"]:
        return f"cloud busy: {running(led)} running >= max_parallel {led['max_parallel']}"
    return None


def add(led: dict, id: str, project: str, sid: str, model: str, now: dt.datetime) -> dict:
    e = {"id": id, "project": project, "sid": sid, "model": model, "sent": now.isoformat(),
         "state": "running", "usd": None, "usd_source": None}
    led["entries"].append(e)
    return e


def find(led: dict, sid: str) -> dict:
    for e in led["entries"]:
        if e["sid"] == sid:
            return e
    raise CloudError(f"no cloud entry with sid {sid}")


def charge_usd(model: str, u: U.Usage | None, reserve: float) -> tuple[float, str]:
    """(usd, source): WF-USAGE x PRICES + 15%; missing/unknown model -> reserve, 'est'."""
    c = U.cost(model, u) if u is not None else None
    return (c * OVERHEAD, "self") if c is not None else (reserve, "est")


def end(led: dict, sid: str, state: str, u: U.Usage | None = None) -> dict:
    """Running entry ends: state, charge added to spent."""
    if state not in ENDED:
        raise CloudError(f"bad end state {state}")
    e = find(led, sid)
    if e["state"] != "running":
        raise CloudError(f"entry {sid} is {e['state']}, not running")
    if state == "lost":
        usd, src = led["reserve_per_task"], "est"
    else:
        usd, src = charge_usd(e["model"], u, led["reserve_per_task"])
    e.update(state=state, usd=usd, usd_source=src)
    led["spent"] += usd
    return e


def expire(led: dict, now: dt.datetime) -> list[dict]:
    """Running > 24h (no WF-RESULT seen) -> lost, charged reserve. Returns the lost entries."""
    out = []
    for e in led["entries"]:
        if e["state"] == "running" and now - dt.datetime.fromisoformat(e["sent"]) > LOST_AFTER:
            out.append(end(led, e["sid"], "lost"))
    return out


def set_balance(led: dict, bal: float, now: dt.datetime) -> dict:
    """Owner reconcile (balance read from claude.ai): spent = budget - balance."""
    old = led["spent"]
    led["spent"] = led["budget"] - bal
    row = {"id": "-", "project": "-", "sid": "-", "model": "-", "sent": now.isoformat(), "state": "done",
           "usd": led["spent"] - old, "usd_source": "owner"}
    led["entries"].append(row)
    return row


def summary(led: dict) -> str:
    return (f"budget ${led['budget']:.2f}  spent ${led['spent']:.2f}  balance ${balance(led):.2f}  "
            f"running {running(led)}/{led['max_parallel']}  (reserve ${led['reserve_per_task']:.2f}/task)")



# --- archive ended sessions (spec §4.3 Archive): the CLI's internal archiveRemoteSession, CLI 2.1.291 ---
API = "https://api.anthropic.com"
_ARCH_SID = re.compile(r"session_[A-Za-z0-9]+")


def archive_request(sid: str, creds: str, version: str) -> tuple[str, dict]:
    """(url, headers) for POST .../archive; creds = the CLI's .credentials.json text. Never logged."""
    if not _ARCH_SID.fullmatch(sid):
        raise CloudError(f"{sid!r} is not a cloud session id")
    try:
        d = json.loads(creds)
        tok = d["claudeAiOauth"]["accessToken"]
    except (ValueError, KeyError, TypeError):
        tok = None
    if not tok:
        raise CloudError("no claude.ai login token (run claude once)")
    h = {"Authorization": f"Bearer {tok}", "Content-Type": "application/json",
         "anthropic-version": "2023-06-01", "User-Agent": f"claude-code/{version}"}
    if d.get("trustedDeviceToken"):
        h["X-Trusted-Device-Token"] = d["trustedDeviceToken"]
    return f"{API}/v1/code/sessions/{sid}/archive", h


def archive_problem(status: int, body: str) -> str | None:
    """None = archived (200, or 409 = already); else one short why."""
    if status in (200, 409):
        return None
    if status == 401:
        return "HTTP 401 (login expired? run claude once)"
    return f"HTTP {status}: {' '.join(body.split())}"[:120]


def unarchived(led: dict) -> list[str]:
    """sids of ended sessions not archived yet (owner reconcile rows and pending sids skipped)."""
    return [e["sid"] for e in led["entries"]
            if e["state"] in ENDED and _ARCH_SID.fullmatch(e["sid"]) and not e.get("archived")]

# --- pty driver helpers (spec §3 F1-F3, F7, F8) ---------------------------------------------------
import re  # noqa: E402

_ANSI = re.compile(r"\x1b(\[[0-?]*[ -/]*[@-~]|\][^\x07\x1b]*(\x07|\x1b\\)|[PX^_][^\x1b]*\x1b\\|[@-Z\\-_])")
_SID = re.compile(r"(session_[A-Za-z0-9]{8,})")


def strip_ansi(text: str) -> str:
    """Terminal output -> plain text: cursor-forward (CSI n C) -> n spaces, cursor-to-column (CSI n G)
    -> one space (the TUI draws word gaps that way), other escape sequences dropped, CR -> LF."""
    text = re.sub(r"\x1b\[(\d*)C", lambda m: " " * int(m.group(1) or 1), text)
    text = re.sub(r"\x1b\[\d*G", " ", text)
    return _ANSI.sub("", text).replace("\r\n", "\n").replace("\r", "\n")


def parse_sid(text: str) -> str | None:
    """Session id from `claude --cloud` output: 'Resume with: claude --teleport <sid>', else the View URL."""
    plain = strip_ansi(text)
    for pat in (r"Resume with:\s*claude --teleport\s+" + _SID.pattern, r"View:\s*\S*?/code/" + _SID.pattern):
        m = re.findall(pat, plain)
        if m:
            return m[-1]
    return None


def last_lines(text: str, n: int = 5) -> list[str]:
    return [ln.strip() for ln in strip_ansi(text).splitlines() if ln.strip()][-n:]


def trust(claude_json: str, path: str) -> str:
    """~/.claude.json text with projects[path].hasTrustDialogAccepted = true (F3). Other keys kept."""
    try:
        data = json.loads(claude_json) if claude_json.strip() else {}
        projects = data.setdefault("projects", {})
        if not isinstance(projects, dict):
            raise TypeError("projects is not an object")
    except (ValueError, TypeError, AttributeError) as e:
        raise CloudError(f"~/.claude.json unreadable: {e}")
    projects.setdefault(path, {})["hasTrustDialogAccepted"] = True
    return json.dumps(data, indent=2) + "\n"


TRUST_PROMPT = re.compile(r"trust\s*(the\s*files\s*in\s*)?this\s*folder|Do\s*you\s*trust", re.I)
RESUMED = re.compile(r"Session\s*resumed", re.I)     # teleport reached the prompt (F7)


# --- send: snapshot + prompt (spec §4.1, §4.2) -----------------------------------------------------

MODEL = "claude-opus-5-5"     # what cloud sessions run (F4): ledger rows price with it
CAP = 90_000_000              # bytes: packed snapshot limit (upload ≤ 100 MB, F1)
NEVER = (".wf", ".worktrees", "out")     # + TASKS / archive: never in a snapshot (cloud_include re-adds)


def include_problem(path: str) -> str | None:
    """Why a cloud_include entry is unusable (absolute, leaves the project, empty), None = fine."""
    parts = path.replace("\\", "/").strip("/").split("/")
    if not path.strip() or path.startswith("/") or ".." in parts or parts in (["."], [".git"]) or parts[0] == ".git":
        return f"cloud_include '{path}': must be a path inside the project"
    return None


def size_refusal(size: int, cap: int = CAP) -> str | None:
    if size > cap:
        return f"snapshot {size / 1e6:.0f} MB > {cap / 1e6:.0f} MB; trim cloud_include"
    return None


def fill(template: str, id: str, base: str, task: str, recipe: str, note: str | None) -> str:
    """templates/cloud-prompt.md with its {{…}} fields filled."""
    values = {"id": id, "base": base, "task": task.strip(), "recipe": recipe.strip() or "(none: see CLAUDE.md)",
              "note": f"\n{note.strip()}\n" if note and note.strip() else ""}
    unknown = sorted(set(re.findall(r"\{\{(\w+)\}\}", template)) - set(values))
    if unknown:
        raise CloudError(f"cloud prompt template: unknown field {{{{{unknown[0]}}}}}")
    out = re.sub(r"\{\{(\w+)\}\}", lambda m: values[m.group(1)], template)   # one pass: task text kept as is
    return out


# The export renders the prompt as '❯ first line' + '  ' continuations and every assistant message as
# '● first line' + '  ' continuations (F12): only an assistant message can start with '● WF-RESULT'.
RESULT = re.compile(r"^● WF-RESULT\b")


def final_message(export: str) -> list[str] | None:
    """Lines of the last assistant message that starts with WF-RESULT (bullet / indent and trailing blanks
    removed), None = no result yet. The prompt never matches, whatever it contains."""
    lines = export.replace("\r\n", "\n").split("\n")
    starts = [i for i, ln in enumerate(lines) if RESULT.match(ln)]
    if not starts:
        return None
    out = []
    for ln in lines[starts[-1]:]:
        if out and ln[:2] not in ("  ", ""):
            break                                   # next message / prompt
        out.append(ln[2:].rstrip())
    while out and not out[-1]:
        out.pop()
    return out


def keyless_message(export: str) -> list[str] | None:
    """No WF-RESULT anywhere, but the last assistant message ends in a bare WF-PATCH-END line (the session
    dropped the key words): its lines (as final_message), else None. A non-empty prompt after it (a redo
    sent) -> None: still running."""
    lines = export.replace("\r\n", "\n").split("\n")
    blocks: list[tuple[str, list[str]]] = []
    for ln in lines:
        if ln[:2] in ("  ", "") and blocks:
            blocks[-1][1].append(ln[2:].rstrip())
        elif ln.strip():
            blocks.append((ln[:1], [ln[2:].rstrip()]))
    found = None
    for kind, body in blocks:
        if kind == "❯" and any(b.strip() for b in body):
            found = None                            # a later prompt: the message before it is answered
        elif kind == "●":
            text = [b for b in body if b.strip()]
            if text and text[-1].strip() == "WF-PATCH-END":
                found = body
    if found is None:
        return None
    out = list(found)
    while out and not out[-1]:
        out.pop()
    return out


def keyless_result(lines: list[str]) -> Result:
    """keyless_message() lines -> a handback Result: usage + patch (sha256=HEX bytes=N header anywhere before
    WF-PATCH-END, base64 after it) when found, so the patch can be kept."""
    text = " ".join(ln.strip() for ln in lines[:-1])
    r = Result("handback", "(no WF-RESULT key)")
    if m := _USAGE.search(text):
        r.usage = U.Usage(turns=1, inp=int(m[1]), cw5=int(m[2]), cr=int(m[3]), out=int(m[4]))
        r.model = m[5]
    blob = "".join(text.split())
    if m := _HEADER.search(blob):
        r.sha, r.nbytes, r.b64, r.has_patch = m[1].lower(), int(m[2]), blob[m.end():], True
    return r


# --- pull: result + patch (spec §4.3) ---------------------------------------------------------------

KEYS = ("WF-RESULT", "WF-REPORT", "WF-USAGE", "WF-PATCH-BEGIN", "WF-PATCH-END")
STATES = ("done", "awaiting", "handback")
_USAGE = re.compile(r"in=(\d+)\s+cw=(\d+)\s+cr=(\d+)\s+out=(\d+)(?:\s+model=(\S+))?")
# whitespace removed: the export may wrap the header anywhere; gzip base64 starts 'H4sI', never a digit
_HEADER = re.compile(r"sha256=([0-9a-fA-F]{64})bytes=(\d+)")


@dataclasses.dataclass
class Result:
    state: str
    report: str
    usage: U.Usage | None = None
    model: str | None = None
    sha: str | None = None
    nbytes: int | None = None
    b64: str = ""
    has_patch: bool = False


def parse_result(lines: list[str]) -> Result:
    """final_message() lines -> Result. A line starting with a key opens its section, other lines continue
    the open one (the export wraps long lines). CloudError: bad state / missing patch markers."""
    sec: dict[str, list[str]] = {}
    cur = None
    for ln in lines:
        s = ln.strip()
        key = next((k for k in KEYS if s == k or s.startswith(k + " ")), None)
        if key:
            cur = key
            sec.setdefault(key, []).append(s[len(key):].strip())
        elif cur:
            sec[cur].append(s)
    state = " ".join(sec.get("WF-RESULT", [""])).strip().lower()
    if state not in STATES:
        raise CloudError(f"bad WF-RESULT '{state}' (want {'/'.join(STATES)})")
    r = Result(state, " ".join(" ".join(sec.get("WF-REPORT", [])).split()) or "(no report)")
    if m := _USAGE.search(" ".join(sec.get("WF-USAGE", []))):
        r.usage = U.Usage(turns=1, inp=int(m[1]), cw5=int(m[2]), cr=int(m[3]), out=int(m[4]))
        r.model = m[5]
    if "WF-PATCH-BEGIN" in sec:
        if "WF-PATCH-END" not in sec:
            raise CloudError("WF-PATCH-BEGIN without WF-PATCH-END (truncated message)")
        blob = "".join("".join(sec["WF-PATCH-BEGIN"]).split())
        m = _HEADER.match(blob)
        if not m:
            raise CloudError("WF-PATCH-BEGIN needs sha256=HEX bytes=N")
        r.sha, r.nbytes, r.b64, r.has_patch = m[1].lower(), int(m[2]), blob[m.end():], True
    return r


def decode_patch(r: Result) -> bytes:
    """base64 -> check bytes + sha256 -> gunzip. CloudError on any mismatch."""
    if not r.has_patch:
        raise CloudError("no WF-PATCH-BEGIN/END in the result")
    try:
        gz = base64.b64decode(r.b64, validate=True)
    except (binascii.Error, ValueError) as e:
        raise CloudError(f"patch base64 broken: {e}")
    if len(gz) != r.nbytes:
        raise CloudError(f"patch bytes {len(gz)} != {r.nbytes} announced")
    if hashlib.sha256(gz).hexdigest() != r.sha:
        raise CloudError("patch sha256 mismatch")
    try:
        return gzip.decompress(gz)
    except (OSError, EOFError) as e:
        raise CloudError(f"patch gzip broken: {e}")


_DIFF = re.compile(r'^diff --git "?a/(.+?)"? "?b/(.+?)"?$')
_MOVE = re.compile(r"^(?:rename|copy) (?:from|to) (.+)$")


def patch_files(patch: str) -> list[str]:
    """Repo-relative paths a format-patch touches (both sides of renames/copies), in order, unique."""
    out = []
    for ln in patch.split("\n"):
        if m := _DIFF.match(ln):
            out += [m[1], m[2]]
        elif m := _MOVE.match(ln):
            out.append(m[1].strip('"'))
    return list(dict.fromkeys(out))


def patch_problem(files: list[str], project: str, never: list[str]) -> str | None:
    """Why a patch may not be applied, None = fine. files: repo-relative; project: the project folder
    relative to the repo ('' = repo root); never: project-relative paths the cloud may not touch
    (TASKS, archive, .wf, .worktrees, out, cloud_include)."""
    pre = project.strip("/")
    for f in files:
        parts = f.split("/")
        if not f or f.startswith("/") or ".." in parts or parts[0] == ".git":
            return f"patch touches {f}: outside the tree"
        if pre and not (f + "/").startswith(pre + "/"):
            return f"patch touches {f}: outside the project folder {pre}"
        rel = f[len(pre) + 1:] if pre else f
        for n in never:
            n = n.strip("/")
            if rel == n or rel.startswith(n + "/"):
                return f"patch touches {f}: {n} is never sent (wf bookkeeping / out / cloud_include)"
    return None


# ------------------------------------------------------------------ fit (spec §4.5)

# (reason, regex) over title + body; `Cloud: yes` skips them (Model opus still required). wf res before the generic wf-step rule.
UNFIT = (
    ("names `wf res` (local resource ledger)", re.compile(r"\bwf res\b")),
    ("names `wf` commands as steps", re.compile(r"\bwf (?!res\b)[a-z][a-z-]*")),
    ("needs a GUI", re.compile(r"\bGUI\b|\bscreenshot", re.I)),
    ("needs a LAN host", re.compile(r"\bLAN\b|\b192\.168\.\d+\.\d+|\b10\.\d+\.\d+\.\d+|\b[\w-]+\.local\b|\bssh [\w@.-]+")),
    ("needs a live server", re.compile(r"live server|running server|\blocalhost\b|127\.0\.0\.1", re.I)),
)


def fit(item: TK.Item, cfg) -> str | None:
    """None when the task fits the cloud lane, else the reason it does not (spec §4.5)."""
    if not cfg.cloud:
        return "project not opted in (cloud = true)"
    if item.cloud == "no":
        return "Cloud: no"
    if not item.runner_ready:
        return "not runner-ready (needs Done, not owner-bound)"
    if item.sessions in ("owner", "solo"):
        return f"Sessions: {item.sessions}"
    if item.effort not in TK.EFFORTS or TK.EFFORTS.index(item.effort) > TK.EFFORTS.index(cfg.slice_above):
        return "slice job (effort above slice_above)"
    if item.model != "opus":
        return f"Model {item.model} stays local"
    if item.cloud == "yes":
        return None
    text = "\n".join([item.text, *item.body])
    for reason, rx in UNFIT:
        if rx.search(text):
            return reason
    return None


def pick_key(item: TK.Item) -> tuple:
    """wf orch pick cloud order (spec §4.5): Cloud: yes first, then prio, then larger effort (all fits are opus)."""
    eff = TK.EFFORTS.index(item.effort) if item.effort in TK.EFFORTS else -1
    return (item.cloud != "yes", item.prio if item.prio is not None else 9, -eff)