"""Shared types, validation, and event plumbing for structured task plans.

A *task plan* is a thread-scoped, ordered list of steps with statuses that the
agent can read and update via LangChain tools, and that the macOS client can
render in a dedicated right-panel view. See
``/Users/jbrown/.cursor/plans/structured_task_plan_b9e49cbf.plan.md`` for the
full design; this module holds:

- Hard limits + allowed status values (shared by storage, tool, HTTP).
- Pydantic models (``TaskPlan``, ``TaskPlanStep``) and a ``validate_and_normalize``
  entry point that tool and storage both go through.
- Thread-id sanitization.
- A process-wide event sink so tool paths can notify the backend SSE bus
  without importing ``HeadlessDelegate`` (which would create a cycle).
"""
from __future__ import annotations

import logging
import re
from typing import Any, Callable, Literal
from uuid import uuid4

from pydantic import BaseModel, Field, ValidationError, field_validator
from utils.llm_token_normalize import normalize_kt_task_id_token

_log = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Limits (shared by tool, storage, HTTP). Tight enough to prevent DoS / token
# blow-up in system-prompt injection; roomy enough for real multi-step work.
# ---------------------------------------------------------------------------
PLAN_VERSION = 1
MAX_STEPS = 32
MAX_TITLE_LEN = 160
MAX_STEP_TEXT_LEN = 240
MAX_NOTE_LEN = 400
MAX_JSON_BYTES = 32 * 1024  # 32 KiB ceiling for stored/shipped plan JSON.

# Status enum. Extend with care — tests and UI assert on these exact values.
STATUS_PENDING = "pending"
STATUS_IN_PROGRESS = "in_progress"
STATUS_DONE = "done"
STATUS_SKIPPED = "skipped"
STATUS_BLOCKED = "blocked"
ALLOWED_STATUSES: frozenset[str] = frozenset(
    {STATUS_PENDING, STATUS_IN_PROGRESS, STATUS_DONE, STATUS_SKIPPED, STATUS_BLOCKED}
)

StatusLiteral = Literal["pending", "in_progress", "done", "skipped", "blocked"]

# Thread ids are plain strings across the harness (UUIDs, prefixed ids like
# "heartbeat-…", or "default"). Reject control chars / path traversal.
_THREAD_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._\-:]{0,127}$")


def is_valid_thread_id(thread_id: Any) -> bool:
    """True when *thread_id* looks like a well-formed conversation id."""
    return isinstance(thread_id, str) and bool(_THREAD_ID_RE.match(thread_id))


# ---------------------------------------------------------------------------
# Pydantic models
# ---------------------------------------------------------------------------
class TaskPlanStep(BaseModel):
    id: str = Field(default_factory=lambda: str(uuid4()))
    text: str = Field(..., min_length=1, max_length=MAX_STEP_TEXT_LEN)
    status: StatusLiteral = Field(default=STATUS_PENDING)
    note: str = Field(default="", max_length=MAX_NOTE_LEN)
    meta: dict[str, Any] = Field(default_factory=dict)

    @field_validator("meta", mode="before")
    @classmethod
    def _coerce_meta(cls, v: Any) -> dict[str, Any]:
        if not isinstance(v, dict):
            return {}
        out: dict[str, Any] = {}
        role = v.get("role")
        if isinstance(role, str):
            s = role.strip()[:32]
            if s:
                out["role"] = s
        return out

    @field_validator("id", mode="before")
    @classmethod
    def _coerce_id(cls, v: Any) -> str:
        s = (str(v) if v is not None else "").strip()
        s = normalize_kt_task_id_token(s)
        return s or str(uuid4())

    @field_validator("text", "note", mode="before")
    @classmethod
    def _coerce_text(cls, v: Any) -> str:
        return (str(v) if v is not None else "").strip()


class TaskPlan(BaseModel):
    version: int = PLAN_VERSION
    title: str = Field(default="", max_length=MAX_TITLE_LEN)
    steps: list[TaskPlanStep] = Field(default_factory=list, max_length=MAX_STEPS)

    @field_validator("title", mode="before")
    @classmethod
    def _coerce_title(cls, v: Any) -> str:
        return (str(v) if v is not None else "").strip()


