"""Canopy reader: read one AI coding agent terminal screen and build a status card from its own lines.

Self-contained: needs only `onnxruntime`, `tokenizers` and `numpy`. Put this file next to (or point it at)
a model folder holding `reader.json`, `tokenizer.json` and `model.int8.onnx`.

    from canopy_reader import CanopyReader
    reader = CanopyReader("path/to/canopy-reader-17m")
    reading = reader.read(state, lines)      # every head, as line indices and classes
    card = reader.card(reading, lines)       # the card as text

`lines` is the screen as a list of lines, bottom last (wrapping undone, chrome stripped: see the model
card). `state` is a small dict: {"agent_program": "claude", "terminal_title": "...", "daintree_observed":
{"agent_state": "waiting", "waiting_reason": "prompt", "seconds_in_state": 12}}; only `agent_program` is
required.
"""
from __future__ import annotations

import json
import math
import re
from pathlib import Path

import numpy as np


def _strip(s: str) -> str:
    return s.strip()


def line_candidates(lines: list[str], rows: int) -> list[int]:
    """The bottom `rows` lines longer than three characters."""
    start = max(0, len(lines) - rows)
    return [i for i in range(start, len(lines)) if len(lines[i].strip()) > 3]


def _sigmoid(x: float) -> float:
    return 1 / (1 + math.exp(-x))


def _softmax(xs) -> list[float]:
    m = max(xs)
    e = [math.exp(x - m) for x in xs]
    s = sum(e)
    return [x / s for x in e]


