"""RegressionService — capture correlation chains as expected-behavior baselines.

A "regression test" is a user-curated pointer to a correlation_id chain
that the user wants to preserve as expected behavior. Unlike unit tests,
regression tests run against the live agent/vault/LLM — the goal is to
catch drift in behavior over time, not correctness bugs.

Captured baseline includes:
  - The initial user input that triggered the chain (best-effort)
  - The tool sequence that ran (names + outcomes + durations)
  - The set of actors involved (main, librarian, etc.)
  - The final status (all-ok / any-failure / any-zero-effect)

Replay re-posts the original input through the agent and records the new
correlation_id + a shallow diff of tool sequences. The diff is
intentionally coarse — anything richer requires either hermetic vault
state or a semantic equivalence judge, both out of scope for MVP.

Rule 13 compliance: every capture/replay/delete emits an activity event
through the caller's middleware. Contract action: ``capture_regression``
(a new action; default confirm).
"""

from __future__ import annotations

import json
import logging
from dataclasses import asdict, dataclass, field
from datetime import UTC, datetime
from typing import TYPE_CHECKING, Any

import aiosqlite

if TYPE_CHECKING:
    from agntspace.activity.decisions import DecisionLogger
    from agntspace.agent.agent import Agent

log = logging.getLogger(__name__)


@dataclass
class RegressionBaseline:
    """Frozen snapshot of what the chain did the first time.

    Only fields the replay comparator needs are captured — we intentionally
    do NOT store the full decision log or tool outputs (they can be looked
    up by correlation_id when needed).
    """

    correlation_id: str
    user_input: str = ""
    actors: list[str] = field(default_factory=list)
    tool_sequence: list[dict] = field(default_factory=list)
    final_status: str = "unknown"  # all_ok / partial_failure / any_failure / zero_effect / unknown
    decision_count: int = 0
    duration_ms: int = 0


@dataclass
class RegressionTest:
    """Row shape returned by list/get."""

    id: int
    name: str
    description: str
    correlation_id: str
    captured_at: str
    baseline: dict
    archived: bool
    last_replay_correlation_id: str | None
    last_replay_status: str | None
    last_replay_at: str | None
    last_replay_diff: dict | None