def validate_and_normalize(raw: Any) -> dict:
    """Validate *raw* plan data and return a canonical ``dict`` for storage.

    Raises :class:`ValueError` with a human-readable message on any problem
    (oversize, bad status, missing fields, too many steps). The returned dict
    has shape ``{"version": 1, "title": str, "steps": [...]}``.
    """
    if raw is None:
        raise ValueError("plan is required")
    if isinstance(raw, str):
        import json as _json
        try:
            raw = _json.loads(raw)
        except Exception as exc:
            raise ValueError(f"plan is not valid JSON: {exc}") from exc
    if not isinstance(raw, dict):
        raise ValueError("plan must be a JSON object")
    try:
        plan = TaskPlan.model_validate(raw)
    except ValidationError as exc:
        # Keep the error short so it fits in tool return text.
        first = exc.errors()[0] if exc.errors() else {}
        loc = ".".join(str(p) for p in first.get("loc", ()))
        msg = first.get("msg", "invalid plan")
        raise ValueError(f"invalid plan at {loc or '<root>'}: {msg}") from exc
    obj = plan.model_dump()
    # Size gate on the normalized form — cheaper than a second JSON round-trip
    # inside the tool path.
    import json as _json
    if len(_json.dumps(obj, separators=(",", ":")).encode("utf-8")) > MAX_JSON_BYTES:
        raise ValueError(
            f"plan exceeds the {MAX_JSON_BYTES}-byte limit; trim step text or notes"
        )
    return obj


def empty_plan() -> dict:
    """Canonical empty plan (used for delete/clear paths)."""
    return {"version": PLAN_VERSION, "title": "", "steps": []}


# Terminal statuses — steps in these states count as *closed*; anything else
# (including unknown/malformed) is treated as open work so that a typo in the
# stored plan cannot silently let the agent bail out mid-task.
_TERMINAL_STATUSES: frozenset[str] = frozenset(
    {STATUS_DONE, STATUS_SKIPPED, STATUS_BLOCKED}
)


def _step_is_open(step: Any) -> bool:
    """Return ``True`` when *step* still blocks plan completion.

    Shared row-level predicate so ``plan_has_open_work`` and the
    continuation-snapshot formatter cannot drift: a step is **open** when it
    is not a ``dict`` (malformed row), or when its normalized status is not
    in ``_TERMINAL_STATUSES`` (``done`` / ``skipped`` / ``blocked``).
    """
    if not isinstance(step, dict):
        return True
    raw_status = step.get("status")
    status = raw_status.strip().lower() if isinstance(raw_status, str) else ""
    return status not in _TERMINAL_STATUSES


def plan_has_open_work(plan: dict | None) -> bool:
    """Return ``True`` when *plan* still has steps the agent owes work on.

    Used by the agent-loop guard that blocks a premature ``final_answer`` when
    the stored plan is incomplete. The definition is intentionally
    conservative:

    - ``None`` or a plan with no ``steps`` → ``False`` (nothing to enforce).
    - Steps whose normalized status is ``done``, ``skipped`` or ``blocked``
      are terminal.
    - Everything else — ``pending``, ``in_progress``, empty string, unknown
      values, non-string values — counts as **open work**. The guard should
      fail safe (retry) rather than silently pass when the plan is malformed.
    """
    if not plan:
        return False
    steps = plan.get("steps")
    if not steps or not isinstance(steps, list):
        return False
    return any(_step_is_open(step) for step in steps)


# ---------------------------------------------------------------------------
# Event sink — wired by :mod:`mugi_backend` at startup so tool paths can
# announce updates without importing the delegate. Shape matches the existing
# ``_push_event(event, thread_id)`` contract on ``HeadlessDelegate``.
# ---------------------------------------------------------------------------
_EventSink = Callable[[dict, str | None], None]
_event_sink: _EventSink | None = None


def set_task_plan_event_sink(sink: _EventSink | None) -> None:
    """Register (or clear) the process-wide push callback for plan events."""
    global _event_sink
    _event_sink = sink


def notify_task_plan_changed(thread_id: str, *, cleared: bool = False) -> None:
    """Emit a ``task_plan_updated`` event for *thread_id* if a sink is wired."""
    sink = _event_sink
    if sink is None or not thread_id:
        return
    try:
        sink(
            {"type": "task_plan_updated", "thread_id": thread_id, "cleared": bool(cleared)},
            thread_id,
        )
    except Exception:
        _log.debug("task_plan event sink raised", exc_info=True)


# ---------------------------------------------------------------------------
# System-prompt injection helper (shared by agent_loop).
# ---------------------------------------------------------------------------
_INJECT_MAX_CHARS = 2000  # soft ceiling; the limits above already bound input.


