"""One episode: the free screen, up to eight replies, and the hour meter.

Pure and deterministic in the world it is given. It never raises on input: a
bad reply is a rejection with a one-line reason, charged a turn.

Two discipline rules the reviewers demanded, both enforced here rather than in
the grader:

  * **Shown ids.** A disposition may cite only ids the engine has actually
    rendered — KYC fields and triggering transactions from the screen, plus
    whatever lookups returned. "Closed per the probate letter on file"
    without opening the file is a rejected reply, not a scored one.
  * **Abort versus truncate.** `aborted` is exactly MAX_TURNS turns spent
    with no accepted disposition, and scores 0.000. Anything that stops
    earlier with no disposition is `truncated` and is not scored at all.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

from .generator import MAX_TURNS, screen_of
from .protocol import (
    ACTIONS,
    CLOSE_REASONS,
    LOOKUPS,
    TYPOLOGIES,
    _MONTH,
    parse_reply,
)
from .reward import DOC_MINUTES, LOOKUP_MINUTES, SCREEN_MINUTES, TURN_MINUTES, Episode

LOOKUP_PRICE = {"history": LOOKUP_MINUTES, "counterparty": LOOKUP_MINUTES, "prior_alerts": LOOKUP_MINUTES, "documents": DOC_MINUTES}


@dataclass
class Step:
    accepted: bool
    action: dict[str, Any] | None      # canonical form of an accepted action
    reason: str | None                 # why a rejected reply was rejected
    result: dict[str, Any] | None      # what an accepted lookup returned
    minutes_spent: int
    turns_spent: int
    done: bool


@dataclass
class Engine:
    world: dict[str, Any]
    turns: int = 0
    rejected: int = 0
    minutes: int = SCREEN_MINUTES
    actions: list[dict[str, Any]] = field(default_factory=list)
    lookups: list[dict[str, Any]] = field(default_factory=list)
    disposition: dict[str, Any] | None = None
    shown: set[str] = field(default_factory=set)
    last_rejection: str | None = None

    def __post_init__(self) -> None:
        screen = self.screen()
        self.shown.update(f"kyc.{k}" for k in screen["customer"] if k not in ("id", "name"))
        self.shown.add(screen["alert"]["id"])
        for t in screen["triggering_transactions"]:
            self.shown.add(t["id"])
            if t["cp"]:
                self.shown.add(t["cp"])
        for cp in screen["screen_counterparties"]:
            self.shown.add(cp["id"])

    # ------------------------------------------------------------ views --

    def screen(self) -> dict[str, Any]:
        screen = screen_of(self.world)
        if screen["screen_order"] == "kyc_first":
            screen["summary"] = list(reversed(screen["summary"]))
        return screen

    def view(self) -> dict[str, Any]:
        """Everything the model has been shown, and nothing it has not.

        The scripted policies take this and only this. There is no label,
        typology, planted set, overlay or reference in it — `tests/` scans
        the keys to make sure.
        """
        return {
            "screen": self.screen(),
            "lookups": list(self.lookups),
            "turns_spent": self.turns,
            "turns_left": MAX_TURNS - self.turns,
            "rejected": self.rejected,
            "minutes_spent": self.minutes,
            "last_rejection": self.last_rejection,
        }

    @property
    def done(self) -> bool:
        return self.disposition is not None or self.turns >= MAX_TURNS

    @property
    def aborted(self) -> bool:
        return self.disposition is None and self.turns >= MAX_TURNS

    # ------------------------------------------------------------- step --

    def step(self, reply: str | None) -> Step:
        if self.done:
            return Step(False, None, "the episode is over", None, self.minutes, self.turns, True)
        obj, reason = parse_reply(reply)
        action = None
        if obj is not None:
            action, reason = self._validate(obj)
        self.turns += 1
        self.minutes += TURN_MINUTES
        if action is None:
            self.rejected += 1
            self.last_rejection = reason
            return Step(False, None, reason, None, self.minutes, self.turns, self.done)

        self.last_rejection = None
        self.actions.append(action)
        result = None
        if action["action"] == "lookup":
            self.minutes += LOOKUP_PRICE[action["what"]]
            result = self._lookup(action)
            self.lookups.append({"what": action["what"], "key": action.get("month") or action.get("id"), "result": result})
        else:
            self.disposition = action
        return Step(True, action, None, result, self.minutes, self.turns, self.done)

    def _validate(self, obj: dict[str, Any]) -> tuple[dict[str, Any] | None, str | None]:
        action = obj.get("action")
        if action not in ACTIONS:
            return None, f"unknown action {action!r}; use lookup, close or escalate"
        if action == "lookup":
            what = obj.get("what")
            if what not in LOOKUPS:
                return None, f"unknown lookup {what!r}; use history, counterparty, prior_alerts or documents"
            if what == "history":
                month = obj.get("month")
                if not isinstance(month, str) or not _MONTH.match(month) or month not in self.world["months"]:
                    return None, f"month {month!r} is not in the twelve-month window"
                if any(l["what"] == "history" and l["key"] == month for l in self.lookups):
                    return None, f"history for {month} has already been looked up"
                return {"action": "lookup", "what": "history", "month": month}, None
            if what == "counterparty":
                cid = obj.get("id")
                if not isinstance(cid, str) or cid not in self.shown or not cid.startswith("CP-"):
                    return None, f"counterparty {cid!r} has not appeared on your screen or in a fetched month"
                if any(l["what"] == "counterparty" and l["key"] == cid for l in self.lookups):
                    return None, f"counterparty {cid} has already been looked up"
                return {"action": "lookup", "what": "counterparty", "id": cid}, None
            if any(l["what"] == what for l in self.lookups):
                return None, f"{what} has already been looked up"
            return {"action": "lookup", "what": what}, None

        cites = obj.get("cites")
        if not isinstance(cites, list) or not cites:
            return None, "cites must be a non-empty list of ids"
        seen: list[str] = []
        for c in cites:
            if not isinstance(c, str):
                return None, "cites must be a non-empty list of ids"
            if c not in seen:
                seen.append(c)
        for c in seen:
            if c not in self.shown:
                return None, f"cites an id you have not been shown: {c}"
        if action == "close":
            reason = obj.get("reason")
            if reason not in CLOSE_REASONS:
                return None, f"unknown close reason {reason!r}"
            note = obj.get("note")
            return {"action": "close", "reason": reason, "cites": seen, "note": note if isinstance(note, str) else ""}, None
        typology = obj.get("typology")
        if typology not in TYPOLOGIES:
            return None, f"unknown typology {typology!r}"
        narrative = obj.get("narrative")
        return {"action": "escalate", "typology": typology, "cites": seen, "narrative": narrative if isinstance(narrative, str) else ""}, None

    def _lookup(self, action: dict[str, Any]) -> dict[str, Any]:
        w = self.world
        what = action["what"]
        if what == "history":
            rows = [t for t in w["transactions"] if t["month"] == action["month"]]
            for t in rows:
                self.shown.add(t["id"])
                if t["cp"]:
                    self.shown.add(t["cp"])
            return {"month": action["month"], "transactions": rows}
        if what == "counterparty":
            cp = next(c for c in w["counterparties"] if c["id"] == action["id"])
            self.shown.update(cp["transactions"])
            return dict(cp)
        if what == "prior_alerts":
            self.shown.update(p["id"] for p in w["prior_alerts"])
            return {"alerts": list(w["prior_alerts"])}
        self.shown.update(d["id"] for d in w["documents"])
        return {"documents": list(w["documents"])}

    # ---------------------------------------------------------- outcome --

    @property
    def outcome(self) -> str | None:
        """solved | failed | aborted, or None while running / if truncated."""
        if self.disposition is not None:
            label = self.world["label"]
            right = (label == "benign") == (self.disposition["action"] == "close")
            return "solved" if right else "failed"
        return "aborted" if self.aborted else None

    def episode(self, truncated: bool | None = None) -> Episode:
        if truncated is None:
            truncated = not self.done
        d = self.disposition
        return Episode(
            label=self.world["label"],
            planted=self.world["planted"],
            disposition=None if d is None else d["action"],
            cites=[] if d is None else list(d["cites"]),
            minutes_spent=self.minutes,
            reference_minutes=self.world.get("reference_minutes"),
            turns_spent=self.turns,
            rejected=self.rejected,
            truncated=truncated,
            typology=None if d is None else d.get("typology"),
            true_typology=self.world.get("typology"),
            lookups=len(self.lookups),
        )