class RegressionService:
    """Curation + replay layer over the existing decision log.

    Wired post-init in main.py; requires the DecisionLogger and the DB.
    Agent reference is optional — needed only for automatic replay via
    ``agent.chat()``.
    """

    def __init__(
        self,
        *,
        db: aiosqlite.Connection,
        decisions: DecisionLogger,
        agent: Agent | None = None,
    ) -> None:
        self._db = db
        self._decisions = decisions
        self._agent = agent

    # ─── Baseline extraction ───

    async def _extract_baseline(self, correlation_id: str) -> RegressionBaseline:
        """Walk the decision log for a correlation_id, build a baseline.

        Never raises — returns a minimal baseline with unknown fields on
        any query failure. This is the read path; we want it robust.
        """
        baseline = RegressionBaseline(correlation_id=correlation_id)
        try:
            rows = await self._decisions.query(
                correlation_id=correlation_id, limit=1000,
            )
        except Exception:
            log.debug("regression baseline: decision query failed", exc_info=True)
            return baseline

        if not rows:
            return baseline

        # Query returns newest-first; we want chronological for tool order.
        # Rows may be Decision dataclasses OR dicts depending on backend.
        def _get(r: Any, key: str, default: Any = "") -> Any:
            if hasattr(r, key):
                return getattr(r, key, default)
            if isinstance(r, dict):
                return r.get(key, default)
            return default

        def _ts(r: Any) -> str:
            return str(_get(r, "ts", ""))

        chronological = sorted(rows, key=_ts)
        baseline.decision_count = len(chronological)

        actors: set[str] = set()
        tool_sequence: list[dict] = []
        failure_count = 0
        zero_effect_count = 0

        for r in chronological:
            actor = str(_get(r, "actor", ""))
            if actor:
                actors.add(actor)

            action_type = str(_get(r, "action_type", ""))
            action_target = str(_get(r, "action_target", ""))
            result_status = str(_get(r, "result_status", "") or "")
            duration = _get(r, "duration_ms", 0) or 0

            if action_type in ("tool_call", "tool", "skill_invoke"):
                tool_sequence.append({
                    "tool": action_target or action_type,
                    "status": result_status,
                    "duration_ms": int(duration) if duration else 0,
                })

            if result_status in ("failed", "blocked", "error"):
                failure_count += 1
            elif result_status == "zero_effect":
                zero_effect_count += 1

            # First row's details sometimes contain the user input for
            # chat-originated chains. Best effort, never blocks capture.
            if not baseline.user_input:
                details = _get(r, "details", {})
                if isinstance(details, dict):
                    candidate = details.get("user_input") or details.get("prompt") or details.get("message") or ""
                    if isinstance(candidate, str) and candidate.strip():
                        baseline.user_input = candidate.strip()[:2000]

        baseline.actors = sorted(actors)
        baseline.tool_sequence = tool_sequence

        if failure_count > 0 and failure_count < len(chronological):
            baseline.final_status = "partial_failure"
        elif failure_count > 0:
            baseline.final_status = "any_failure"
        elif zero_effect_count > 0:
            baseline.final_status = "zero_effect"
        elif chronological:
            baseline.final_status = "all_ok"

        # Best-effort duration: span from first to last timestamp.
        try:
            first_ts = datetime.fromisoformat(_ts(chronological[0]).replace("Z", "+00:00"))
            last_ts = datetime.fromisoformat(_ts(chronological[-1]).replace("Z", "+00:00"))
            baseline.duration_ms = int((last_ts - first_ts).total_seconds() * 1000)
        except Exception:
            baseline.duration_ms = 0

        return baseline

    # ─── Mutations ───

    async def capture(
        self,
        *,
        correlation_id: str,
        name: str,
        description: str = "",
    ) -> RegressionTest:
        """Capture a correlation chain as a regression baseline.

        Raises ``ValueError`` when the correlation_id has no decisions —
        nothing to test against.
        """
        if not correlation_id or not correlation_id.strip():
            raise ValueError("correlation_id required")
        if not name or not name.strip():
            raise ValueError("name required")

        baseline = await self._extract_baseline(correlation_id)
        if baseline.decision_count == 0:
            raise ValueError(
                f"correlation_id {correlation_id!r} has no decisions — nothing to capture",
            )

        now_iso = datetime.now(UTC).isoformat()
        payload = json.dumps(asdict(baseline))

        cursor = await self._db.execute(
            """INSERT INTO regression_tests
               (name, description, correlation_id, captured_at, baseline_snapshot)
               VALUES (?, ?, ?, ?, ?)""",
            (name.strip(), description.strip(), correlation_id.strip(), now_iso, payload),
        )
        await self._db.commit()
        row_id = int(cursor.lastrowid or 0)

        return RegressionTest(
            id=row_id,
            name=name.strip(),
            description=description.strip(),
            correlation_id=correlation_id.strip(),
            captured_at=now_iso,
            baseline=asdict(baseline),
            archived=False,
            last_replay_correlation_id=None,
            last_replay_status=None,
            last_replay_at=None,
            last_replay_diff=None,
        )

    async def delete(self, regression_id: int) -> bool:
        cursor = await self._db.execute(
            "DELETE FROM regression_tests WHERE id = ?", (regression_id,),
        )
        await self._db.commit()
        return (cursor.rowcount or 0) > 0

    async def set_archived(self, regression_id: int, archived: bool) -> bool:
        cursor = await self._db.execute(
            "UPDATE regression_tests SET archived = ? WHERE id = ?",
            (1 if archived else 0, regression_id),
        )
        await self._db.commit()
        return (cursor.rowcount or 0) > 0

    # ─── Reads ───

    async def list(
        self,
        *,
        archived: bool = False,
        limit: int = 50,
        offset: int = 0,
    ) -> list[RegressionTest]:
        cursor = await self._db.execute(
            """SELECT id, name, description, correlation_id, captured_at,
                      baseline_snapshot, archived, last_replay_correlation_id,
                      last_replay_status, last_replay_at, last_replay_diff
               FROM regression_tests
               WHERE archived = ?
               ORDER BY captured_at DESC
               LIMIT ? OFFSET ?""",
            (1 if archived else 0, limit, offset),
        )
        rows = await cursor.fetchall()
        return [self._row_to_test(r) for r in rows]

    async def get(self, regression_id: int) -> RegressionTest | None:
        cursor = await self._db.execute(
            """SELECT id, name, description, correlation_id, captured_at,
                      baseline_snapshot, archived, last_replay_correlation_id,
                      last_replay_status, last_replay_at, last_replay_diff
               FROM regression_tests WHERE id = ?""",
            (regression_id,),
        )
        row = await cursor.fetchone()
        return self._row_to_test(row) if row else None

    async def count(self, *, archived: bool = False) -> int:
        cursor = await self._db.execute(
            "SELECT COUNT(*) FROM regression_tests WHERE archived = ?",
            (1 if archived else 0,),
        )
        row = await cursor.fetchone()
        return int(row[0]) if row else 0

    def _row_to_test(self, row: Any) -> RegressionTest:
        try:
            baseline = json.loads(row[5]) if row[5] else {}
        except Exception:
            baseline = {}
        try:
            diff = json.loads(row[10]) if row[10] else None
        except Exception:
            diff = None
        return RegressionTest(
            id=int(row[0]),
            name=str(row[1]),
            description=str(row[2] or ""),
            correlation_id=str(row[3]),
            captured_at=str(row[4]),
            baseline=baseline,
            archived=bool(row[6]),
            last_replay_correlation_id=row[7],
            last_replay_status=row[8],
            last_replay_at=row[9],
            last_replay_diff=diff,
        )

    # ─── Replay ───

    async def replay(self, regression_id: int) -> dict:
        """Re-run the captured regression and compare the new run to the baseline.

        MVP behavior:
          - If baseline has a user_input and we have a wired Agent, send
            the input through chat() under a new correlation scope and
            record the new correlation_id + tool-sequence diff.
          - If no user_input or no agent, mark status=manual_only so the
            user knows they need to retrigger by hand.

        Never raises — every failure is persisted as status=failed with
        a human-readable reason so the regression history stays
        browseable.
        """
        test = await self.get(regression_id)
        if test is None:
            return {"status": "not_found", "regression_id": regression_id}

        baseline = test.baseline or {}
        user_input = str(baseline.get("user_input", "") or "")
        status: str
        new_corr: str | None = None
        diff: dict | None = None
        reason = ""

        if not user_input:
            status = "manual_only"
            reason = "baseline has no captured user_input — replay requires manual trigger"
        elif self._agent is None:
            status = "manual_only"
            reason = "agent not wired into regression service"
        else:
            try:
                from agntspace.activity.decisions import (
                    correlation_scope,
                    new_correlation_id,
                )

                new_corr = new_correlation_id("regression-replay")
                with correlation_scope(new_corr):
                    # chat() is synchronous-style on the agent (returns full reply).
                    # We invoke it, then let the regular chain build up in the
                    # decision log. Any exception is captured below.
                    chat = getattr(self._agent, "chat", None)
                    if chat is None:
                        status = "failed"
                        reason = "agent has no chat() method"
                    else:
                        await chat(user_input)
                        status = "replayed"
                        # Build the diff between baseline and replay.
                        replay_baseline = await self._extract_baseline(new_corr)
                        diff = _diff_baselines(baseline, asdict(replay_baseline))
                        if diff.get("tool_sequence_changed") or diff.get("final_status_changed"):
                            status = "drift"
            except Exception as exc:
                status = "failed"
                reason = f"{type(exc).__name__}: {exc}"

        now_iso = datetime.now(UTC).isoformat()
        await self._db.execute(
            """UPDATE regression_tests
               SET last_replay_correlation_id = ?,
                   last_replay_status = ?,
                   last_replay_at = ?,
                   last_replay_diff = ?
               WHERE id = ?""",
            (new_corr, status, now_iso, json.dumps(diff) if diff else None, regression_id),
        )
        await self._db.commit()

        return {
            "status": status,
            "regression_id": regression_id,
            "new_correlation_id": new_corr,
            "reason": reason,
            "diff": diff,
        }


