"""RCA state tracking and persistence for MCP/debug workflows."""

from __future__ import annotations

import hashlib
import json
import os
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Literal, cast

from filelock import FileLock

from ._debug_session import coerce_crash_frame_index, normalize_snapshot_view
from .gates import evaluate_gate, normalize_gate_mode
from .rca_store import (
    has_any_records,
    history,
    latest_record,
    store_rca_record,
    store_rca_records_bulk,
)

RCAState = Literal[
    "captured",
    "evidence_bound",
    "hypothesized",
    "planned",
    "changed",
    "verified",
    "resolved",
    "stalled",
]
RCAAction = Literal[
    "capture",
    "bind_evidence",
    "hypothesize",
    "plan",
    "change",
    "verify",
    "resolve",
]
RCABackend = Literal["sqlite", "jsonl"]

_RCA_FILE = "rca_attempts.jsonl"
_RCA_LOCK_FILE = ".llmdebug.rca.lock"
_RCA_BACKEND_ENV = "LLMDEBUG_RCA_BACKEND"
_DEFAULT_BACKEND: RCABackend = "sqlite"
_DEFAULT_MAX_RECORDS = 5000
_DEFAULT_STATE: RCAState = "captured"
_VALID_ACTIONS: set[str] = {
    "capture",
    "bind_evidence",
    "hypothesize",
    "plan",
    "change",
    "verify",
    "resolve",
}
_STATE_TRANSITIONS: dict[str, RCAState] = {
    "capture": "captured",
    "bind_evidence": "evidence_bound",
    "hypothesize": "hypothesized",
    "plan": "planned",
    "change": "changed",
    "verify": "verified",
    "resolve": "resolved",
}


def _now_utc() -> str:
    return datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")


def _store_path(out_dir: str) -> Path:
    return Path(out_dir) / _RCA_FILE


@contextmanager
def _store_lock(path: Path) -> Any:
    """Serialize RCA writes and migration operations."""
    lock = FileLock(path.parent / _RCA_LOCK_FILE)
    with lock:
        yield


def _backend() -> RCABackend:
    raw = os.getenv(_RCA_BACKEND_ENV, _DEFAULT_BACKEND).strip().lower()
    if raw in {"sqlite", "jsonl"}:
        return cast(RCABackend, raw)
    return _DEFAULT_BACKEND


def _normalize_action(action: str) -> RCAAction:
    lowered = action.strip().lower()
    if lowered in _VALID_ACTIONS:
        return cast(RCAAction, lowered)
    import sys

    sys.stderr.write(
        f"llmdebug: warning: unrecognized RCA action {action!r}, "
        f"defaulting to 'bind_evidence'. Valid: {sorted(_VALID_ACTIONS)}\n"
    )
    return "bind_evidence"


def _iter_jsonl_records(path: Path) -> list[dict[str, Any]]:
    if not path.exists():
        return []
    records: list[dict[str, Any]] = []
    try:
        raw = path.read_text(encoding="utf-8")
    except OSError:
        return []
    for line in raw.splitlines():
        item = line.strip()
        if not item:
            continue
        try:
            data = json.loads(item)
        except json.JSONDecodeError:
            continue
        if isinstance(data, dict):
            records.append(cast(dict[str, Any], data))
    return records


def _iter_jsonl_lines_reverse(path: Path, *, block_size: int = 8192) -> Iterator[str]:
    """Yield JSONL lines in reverse order without loading the whole file."""
    with path.open("rb") as handle:
        handle.seek(0, os.SEEK_END)
        position = handle.tell()
        pending = b""

        while position > 0:
            read_size = min(block_size, position)
            position -= read_size
            handle.seek(position)
            chunk = handle.read(read_size)
            pending = chunk + pending
            lines = pending.split(b"\n")
            pending = lines[0]
            for raw_line in reversed(lines[1:]):
                line = raw_line.strip()
                if line:
                    yield line.decode("utf-8", errors="replace")

        if pending.strip():
            yield pending.strip().decode("utf-8", errors="replace")


def _latest_jsonl_record(path: Path, *, session_id: str | None = None) -> dict[str, Any] | None:
    if not path.exists():
        return None
    try:
        for line in _iter_jsonl_lines_reverse(path):
            try:
                data = json.loads(line)
            except json.JSONDecodeError:
                continue
            if not isinstance(data, dict):
                continue
            data_dict = cast(dict[str, Any], data)
            if session_id and data_dict.get("session_id") != session_id:
                continue
            return data_dict
    except OSError:
        return None
    return None


