"""Session-level trace analysis and adaptive ranking."""

from __future__ import annotations

from collections import Counter
from dataclasses import dataclass
from typing import Any

from agent_debugger_sdk.core.events import Checkpoint
from agent_debugger_sdk.core.events import EventType
from agent_debugger_sdk.core.events import TraceEvent


def _event_value(event: TraceEvent | None, key: str, default: Any = None) -> Any:
    if event is None:
        return default
    if hasattr(event, key):
        return getattr(event, key)
    return event.data.get(key, default)


def _mean(values: list[float]) -> float:
    if not values:
        return 0.0
    return sum(values) / len(values)


@dataclass
class TraceIntelligence:
    """Compute replay-centric analysis from session events."""

    severity_weights: dict[EventType, float] | None = None

    def __post_init__(self) -> None:
        if self.severity_weights is None:
            self.severity_weights = {
                EventType.ERROR: 1.0,
                EventType.POLICY_VIOLATION: 0.96,
                EventType.REFUSAL: 0.92,
                EventType.BEHAVIOR_ALERT: 0.88,
                EventType.SAFETY_CHECK: 0.8,
                EventType.DECISION: 0.72,
                EventType.CHECKPOINT: 0.65,
                EventType.TOOL_RESULT: 0.58,
                EventType.LLM_RESPONSE: 0.52,
                EventType.PROMPT_POLICY: 0.48,
                EventType.AGENT_TURN: 0.44,
                EventType.TOOL_CALL: 0.4,
                EventType.LLM_REQUEST: 0.35,
                EventType.AGENT_START: 0.2,
                EventType.AGENT_END: 0.2,
            }

    def fingerprint(self, event: TraceEvent) -> str:
        """Return a coarse fingerprint used for recurrence clustering."""
        match event.event_type:
            case EventType.ERROR:
                return f"error:{_event_value(event, 'error_type', 'unknown')}:{_event_value(event, 'error_message', '')}"
            case EventType.TOOL_RESULT:
                return f"tool:{_event_value(event, 'tool_name', 'unknown')}:{bool(_event_value(event, 'error'))}"
            case EventType.REFUSAL:
                return f"refusal:{_event_value(event, 'policy_name', 'unknown')}:{_event_value(event, 'risk_level', 'medium')}"
            case EventType.POLICY_VIOLATION:
                return f"policy:{_event_value(event, 'policy_name', 'unknown')}:{_event_value(event, 'violation_type', 'unknown')}"
            case EventType.BEHAVIOR_ALERT:
                return f"alert:{_event_value(event, 'alert_type', 'unknown')}"
            case EventType.SAFETY_CHECK:
                return f"safety:{_event_value(event, 'policy_name', 'unknown')}:{_event_value(event, 'outcome', 'pass')}"
            case EventType.DECISION:
                return f"decision:{_event_value(event, 'chosen_action', 'unknown')}"
            case _:
                return f"{event.event_type}:{event.name}"

    def severity(self, event: TraceEvent) -> float:
        """Compute an event severity score."""
        severity = self.severity_weights.get(event.event_type, 0.3)
        if event.event_type == EventType.DECISION:
            confidence = float(_event_value(event, "confidence", 0.5) or 0.5)
            evidence = _event_value(event, "evidence", []) or []
            severity += (1 - confidence) * 0.25
            if not evidence:
                severity += 0.08
        if event.event_type == EventType.TOOL_RESULT and _event_value(event, "error"):
            severity += 0.28
        if event.event_type == EventType.SAFETY_CHECK and _event_value(event, "outcome", "pass") != "pass":
            severity += 0.15
        if event.event_type == EventType.LLM_RESPONSE:
            severity += min(float(_event_value(event, "cost_usd", 0.0) or 0.0) / 0.05, 0.12)
        return min(severity, 1.0)

    def retention_tier(
        self,
        *,
        replay_value: float,
        high_severity_count: int,
        failure_cluster_count: int,
        behavior_alert_count: int,
    ) -> str:
        """Assign a coarse retention tier for a session or checkpoint."""
        if replay_value >= 0.72 or high_severity_count > 0 or failure_cluster_count >= 2:
            return "full"
        if replay_value >= 0.42 or behavior_alert_count > 0 or failure_cluster_count > 0:
            return "summarized"
        return "downsampled"

    def build_live_summary(self, events: list[TraceEvent], checkpoints: list[Checkpoint]) -> dict[str, Any]:
        """Build a live monitoring summary from the current persisted session state."""
        if not events:
            return {
                "event_count": 0,
                "checkpoint_count": len(checkpoints),
                "latest": {
                    "decision_event_id": None,
                    "tool_event_id": None,
                    "safety_event_id": None,
                    "turn_event_id": None,
                    "policy_event_id": None,
                    "checkpoint_id": checkpoints[-1].id if checkpoints else None,
                },
                "rolling_summary": "Awaiting richer live summaries",
                "recent_alerts": [],
            }

        latest_decision = next((event for event in reversed(events) if event.event_type == EventType.DECISION), None)
        latest_tool = next(
            (
                event
                for event in reversed(events)
                if event.event_type in {EventType.TOOL_CALL, EventType.TOOL_RESULT}
            ),
            None,
        )
        latest_safety = next(
            (
                event
                for event in reversed(events)
                if event.event_type in {EventType.SAFETY_CHECK, EventType.REFUSAL, EventType.POLICY_VIOLATION}
            ),
            None,
        )
        latest_turn = next((event for event in reversed(events) if event.event_type == EventType.AGENT_TURN), None)
        latest_policy = next((event for event in reversed(events) if event.event_type == EventType.PROMPT_POLICY), None)

        recent_events = events[-12:]
        recent_alerts: list[dict[str, Any]] = [
            {
                "alert_type": _event_value(event, "alert_type", "behavior_alert"),
                "severity": _event_value(event, "severity", "medium"),
                "signal": _event_value(event, "signal", event.name),
                "event_id": event.id,
                "source": "captured",
            }
            for event in recent_events
            if event.event_type == EventType.BEHAVIOR_ALERT
        ]

        recent_tool_calls = [event for event in recent_events if event.event_type == EventType.TOOL_CALL]
        last_three_tool_calls = recent_tool_calls[-3:]
        if len(last_three_tool_calls) == 3:
            tool_name = _event_value(last_three_tool_calls[-1], "tool_name", "")
            if tool_name and all(_event_value(event, "tool_name", "") == tool_name for event in last_three_tool_calls):
                recent_alerts.append(
                    {
                        "alert_type": "tool_loop",
                        "severity": "high",
                        "signal": f"Three consecutive calls to {tool_name}",
                        "event_id": last_three_tool_calls[-1].id,
                        "source": "derived",
                    }
                )

        recent_guardrails = [
            event
            for event in recent_events
            if (
                event.event_type == EventType.REFUSAL
                or event.event_type == EventType.POLICY_VIOLATION
                or (
                    event.event_type == EventType.SAFETY_CHECK
                    and _event_value(event, "outcome", "pass") != "pass"
                )
            )
        ]
        if len(recent_guardrails) >= 2:
            recent_alerts.append(
                {
                    "alert_type": "guardrail_pressure",
                    "severity": "high" if len(recent_guardrails) >= 3 else "medium",
                    "signal": f"{len(recent_guardrails)} recent blocked or warned actions",
                    "event_id": recent_guardrails[-1].id,
                    "source": "derived",
                }
            )

        recent_policies = [event for event in recent_events if event.event_type == EventType.PROMPT_POLICY]
        unique_policies = {
            _event_value(event, "template_id", event.name)
            for event in recent_policies
            if _event_value(event, "template_id", event.name)
        }
        if len(unique_policies) >= 2:
            recent_alerts.append(
                {
                    "alert_type": "policy_shift",
                    "severity": "medium",
                    "signal": f"{len(unique_policies)} prompt policies active in the recent window",
                    "event_id": recent_policies[-1].id,
                    "source": "derived",
                }
            )

        recent_decisions = [event for event in recent_events if event.event_type == EventType.DECISION]
        last_two_decisions = recent_decisions[-2:]
        if len(last_two_decisions) == 2:
            previous_action = _event_value(last_two_decisions[0], "chosen_action", last_two_decisions[0].name)
            latest_action = _event_value(last_two_decisions[1], "chosen_action", last_two_decisions[1].name)
            if previous_action != latest_action:
                recent_alerts.append(
                    {
                        "alert_type": "strategy_change",
                        "severity": "medium",
                        "signal": f'Decision shifted from "{previous_action}" to "{latest_action}"',
                        "event_id": last_two_decisions[-1].id,
                        "source": "derived",
                    }
                )

        rolling_summary = (
            (_event_value(latest_turn, "state_summary", "") if latest_turn else "")
            or (_event_value(latest_policy, "state_summary", "") if latest_policy else "")
            or (_event_value(latest_decision, "reasoning", "") if latest_decision else "")
            or (recent_alerts[-1]["signal"] if recent_alerts else "Awaiting richer live summaries")
        )

        return {
            "event_count": len(events),
            "checkpoint_count": len(checkpoints),
            "latest": {
                "decision_event_id": latest_decision.id if latest_decision else None,
                "tool_event_id": latest_tool.id if latest_tool else None,
                "safety_event_id": latest_safety.id if latest_safety else None,
                "turn_event_id": latest_turn.id if latest_turn else None,
                "policy_event_id": latest_policy.id if latest_policy else None,
                "checkpoint_id": checkpoints[-1].id if checkpoints else None,
            },
            "rolling_summary": rolling_summary,
            "recent_alerts": recent_alerts[-6:],
        }

    def analyze_session(self, events: list[TraceEvent], checkpoints: list[Checkpoint]) -> dict[str, Any]:
        """Analyze session events for replay, clustering, and anomaly signals."""
        if not events:
            return {
                "event_rankings": [],
                "failure_clusters": [],
                "representative_failure_ids": [],
                "high_replay_value_ids": [],
                "behavior_alerts": [],
                "checkpoint_rankings": [],
                "session_replay_value": 0.0,
                "retention_tier": "downsampled",
                "session_summary": {
                    "failure_count": 0,
                    "behavior_alert_count": 0,
                    "high_severity_count": 0,
                    "checkpoint_count": 0,
                },
                "live_summary": self.build_live_summary(events, checkpoints),
            }

        fingerprints = [self.fingerprint(event) for event in events]
        counts = Counter(fingerprints)
        checkpoint_event_ids = {checkpoint.event_id for checkpoint in checkpoints}
        event_rankings: list[dict[str, Any]] = []

        consecutive_tool_loop = 0
        previous_tool_name = None
        behavior_alerts: list[dict[str, Any]] = []

        for index, event in enumerate(events):
            fingerprint = fingerprints[index]
            severity = self.severity(event)
            recurrence_count = counts[fingerprint]
            recurrence = min((recurrence_count - 1) / max(len(events), 1), 1.0)
            novelty = 1.0 / recurrence_count
            replay_value = severity * 0.55
            replay_value += 0.15 if event.id in checkpoint_event_ids else 0.0
            replay_value += 0.1 if event.event_type in {EventType.DECISION, EventType.REFUSAL, EventType.POLICY_VIOLATION} else 0.0
            replay_value += 0.1 if bool(_event_value(event, "upstream_event_ids", getattr(event, "upstream_event_ids", []))) else 0.0
            replay_value += 0.1 if bool(_event_value(event, "evidence_event_ids", [])) else 0.0
            composite = min(1.0, severity * 0.45 + novelty * 0.2 + recurrence * 0.15 + replay_value * 0.2)

            event_rankings.append(
                {
                    "event_id": event.id,
                    "event_type": str(event.event_type),
                    "fingerprint": fingerprint,
                    "severity": round(severity, 4),
                    "novelty": round(novelty, 4),
                    "recurrence": round(recurrence, 4),
                    "replay_value": round(min(replay_value, 1.0), 4),
                    "composite": round(composite, 4),
                }
            )

            if event.event_type == EventType.TOOL_CALL:
                tool_name = _event_value(event, "tool_name", "")
                if tool_name and tool_name == previous_tool_name:
                    consecutive_tool_loop += 1
                else:
                    consecutive_tool_loop = 1
                previous_tool_name = tool_name
                if consecutive_tool_loop >= 3:
                    behavior_alerts.append(
                        {
                            "alert_type": "tool_loop",
                            "severity": "high",
                            "signal": f"Repeated tool loop for {tool_name}",
                            "event_id": event.id,
                        }
                    )
            else:
                previous_tool_name = None
                consecutive_tool_loop = 0

        clusters: dict[str, dict[str, Any]] = {}
        for ranking in event_rankings:
            if ranking["severity"] < 0.78:
                continue
            cluster = clusters.setdefault(
                ranking["fingerprint"],
                {
                    "fingerprint": ranking["fingerprint"],
                    "count": 0,
                    "event_ids": [],
                    "representative_event_id": ranking["event_id"],
                    "max_composite": ranking["composite"],
                },
            )
            cluster["count"] += 1
            cluster["event_ids"].append(ranking["event_id"])
            if ranking["composite"] > cluster["max_composite"]:
                cluster["max_composite"] = ranking["composite"]
                cluster["representative_event_id"] = ranking["event_id"]

        failure_clusters = sorted(clusters.values(), key=lambda item: (-item["count"], -item["max_composite"]))
        representative_failure_ids = [cluster["representative_event_id"] for cluster in failure_clusters]
        high_replay_value_ids = [
            ranking["event_id"]
            for ranking in sorted(event_rankings, key=lambda item: item["composite"], reverse=True)[:12]
        ]
        ranking_by_event_id = {ranking["event_id"]: ranking for ranking in event_rankings}
        checkpoint_rankings: list[dict[str, Any]] = []
        total_cost = sum(
            float(_event_value(event, "cost_usd", 0.0) or 0.0)
            for event in events
            if event.event_type == EventType.LLM_RESPONSE
        )
        high_severity_count = sum(1 for ranking in event_rankings if ranking["severity"] >= 0.9)
        top_composites = [ranking["composite"] for ranking in sorted(event_rankings, key=lambda item: item["composite"], reverse=True)[:5]]
        checkpoint_values: list[float] = []

        max_sequence = max((checkpoint.sequence for checkpoint in checkpoints), default=0)
        for checkpoint in checkpoints:
            event_ranking = ranking_by_event_id.get(checkpoint.event_id)
            event_replay = float(event_ranking["replay_value"]) if event_ranking else 0.0
            event_composite = float(event_ranking["composite"]) if event_ranking else 0.0
            sequence_weight = checkpoint.sequence / max(max_sequence, 1)
            restore_value = min(1.0, event_replay * 0.45 + event_composite * 0.2 + checkpoint.importance * 0.2 + sequence_weight * 0.15)
            checkpoint_values.append(restore_value)
            checkpoint_rankings.append(
                {
                    "checkpoint_id": checkpoint.id,
                    "event_id": checkpoint.event_id,
                    "sequence": checkpoint.sequence,
                    "importance": round(checkpoint.importance, 4),
                    "replay_value": round(event_replay, 4),
                    "restore_value": round(restore_value, 4),
                    "retention_tier": self.retention_tier(
                        replay_value=restore_value,
                        high_severity_count=1 if event_ranking and event_ranking["severity"] >= 0.92 else 0,
                        failure_cluster_count=1 if checkpoint.event_id in representative_failure_ids else 0,
                        behavior_alert_count=0,
                    ),
                }
            )

        checkpoint_rankings.sort(key=lambda item: (-item["restore_value"], -item["importance"], -item["sequence"]))
        session_replay_value = min(
            1.0,
            _mean(top_composites) * 0.55
            + min(len(representative_failure_ids) / 4, 1.0) * 0.2
            + min(len(behavior_alerts) / 3, 1.0) * 0.1
            + _mean(checkpoint_values) * 0.1
            + min(total_cost / 0.25, 1.0) * 0.05,
        )
        retention_tier = self.retention_tier(
            replay_value=session_replay_value,
            high_severity_count=high_severity_count,
            failure_cluster_count=len(failure_clusters),
            behavior_alert_count=len(behavior_alerts),
        )

        return {
            "event_rankings": event_rankings,
            "failure_clusters": failure_clusters,
            "representative_failure_ids": representative_failure_ids,
            "high_replay_value_ids": high_replay_value_ids,
            "behavior_alerts": behavior_alerts,
            "checkpoint_rankings": checkpoint_rankings,
            "session_replay_value": round(session_replay_value, 4),
            "retention_tier": retention_tier,
            "session_summary": {
                "failure_count": len(representative_failure_ids),
                "behavior_alert_count": len(behavior_alerts),
                "high_severity_count": high_severity_count,
                "checkpoint_count": len(checkpoints),
            },
            "live_summary": self.build_live_summary(events, checkpoints),
        }