class CanopyReader:
    def __init__(self, folder: str | Path, variant: str = "int8", threads: int = 1):
        import onnxruntime as ort
        from tokenizers import Tokenizer
        folder = Path(folder)
        # config.json first: Hugging Face counts a download per fetch of it (reader.json is the same file).
        meta_file = folder / "config.json" if (folder / "config.json").exists() else folder / "reader.json"
        self.meta = json.loads(meta_file.read_text(encoding="utf-8"))
        self.tok = Tokenizer.from_file(str(folder / "tokenizer.json"))
        opts = ort.SessionOptions()
        opts.intra_op_num_threads = threads
        opts.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        self.session = ort.InferenceSession(str(folder / f"model.{variant}.onnx"), opts, providers=["CPUExecutionProvider"])
        self.outputs = [o.name for o in self.session.get_outputs()]
        self.cls, self.sep, self.pad = (self.tok.token_to_id(t) for t in ("[CLS]", "[SEP]", "[PAD]"))

    def _ids(self, text: str) -> list[int]:
        return self.tok.encode(text, add_special_tokens=False).ids

    def encode(self, state: dict, lines: list[str]) -> dict:
        """`[CLS] state-json [SEP]`, then `[i] line\\n` for each non-blank line, filled from the bottom of the
        screen up until `max_len` tokens, then `[SEP]`. Each line's first token is where its heads read."""
        m = self.meta
        header = self._ids(json.dumps(state, ensure_ascii=False, separators=(",", ":")))[: m["header_tokens"]]
        budget = m["max_len"] - len(header) - 3
        pieces = []
        for i in range(len(lines) - 1, -1, -1):
            if not lines[i].strip():
                continue
            ids = self._ids(f"[{i}] {lines[i]}\n")[: m["line_tokens"]]
            if len(ids) > budget:
                break
            budget -= len(ids)
            pieces.append((i, ids))
        pieces.reverse()
        ids = [self.cls, *header, self.sep]
        index, pos = [], []
        for i, p in pieces:
            index.append(i)
            pos.append(len(ids))
            ids.extend(p)
        ids.append(self.sep)
        q = set(line_candidates(lines, m["question_candidate_rows"]))
        s = set(line_candidates(lines, m["status_candidate_rows"]))
        return {"ids": ids, "index": index, "pos": pos, "q": [i in q for i in index], "s": [i in s for i in index]}

    def read(self, state: dict, lines: list[str]) -> dict:
        e = self.encode(state, lines)
        k = max(1, len(e["index"]))
        pad = lambda xs, v: xs + [v] * (k - len(xs))  # noqa: E731
        feed = {
            "input_ids": np.array([e["ids"]], dtype=np.int64),
            "attention_mask": np.ones((1, len(e["ids"])), dtype=np.int64),
            "line_pos": np.array([pad(e["pos"], 0)], dtype=np.int64),
            "line_mask": np.array([pad([True] * len(e["index"]), False)]),
            "q_mask": np.array([pad(e["q"], False)]),
            "s_mask": np.array([pad(e["s"], False)]),
        }
        out = dict(zip(self.outputs, self.session.run(self.outputs, feed)))
        return self.decode({n: v[0].astype(np.float64).tolist() for n, v in out.items()}, e)

    def decode(self, out: dict, e: dict) -> dict:
        m, idx, t = len(e["index"]), e["index"], self.meta["thresholds"]

        def pick(name):
            scores = out[name]
            cand = scores[:m] + [scores[-1]]  # the last column is "none"
            j = max(range(len(cand)), key=lambda x: (cand[x], -x))
            return -1 if j == m else idx[j]

        logit = out["summary_lines"][:m]
        p_sum = [_sigmoid(x) for x in logit]
        summary = []
        for j in sorted(range(m), key=lambda j: (-logit[j], j))[:3]:  # rank by logit: no saturation ties
            if p_sum[j] >= (t["summary_more"] if summary else t["summary_first"]):
                summary.append(idx[j])
            else:
                break
        first = pick("status_line")
        second_th = t.get("second")
        card_lines = ([first] if first >= 0 else [])
        if second_th is not None:  # the summary-only card: line two is the best other line scoring about a 2
            card_lines += [idx[j] for j in sorted(range(m), key=lambda j: (-logit[j], j)) if idx[j] != first and p_sum[j] >= second_th][:1]
        probs = _softmax(out["category"])
        cats = self.meta["categories"]
        raw = _sigmoid(out["needs_attention"])
        g = max(raw, sum(probs[cats.index(c)] for c in t["gate_categories"]))
        gate = 0.5 * g / t["gate"] if g < t["gate"] else 0.5 + 0.5 * (g - t["gate"]) / (1 - t["gate"])
        cls = lambda n: self.meta["classes"][n][max(range(len(out[n])), key=lambda i: (out[n][i], -i))]  # noqa: E731
        return {
            "category": cats[max(range(len(probs)), key=lambda i: (probs[i], -i))], "category_probs": probs,
            "needs_attention": gate, "needs_attention_raw": raw,
            "blocked_on_user": _sigmoid(out["blocked_on_user"]), "notable": _sigmoid(out["notable"]),
            "question_line": pick("question_line"), "status_line": pick("status_line"),
            "action_line": pick("action_line"), "task_line": pick("task_line"),
            "summary_lines": summary, "card_lines": card_lines,
            "option_lines": [idx[j] for j in range(m) if e["q"][j] and _sigmoid(out["option_lines"][j]) >= t["option"]],
            "key_lines": [idx[j] for j in range(m) if out["key_lines"][j] > 0],
            **{n: cls(n) for n in self.meta["classes"]},
        }

    # --- the card as text ---------------------------------------------------------------------------

    # A leading glyph is stripped only when whitespace (or the end) follows it, so `**bold**` and `>file` survive.
    LEAD = re.compile(r"^(?:[⏺•●✦⎿└├│┃║■✻✶✢✳✽✺·*❯›>▸▶◆◇-](?:\s+|$))+")
    BOX_TAIL = re.compile(r"\s*[│┃║]\s*$")
    TIMER = re.compile(r"\s*\((?:\d[^)]*)?(?:esc to \w+|thinking|thought for|ctrl\+\w to \w+)[^)]*\)\s*$", re.I)
    END = re.compile(r"""[.!?:;)"'`]$""")

    @classmethod
    def text(cls, line: str | None) -> str | None:
        """A screen line as card text: bullets, tree and box glyphs, spinners and timers off."""
        if line is None:
            return None
        t = cls.TIMER.sub("", cls.BOX_TAIL.sub("", cls.LEAD.sub("", line.strip()))).strip()
        return t if len(t) >= 2 else None

    @staticmethod
    def _short(line: str) -> str | None:
        """A one-character fragment that text() drops for length (a lone "1"), not text the cleanup removed."""
        t = line.strip()
        return t if len(t) == 1 else None

    def _continues(self, prev: str, cur: str) -> bool:
        """Does line `cur` carry on the sentence of `prev` (the agent wrapped it onto a new row)?"""
        if not prev.strip() or not cur.strip() or self.LEAD.match(cur.strip()):
            return False
        pt = self.text(prev)
        if not pt or self.END.search(pt):
            return False
        return bool(re.match(r"^\s{2,}\S", cur) or cur.strip()[:1].islower() or cur.strip()[:1] in "([{'\"`")

    def _sentence(self, lines, i, picks, used):
        # A pick in the middle of a wrapped sentence starts where the sentence starts (up to two rows back).
        k = i
        while k > 0 and i - k < 4 and k - 1 not in used and self._continues(lines[k - 1], lines[k]):
            k -= 1
        for x in range(k, i):
            used.add(x)
        t = " ".join(filter(None, (self.text(lines[x]) or (self._short(lines[x]) if x != i else None) for x in range(k, i + 1))))
        used.add(i)
        j = i + 1
        while t and j <= i + 4 and j < len(lines) and not self.END.search(t) and j not in used:
            nxt = lines[j]
            if not nxt.strip() or self.LEAD.match(nxt.strip()):
                break  # a new bullet is a new item, even when it was picked too
            if j not in picks and not self._continues(lines[j - 1], nxt):
                break
            used.add(j)
            more = self.text(nxt) or self._short(nxt)  # keep a one-character fragment ("1") the sentence carries on into
            if more:
                t += " " + more
            j += 1
        return t

    def card(self, r: dict, lines: list[str]) -> dict:
        asks = r["category"] in ("approval", "question")
        used: set[int] = set()
        line = lambda i: self.text(lines[i]) if i >= 0 else None  # noqa: E731
        if "card_lines" in r and self.meta["thresholds"].get("card") == "summary-only":
            # The summary-only card: no headline (the terminal title is the title), up to two lines.
            headline = None
            picks = list(r["card_lines"])
            if len(picks) == 2 and abs(picks[0] - picks[1]) <= 3:
                picks.sort()  # two lines of one paragraph read in screen order
            summary = [s for i in picks if i not in used for s in [self._sentence(lines, i, picks, used)] if s]
        else:
            headline = self._sentence(lines, r["status_line"], r["summary_lines"], used) if r["status_line"] >= 0 else None
            summary = [s for i in r["summary_lines"] if i not in used for s in [self._sentence(lines, i, r["summary_lines"], used)] if s]
        return {
            "category": r["category"], "needs_attention": r["needs_attention"] >= 0.5, "score": int(r["score"]),
            "headline": headline, "summary": summary,
            "question": line(r["question_line"]) if asks else None,
            "options": [x for x in (self.text(lines[i]) for i in r["option_lines"]) if x] if asks else [],
            "action": line(r["action_line"]) if r["category"] == "approval" else None,
            "task": line(r["task_line"]),
            "tests": r["tests"], "changes": r["changes"],
            "progress": None if r["progress"] == "none" else int(r["progress"]),
        }