"""Divergence detection — compare a replay execution against the original trace."""

from __future__ import annotations

from dataclasses import dataclass, field

from litmus.models import Trace, LLMCall, ToolCall


@dataclass
class StepDiff:
    """Difference found at a specific step."""

    step_index: int
    field: str
    original: str
    replayed: str


@dataclass
class TraceDiff:
    """Result of comparing a replayed trace to the original."""

    original_id: str
    replayed_id: str
    diverged: bool = False
    divergence_step: int | None = None
    step_count_match: bool = True
    final_output_match: bool = True
    diffs: list[StepDiff] = field(default_factory=list)

    @property
    def fidelity_score(self) -> float:
        """1.0 = perfect replay, 0.0 = completely different."""
        if not self.step_count_match:
            return 0.0
        if not self.diffs:
            return 1.0
        # Each diff reduces score proportionally
        total_steps = max(1, self.divergence_step or 1)
        return max(0.0, 1.0 - len(self.diffs) / (total_steps * 3))


def _extract_final_text(trace: Trace) -> str | None:
    """Get the final text output from the last LLM call."""
    for step in reversed(trace.steps):
        if isinstance(step, LLMCall) and step.response and step.response.body:
            body = step.response.body
            # Anthropic format
            content = body.get("content", [])
            for block in content:
                if isinstance(block, dict) and block.get("type") == "text":
                    return block.get("text", "")
            # OpenAI format
            choices = body.get("choices", [])
            for choice in choices:
                msg = choice.get("message", {})
                if msg.get("content"):
                    return msg["content"]
    return None


def _extract_tool_calls(step: LLMCall) -> list[tuple[str, dict]]:
    """Extract tool use blocks from an LLM response."""
    if not step.response or not step.response.body:
        return []
    content = step.response.body.get("content", [])
    calls = []
    for block in content:
        if isinstance(block, dict) and block.get("type") == "tool_use":
            calls.append((block.get("name", ""), block.get("input", {})))
    return calls


def compare(original: Trace, replayed: Trace) -> TraceDiff:
    """Compare a replayed trace against the original.

    Checks:
    - Same number of steps
    - Each step: same provider, same model, same status code
    - LLM steps: same tool calls chosen
    - Final output: same text content
    """
    result = TraceDiff(
        original_id=original.trace_id,
        replayed_id=replayed.trace_id,
    )

    # Step count
    if len(original.steps) != len(replayed.steps):
        result.diverged = True
        result.step_count_match = False
        result.divergence_step = min(len(original.steps), len(replayed.steps))
        result.diffs.append(StepDiff(
            step_index=-1,
            field="step_count",
            original=str(len(original.steps)),
            replayed=str(len(replayed.steps)),
        ))
        return result

    # Compare each step
    for i, (orig_step, replay_step) in enumerate(zip(original.steps, replayed.steps)):
        # Both should be same type
        if type(orig_step) != type(replay_step):
            result.diverged = True
            result.divergence_step = i
            result.diffs.append(StepDiff(
                step_index=i,
                field="step_type",
                original=type(orig_step).__name__,
                replayed=type(replay_step).__name__,
            ))
            continue

        if isinstance(orig_step, LLMCall) and isinstance(replay_step, LLMCall):
            # Check provider
            if orig_step.request.provider != replay_step.request.provider:
                result.diffs.append(StepDiff(
                    step_index=i, field="provider",
                    original=orig_step.request.provider,
                    replayed=replay_step.request.provider,
                ))

            # Check model
            if orig_step.request.model != replay_step.request.model:
                result.diffs.append(StepDiff(
                    step_index=i, field="model",
                    original=str(orig_step.request.model),
                    replayed=str(replay_step.request.model),
                ))

            # Check status code
            orig_status = orig_step.response.status_code if orig_step.response else None
            replay_status = replay_step.response.status_code if replay_step.response else None
            if orig_status != replay_status:
                result.diffs.append(StepDiff(
                    step_index=i, field="status_code",
                    original=str(orig_status),
                    replayed=str(replay_status),
                ))

            # Check tool calls chosen (most important for agent behavior)
            orig_tools = _extract_tool_calls(orig_step)
            replay_tools = _extract_tool_calls(replay_step)
            if orig_tools != replay_tools:
                result.diffs.append(StepDiff(
                    step_index=i, field="tool_calls",
                    original=str(orig_tools),
                    replayed=str(replay_tools),
                ))

    # Check final output
    orig_final = _extract_final_text(original)
    replay_final = _extract_final_text(replayed)
    if orig_final != replay_final:
        result.final_output_match = False
        result.diffs.append(StepDiff(
            step_index=len(original.steps) - 1,
            field="final_output",
            original=str(orig_final)[:200],
            replayed=str(replay_final)[:200],
        ))

    if result.diffs:
        result.diverged = True
        result.divergence_step = result.diffs[0].step_index

    return result