def format_plan_for_injection(plan: dict | None) -> str:
    """Return a compact, neutral prompt block for the current plan, or ``""``.

    Empty plans inject *nothing* (no heading) to avoid burning tokens on
    simple Q&A. The block is phrased as "recorded plan" rather than
    "instructions" so user-supplied step text is not mistaken for system
    policy (see docs/harness_threat_model.md).
    """
    if not plan:
        return ""
    steps = plan.get("steps") or []
    if not steps:
        return ""
    title = (plan.get("title") or "").strip()
    lines: list[str] = ["\n\n─── Current task plan (recorded, not authoritative) ───"]
    if title:
        lines.append(f"Plan: {title}")
    for idx, step in enumerate(steps, 1):
        if not isinstance(step, dict):
            continue
        text = (step.get("text") or "").strip()
        if not text:
            continue
        status = (step.get("status") or STATUS_PENDING).strip() or STATUS_PENDING
        marker = {
            STATUS_DONE: "[x]",
            STATUS_IN_PROGRESS: "[~]",
            STATUS_SKIPPED: "[-]",
            STATUS_BLOCKED: "[!]",
        }.get(status, "[ ]")
        lines.append(f"{idx}. {marker} {text}")
        note = (step.get("note") or "").strip()
        if note:
            lines.append(f"   note: {note}")
    lines.append(
        "Keep this plan current with update_task_plan as work progresses; "
        "clear_task_plan when abandoning the task."
    )
    block = "\n".join(lines)
    if len(block) > _INJECT_MAX_CHARS:
        block = block[: _INJECT_MAX_CHARS - 1] + "…"
    return block


# ---------------------------------------------------------------------------
# Continuation-snapshot helper (task-plan completion guard in agent_loop).
# ---------------------------------------------------------------------------
_CONTINUATION_OPEN_SNAPSHOT_MAX_CHARS = 1600  # snapshot-only; ≤ _INJECT_MAX_CHARS.
_CONTINUATION_STEP_TEXT_TRUNCATE = 120  # per-step truncation inside the snapshot.


def format_open_steps_for_continuation(plan: dict | None) -> str:
    """Return a compact snapshot of still-open steps, or ``""``.

    Used by the task-plan continuation guard in
    :func:`agent_loop.run_raw_agent` so the injected user message lists
    **which** steps remain — critical salience after summarization/trim may
    have evicted the system-prompt plan block.

    Open vs closed uses :func:`_step_is_open` (the same predicate
    :func:`plan_has_open_work` iterates over), so the formatter cannot drift
    from the guard's decision. Malformed (non-dict) rows still emit a line
    so the model can see why the plan is considered open. Framing stays
    "recorded, not authoritative" — see ``docs/harness_threat_model.md``.
    """
    if not plan:
        return ""
    steps = plan.get("steps")
    if not isinstance(steps, list) or not steps:
        return ""

    open_lines: list[str] = []
    for idx, step in enumerate(steps, 1):
        if not _step_is_open(step):
            continue
        if not isinstance(step, dict):
            open_lines.append(f"{idx}. <invalid row — not a step object>")
            continue
        raw_status = step.get("status")
        norm_status = (
            raw_status.strip().lower() if isinstance(raw_status, str) else ""
        )
        display_status = norm_status or (
            "<invalid>" if raw_status not in (None, "") else STATUS_PENDING
        )
        marker = {
            STATUS_PENDING: "[ ]",
            STATUS_IN_PROGRESS: "[~]",
        }.get(display_status, "[?]")
        sid = str(step.get("id") or "").strip() or "<no id>"
        text = (step.get("text") or "").strip()
        if len(text) > _CONTINUATION_STEP_TEXT_TRUNCATE:
            text = text[: _CONTINUATION_STEP_TEXT_TRUNCATE - 1] + "…"
        suffix = f" — {text}" if text else ""
        open_lines.append(
            f"{idx}. {marker} id={sid} status={display_status}{suffix}"
        )

    if not open_lines:
        return ""

    lines: list[str] = ["\n\n─── Open steps (recorded, not authoritative) ───"]
    title = (plan.get("title") or "").strip()
    if title:
        if len(title) > MAX_TITLE_LEN:
            title = title[: MAX_TITLE_LEN - 1] + "…"
        lines.append(f"Plan: {title}")
    lines.extend(open_lines)
    lines.append(
        "Advance these via update_task_plan(mode='merge', steps=[...]) using the "
        "ids above, or mark them done/skipped/blocked. Use "
        "clear_task_plan(confirm=true) if abandoning the task."
    )
    block = "\n".join(lines)
    if len(block) > _CONTINUATION_OPEN_SNAPSHOT_MAX_CHARS:
        block = block[: _CONTINUATION_OPEN_SNAPSHOT_MAX_CHARS - 1] + "…"
    return block