def _diff_baselines(baseline: dict, replay: dict) -> dict:
    """Shallow diff of two captured baselines.

    Pure function — easy to unit test independently of DB or agent.
    """
    baseline_tools = [t.get("tool", "") for t in (baseline.get("tool_sequence") or [])]
    replay_tools = [t.get("tool", "") for t in (replay.get("tool_sequence") or [])]

    baseline_actors = set(baseline.get("actors") or [])
    replay_actors = set(replay.get("actors") or [])

    diff: dict = {
        "baseline_tool_count": len(baseline_tools),
        "replay_tool_count": len(replay_tools),
        "tool_sequence_changed": baseline_tools != replay_tools,
        "added_actors": sorted(replay_actors - baseline_actors),
        "removed_actors": sorted(baseline_actors - replay_actors),
        "baseline_status": baseline.get("final_status", ""),
        "replay_status": replay.get("final_status", ""),
        "final_status_changed": baseline.get("final_status") != replay.get("final_status"),
    }
    # Per-position tool mismatch when counts align.
    if baseline_tools and replay_tools and len(baseline_tools) == len(replay_tools):
        diff["tool_mismatches"] = [
            i for i, (a, b) in enumerate(zip(baseline_tools, replay_tools, strict=False)) if a != b
        ]
    return diff