def _append_record(path: Path, record: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with _store_lock(path):
        with path.open("a", encoding="utf-8") as handle:
            handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")

        max_records = _max_rca_records()
        if max_records <= 0:
            return

        try:
            lines = path.read_text(encoding="utf-8").splitlines()
        except OSError:
            return
        if len(lines) <= max_records:
            return

        trimmed = lines[-max_records:]
        try:
            path.write_text("\n".join(trimmed) + "\n", encoding="utf-8")
        except OSError:
            return


def _max_rca_records() -> int:
    raw_value = os.getenv("LLMDEBUG_RCA_MAX_RECORDS", str(_DEFAULT_MAX_RECORDS))
    try:
        value = int(raw_value)
    except ValueError:
        return _DEFAULT_MAX_RECORDS
    return max(value, 0)


def _ensure_sqlite_migration(out_dir: str) -> None:
    jsonl_path = _store_path(out_dir)
    if not jsonl_path.exists():
        return

    with _store_lock(jsonl_path):
        if not jsonl_path.exists():
            return
        if has_any_records(out_dir=out_dir):
            return

        records = _iter_jsonl_records(jsonl_path)
        if records:
            store_rca_records_bulk(records, out_dir=out_dir)

        backup_path = Path(f"{jsonl_path}.bak")
        try:
            backup_path.unlink(missing_ok=True)
        except OSError:
            pass
        try:
            jsonl_path.replace(backup_path)
        except OSError:
            # Keep original JSONL if backup rotation fails.
            pass


def _classify_hypothesis_delta(previous: str, current: str) -> str:
    prev = previous.strip()
    curr = current.strip()
    if not prev and not curr:
        return "unknown"
    if prev == curr:
        return "unchanged"
    if prev and curr and prev in curr:
        return "refined"
    if not prev and curr:
        return "added"
    if prev and not curr:
        return "removed"
    return "replaced"


def allowed_actions_for_state(state: str) -> list[str]:
    """Return sorted allowed action list for display/UI guidance."""
    matrix: dict[str, list[str]] = {
        "captured": ["bind_evidence", "capture"],
        "evidence_bound": ["hypothesize", "bind_evidence"],
        "hypothesized": ["plan", "bind_evidence", "hypothesize"],
        "planned": ["change", "plan", "bind_evidence"],
        "changed": ["verify", "change", "bind_evidence"],
        "verified": ["resolve", "bind_evidence", "plan"],
        "resolved": ["capture", "bind_evidence"],
        "stalled": ["bind_evidence", "hypothesize", "plan"],
    }
    return matrix.get(state, ["bind_evidence"])


def _default_next_required(state: str) -> list[str]:
    steps: dict[str, list[str]] = {
        "captured": ["Bind concrete evidence from traceback/frames."],
        "evidence_bound": ["State a falsifiable root-cause hypothesis."],
        "hypothesized": ["Write a minimal experiment or patch plan."],
        "planned": ["Apply one focused change tied to the plan."],
        "changed": ["Run verification checks and capture result."],
        "verified": ["Resolve if checks passed, else return to evidence."],
        "stalled": ["Gather new evidence or revise hypothesis before retry."],
    }
    return steps.get(state, [])


def _signature_from_snapshot(view: dict[str, Any]) -> str:
    """Build failure signature from a normalized snapshot view."""
    exc = view.get("exception", {})
    exc_type = str(exc.get("type") or "Exception")
    exc_message = str(exc.get("message") or "")
    frames_raw = view.get("frames")
    frames: list[dict[str, Any]] = (
        cast(list[dict[str, Any]], frames_raw) if isinstance(frames_raw, list) else []
    )
    crash_frame: dict[str, Any] = {}
    if frames:
        crash_idx = coerce_crash_frame_index(view)
        candidate = frames[crash_idx] if 0 <= crash_idx < len(frames) else frames[0]
        if isinstance(candidate, dict):  # pyright: ignore[reportUnnecessaryIsInstance] — defense-in-depth for malformed data
            crash_frame = candidate
    location = {
        "file": str(crash_frame.get("file_rel") or crash_frame.get("file") or "?"),
        "line": int(crash_frame.get("line") or 0),
        "function": str(crash_frame.get("function") or "?"),
    }
    nodeid = ""
    pytest_block = view.get("pytest")
    if isinstance(pytest_block, dict):
        nodeid = str(cast(dict[str, Any], pytest_block).get("nodeid") or "")
    payload = {
        "exception_type": exc_type,
        "exception_message": exc_message,
        "location": location,
        "nodeid": nodeid,
    }
    digest = hashlib.sha256(
        json.dumps(payload, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode(
            "utf-8"
        )
    ).hexdigest()[:16]
    return f"{exc_type}:{digest}"


def _signature_from_text(stdout: str, stderr: str) -> str:
    """Build failure signature from raw text output."""
    text = (stderr or "").strip() or (stdout or "").strip() or "unknown_failure"
    lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
    head = "\n".join(lines[-3:]) if lines else text
    digest = hashlib.sha256(head.encode("utf-8")).hexdigest()[:16]
    return f"TextFailure:{digest}"


def build_failure_signature(
    snapshot: dict[str, Any] | None,
    *,
    stdout: str = "",
    stderr: str = "",
) -> str:
    """Build a stable failure signature from snapshot (preferred) or textual output."""
    if snapshot:
        return _signature_from_snapshot(normalize_snapshot_view(snapshot))
    return _signature_from_text(stdout, stderr)


def derive_session_id(
    snapshot: dict[str, Any] | None,
    *,
    out_dir: str,
    fallback: str = "global",
) -> str:
    """Derive stable session id from out_dir + pytest nodeid/name when available."""
    resolved = str(Path(out_dir).resolve())
    scope = fallback
    if snapshot:
        view = normalize_snapshot_view(snapshot)
        pytest_block = view.get("pytest")
        if isinstance(pytest_block, dict):
            pytest_dict = cast(dict[str, Any], pytest_block)
            raw_nodeid = pytest_dict.get("nodeid")
            if isinstance(raw_nodeid, str):
                nodeid = raw_nodeid.strip()
                if nodeid:
                    scope = nodeid
        if scope == fallback and isinstance(view.get("name"), str):
            name = str(view.get("name")).strip()
            if name:
                scope = name
    return f"{resolved}::{scope}"


def get_rca_history(
    *,
    out_dir: str,
    session_id: str | None = None,
    limit: int = 20,
) -> list[dict[str, Any]]:
    """Load RCA history, newest-first."""
    if _backend() == "sqlite":
        _ensure_sqlite_migration(out_dir)
        return history(out_dir=out_dir, session_id=session_id, limit=limit)

    records = _iter_jsonl_records(_store_path(out_dir))
    if session_id:
        records = [record for record in records if record.get("session_id") == session_id]
    records = list(reversed(records))
    if limit > 0:
        records = records[:limit]
    return records


def get_rca_status(*, out_dir: str, session_id: str | None = None) -> dict[str, Any] | None:
    """Return latest RCA attempt for a session."""
    if _backend() == "sqlite":
        _ensure_sqlite_migration(out_dir)
        return latest_record(out_dir=out_dir, session_id=session_id)

    return _latest_jsonl_record(_store_path(out_dir), session_id=session_id)


def _merge_evidence_refs(
    previous: list[str],
    incoming: list[str],
    *,
    max_items: int = 16,
) -> list[str]:
    merged: list[str] = []
    seen: set[str] = set()
    for value in previous + incoming:
        if not value:
            continue
        item = value.strip()
        if not item or item in seen:
            continue
        seen.add(item)
        merged.append(item)
        if len(merged) >= max_items:
            break
    return merged


def _resolve_previous_values(
    previous: dict[str, Any] | None,
) -> tuple[str, str, str, Any, Any, str]:
    previous_state = str(previous.get("state")) if previous else _DEFAULT_STATE
    previous_hypothesis = str(previous.get("hypothesis") or "") if previous else ""
    previous_plan = str(previous.get("experiment_plan") or "") if previous else ""
    previous_refs: list[Any] = (
        cast(list[Any], previous.get("evidence_refs"))
        if isinstance(previous, dict) and previous.get("evidence_refs")
        else []
    )
    previous_verification = (
        previous.get("verification_result") if isinstance(previous, dict) else None
    )
    previous_change_summary = (
        str(previous.get("proposed_change_summary") or "") if isinstance(previous, dict) else ""
    )
    return (
        previous_state,
        previous_hypothesis,
        previous_plan,
        previous_refs,
        previous_verification,
        previous_change_summary,
    )


def _resolve_merged_evidence_refs(
    previous_refs: Any,
    evidence_refs: list[str] | None,
) -> list[str]:
    incoming_refs = evidence_refs if isinstance(evidence_refs, list) else []
    return _merge_evidence_refs(
        [str(ref) for ref in cast(list[Any], previous_refs)]
        if isinstance(previous_refs, list)
        else [],
        [str(ref) for ref in incoming_refs],
    )


def _resolve_attempt_values(
    *,
    previous_hypothesis: str,
    previous_plan: str,
    previous_change_summary: str,
    previous_verification: Any,
    hypothesis: str | None,
    experiment_plan: str | None,
    proposed_change_summary: str | None,
    verification_result: dict[str, Any] | None,
) -> tuple[str, str, Any, str]:
    hypothesis_value = (hypothesis or "").strip() or previous_hypothesis
    plan_value = (experiment_plan or "").strip() or previous_plan
    verification_value = (
        verification_result if verification_result is not None else previous_verification
    )
    proposed_change_value = (proposed_change_summary or "").strip() or previous_change_summary
    return hypothesis_value, plan_value, verification_value, proposed_change_value


def _derive_attempt_progress(
    *,
    previous: dict[str, Any] | None,
    failure_signature: str,
    previous_hypothesis: str,
    hypothesis_value: str,
    previous_plan: str,
    plan_value: str,
) -> tuple[bool, int, str, bool, bool]:
    same_signature = (
        bool(previous) and str(previous.get("failure_signature") or "") == failure_signature
    )
    repeat_count = int(previous.get("repeat_count") or 1) + 1 if same_signature and previous else 1
    hypothesis_delta = _classify_hypothesis_delta(previous_hypothesis, hypothesis_value)
    plan_advanced = bool(plan_value) and plan_value != previous_plan
    hypothesis_advanced = hypothesis_delta in {"added", "refined", "replaced"}
    return same_signature, repeat_count, hypothesis_delta, plan_advanced, hypothesis_advanced


def _resolve_target_state(
    previous_state: str,
    normalized_action: RCAAction,
    gate_allowed: bool,
) -> RCAState:
    target_state = _STATE_TRANSITIONS.get(normalized_action, "evidence_bound")
    if gate_allowed:
        return target_state
    return cast(
        RCAState,
        previous_state if previous_state in _STATE_TRANSITIONS.values() else "captured",
    )


def _apply_stall_rule(
    *,
    target_state: RCAState,
    gate_warnings: list[str],
    same_signature: bool,
    repeat_count: int,
    stall_threshold: int,
    hypothesis_advanced: bool,
    plan_advanced: bool,
    normalized_action: RCAAction,
) -> tuple[RCAState, list[str]]:
    if (
        same_signature
        and repeat_count >= max(stall_threshold, 2)
        and not hypothesis_advanced
        and not plan_advanced
        and normalized_action in {"bind_evidence", "change", "verify"}
    ):
        gate_warnings.append(
            "Repeated failure signature without hypothesis/plan advancement; state marked stalled."
        )
        return "stalled", gate_warnings
    return target_state, gate_warnings


def _build_transition_history(
    *,
    previous: dict[str, Any] | None,
    now: str,
    normalized_action: RCAAction,
    previous_state: str,
    target_state: RCAState,
    gate_allowed: bool,
    gate_mode: str,
) -> list[dict[str, Any]]:
    previous_history_raw = (
        previous.get("transition_history") if isinstance(previous, dict) else cast(list[Any], [])
    )
    transition_history: list[dict[str, Any]] = []
    if isinstance(previous_history_raw, list):
        transition_history = [
            cast(dict[str, Any], item)
            for item in cast(list[Any], previous_history_raw)
            if isinstance(item, dict)
        ]
    transition_history.append(
        {
            "at": now,
            "action": normalized_action,
            "from": previous_state,
            "to": target_state,
            "gate_allowed": gate_allowed,
            "gate_mode": gate_mode,
        }
    )
    if len(transition_history) > 32:
        transition_history = transition_history[-32:]
    return transition_history


def advance_rca_attempt(
    *,
    out_dir: str,
    session_id: str,
    action: str,
    failure_signature: str,
    evidence_refs: list[str] | None = None,
    hypothesis: str | None = None,
    experiment_plan: str | None = None,
    proposed_change_summary: str | None = None,
    verification_result: dict[str, Any] | None = None,
    gate_mode: str | None = None,
    exploratory: bool = False,
    override_reason: str | None = None,
    stall_threshold: int = 3,
    exploration_budget: int = 2,
) -> dict[str, Any]:
    """Advance RCA state machine and persist attempt row."""
    normalized_mode = normalize_gate_mode(gate_mode)
    normalized_action = _normalize_action(action)
    now = _now_utc()

    previous = get_rca_status(out_dir=out_dir, session_id=session_id)
    (
        previous_state,
        previous_hypothesis,
        previous_plan,
        previous_refs,
        previous_verification,
        previous_change_summary,
    ) = _resolve_previous_values(previous)

    merged_refs = _resolve_merged_evidence_refs(previous_refs, evidence_refs)
    (
        hypothesis_value,
        plan_value,
        verification_value,
        proposed_change_value,
    ) = _resolve_attempt_values(
        previous_hypothesis=previous_hypothesis,
        previous_plan=previous_plan,
        previous_change_summary=previous_change_summary,
        previous_verification=previous_verification,
        hypothesis=hypothesis,
        experiment_plan=experiment_plan,
        proposed_change_summary=proposed_change_summary,
        verification_result=verification_result,
    )
    same_signature, repeat_count, hypothesis_delta, plan_advanced, hypothesis_advanced = (
        _derive_attempt_progress(
            previous=previous,
            failure_signature=failure_signature,
            previous_hypothesis=previous_hypothesis,
            hypothesis_value=hypothesis_value,
            previous_plan=previous_plan,
            plan_value=plan_value,
        )
    )

    gate_payload = {
        "evidence_refs": merged_refs,
        "hypothesis": hypothesis_value,
        "experiment_plan": plan_value,
        "verification_result": verification_value,
        "exploratory": exploratory,
        "override_reason": override_reason or "",
        "hypothesis_advanced": hypothesis_advanced,
    }
    gate = evaluate_gate(
        mode=normalized_mode,
        current_state=previous_state,
        action=normalized_action,
        payload=gate_payload,
        repeated_signature_count=repeat_count,
        exploration_budget=exploration_budget,
    )

    target_state = _resolve_target_state(previous_state, normalized_action, gate.allowed)
    target_state, gate_warnings = _apply_stall_rule(
        target_state=target_state,
        gate_warnings=list(gate.warnings),
        same_signature=same_signature,
        repeat_count=repeat_count,
        stall_threshold=stall_threshold,
        hypothesis_advanced=hypothesis_advanced,
        plan_advanced=plan_advanced,
        normalized_action=normalized_action,
    )
    transition_history = _build_transition_history(
        previous=previous,
        now=now,
        normalized_action=normalized_action,
        previous_state=previous_state,
        target_state=target_state,
        gate_allowed=gate.allowed,
        gate_mode=gate.mode,
    )

    attempt_id = f"{int(datetime.now(timezone.utc).timestamp() * 1000)}-{uuid.uuid4().hex[:8]}"
    created_at = str(previous.get("created_at") or now) if previous else now
    next_required = (
        list(gate.next_required) if gate.next_required else _default_next_required(target_state)
    )

    record: dict[str, Any] = {
        "attempt_id": attempt_id,
        "session_id": session_id,
        "state": target_state,
        "failure_signature": failure_signature,
        "repeat_count": repeat_count,
        "evidence_refs": merged_refs,
        "hypothesis": hypothesis_value or None,
        "hypothesis_delta": hypothesis_delta,
        "experiment_plan": plan_value or None,
        "proposed_change_summary": proposed_change_value or None,
        "verification_result": verification_value,
        "gate_feedback": gate.to_dict(),
        "gate_warnings": gate_warnings,
        "allowed_actions": allowed_actions_for_state(target_state),
        "next_required": next_required,
        "exploratory": exploratory,
        "transition_history": transition_history,
        "created_at": created_at,
        "updated_at": now,
    }

    backend = _backend()
    if backend == "sqlite":
        _ensure_sqlite_migration(out_dir)
        store_rca_record(record, out_dir=out_dir, max_records=_max_rca_records())
    else:
        _append_record(_store_path(out_dir), record)
    return record
