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

from __future__ import annotations

import hashlib
import json
import os
import uuid
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

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

_RCA_FILE = "rca_attempts.jsonl"
_RCA_LOCK_FILE = ".llmdebug.rca.lock"
_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 to avoid interleaved JSONL records."""
    lock = FileLock(path.parent / _RCA_LOCK_FILE)
    with lock:
        yield


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(data)
    return records


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 _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 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:
        view = normalize_snapshot_view(snapshot)
        exc = view.get("exception", {}) if isinstance(view, dict) else {}
        exc_type = str(exc.get("type") or "Exception")
        exc_message = str(exc.get("message") or "")
        frames = view.get("frames") or []
        crash_frame: dict[str, Any] = {}
        if isinstance(frames, list) and frames:
            crash_idx = coerce_crash_frame_index(view)
            if 0 <= crash_idx < len(frames) and isinstance(frames[crash_idx], dict):
                crash_frame = frames[crash_idx]
            elif isinstance(frames[0], dict):
                crash_frame = frames[0]
        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(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}"

    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 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) and isinstance(pytest_block.get("nodeid"), str):
            nodeid = str(pytest_block.get("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."""
    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."""
    history = get_rca_history(out_dir=out_dir, session_id=session_id, limit=1)
    if not history:
        return None
    return history[0]


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 = previous.get("evidence_refs") if isinstance(previous, dict) 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 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 = previous.get("transition_history") if isinstance(previous, dict) else []
    transition_history: list[dict[str, Any]] = []
    if isinstance(previous_history, list):
        transition_history = [item for item in previous_history 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 to JSONL."""
    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,
    }

    _append_record(_store_path(out_dir), record)
    return record
