"""MCP trace data models for capturing and replaying agent tool-call sessions."""

from __future__ import annotations

import json
from datetime import datetime
from pathlib import Path
from typing import Any

from pydantic import BaseModel, field_validator


class MCPToolCall(BaseModel):
    """A single MCP tool invocation with its input, output, and metadata.

    Attributes:
        tool: The tool name that was called (e.g. "search_diagnosis").
        arguments: The arguments passed to the tool.
        result: The tool's response, if available.
        error: Error message if the tool call failed.
        latency_ms: Time taken for the tool call in milliseconds.
        timestamp: When the call was made.
        server: MCP server name (e.g. "di-agentic-service").
    """

    tool: str
    arguments: dict[str, Any] = {}
    result: dict[str, Any] | list[Any] | str | None = None
    error: str | None = None
    latency_ms: float = 0.0
    timestamp: datetime | None = None
    server: str | None = None

    @field_validator("tool", mode="before")
    @classmethod
    def tool_not_empty(cls, v: Any) -> str:
        if not v or not isinstance(v, str) or not v.strip():
            raise ValueError("tool name must be a non-empty string")
        return v.strip()

    @property
    def succeeded(self) -> bool:
        """Whether this tool call completed without error."""
        return self.error is None

    @property
    def has_result(self) -> bool:
        """Whether this tool call produced a result."""
        return self.result is not None

    def to_context_dict(self) -> dict[str, Any]:
        """Convert to the dict format expected by Context.tool_calls.

        Returns:
            Dict with 'name', 'arguments', and optionally 'result'/'error'.
        """
        d: dict[str, Any] = {
            "name": self.tool,
            "arguments": self.arguments,
        }
        if self.result is not None:
            d["result"] = self.result
        if self.error is not None:
            d["error"] = self.error
        return d


class MCPTrace(BaseModel):
    """A recorded sequence of MCP tool calls from an agent session.

    Attributes:
        calls: Ordered list of tool calls made during the session.
        user_input: The original user query that triggered the agent.
        agent_output: The agent's final response to the user.
        session_id: Unique identifier for this trace session.
        metadata: Arbitrary metadata (model, temperature, etc.).
    """

    calls: list[MCPToolCall] = []
    user_input: str | None = None
    agent_output: str | None = None
    session_id: str | None = None
    metadata: dict[str, Any] = {}

    @property
    def tool_names(self) -> list[str]:
        """Ordered list of tool names called in this trace."""
        return [c.tool for c in self.calls]

    @property
    def unique_tools(self) -> set[str]:
        """Set of unique tool names called."""
        return set(self.tool_names)

    @property
    def total_latency_ms(self) -> float:
        """Sum of all tool call latencies."""
        return sum(c.latency_ms for c in self.calls)

    @property
    def total_cost_usd(self) -> float:
        """Sum of cost_usd from tool results that report it."""
        total = 0.0
        for c in self.calls:
            if isinstance(c.result, dict):
                total += c.result.get("cost_usd", 0.0)
        return total

    @property
    def failed_calls(self) -> list[MCPToolCall]:
        """Tool calls that returned errors."""
        return [c for c in self.calls if not c.succeeded]

    @property
    def succeeded_calls(self) -> list[MCPToolCall]:
        """Tool calls that completed successfully."""
        return [c for c in self.calls if c.succeeded]

    def to_tool_calls(self) -> list[dict[str, Any]]:
        """Convert to Context.tool_calls format.

        Returns:
            List of dicts compatible with Context(tool_calls=...).
        """
        return [c.to_context_dict() for c in self.calls]

    def get_calls_by_tool(self, tool_name: str) -> list[MCPToolCall]:
        """Get all calls to a specific tool.

        Args:
            tool_name: Tool name to filter by.

        Returns:
            List of matching tool calls.
        """
        return [c for c in self.calls if c.tool == tool_name]

    def get_call_by_index(self, index: int) -> MCPToolCall | None:
        """Get a tool call by its position in the trace.

        Args:
            index: Zero-based index.

        Returns:
            MCPToolCall or None if index is out of range.
        """
        if 0 <= index < len(self.calls):
            return self.calls[index]
        return None

    @classmethod
    def from_jsonl(cls, path: str | Path) -> MCPTrace:
        """Load a trace from a JSONL file.

        Each line should be a JSON object with at least a "tool" key.
        Optional keys: "arguments", "result", "error", "latency_ms",
        "timestamp", "server".

        A line with "user_input" or "agent_output" sets trace-level fields.

        Args:
            path: Path to the JSONL file.

        Returns:
            Populated MCPTrace.

        Raises:
            FileNotFoundError: If the file doesn't exist.
            ValueError: If a line can't be parsed.
        """
        filepath = Path(path)
        if not filepath.exists():
            raise FileNotFoundError(f"Trace file not found: {path}")

        calls: list[MCPToolCall] = []
        user_input: str | None = None
        agent_output: str | None = None
        session_id: str | None = None
        metadata: dict[str, Any] = {}

        for line_num, line in enumerate(filepath.read_text().splitlines(), 1):
            line = line.strip()
            if not line:
                continue
            try:
                data = json.loads(line)
            except json.JSONDecodeError as e:
                raise ValueError(
                    f"Invalid JSON on line {line_num} of {path}: {e}"
                ) from e

            if "user_input" in data:
                user_input = data["user_input"]
                continue
            if "agent_output" in data:
                agent_output = data["agent_output"]
                continue
            if "session_id" in data and "tool" not in data:
                session_id = data["session_id"]
                metadata.update(data.get("metadata", {}))
                continue

            if "tool" in data:
                calls.append(MCPToolCall(**data))

        return cls(
            calls=calls,
            user_input=user_input,
            agent_output=agent_output,
            session_id=session_id,
            metadata=metadata,
        )

    @classmethod
    def from_list(cls, data: list[dict[str, Any]]) -> MCPTrace:
        """Create a trace from a list of tool-call dicts.

        Args:
            data: List of dicts, each with at least "tool" (or "name") key.

        Returns:
            Populated MCPTrace.
        """
        calls: list[MCPToolCall] = []
        for item in data:
            if "name" in item and "tool" not in item:
                item = {**item, "tool": item.pop("name")}
            calls.append(MCPToolCall(**item))
        return cls(calls=calls)

    def to_jsonl(self, path: str | Path) -> None:
        """Save the trace to a JSONL file.

        Args:
            path: Output file path.
        """
        filepath = Path(path)
        lines: list[str] = []

        if self.session_id or self.metadata:
            header: dict[str, Any] = {}
            if self.session_id:
                header["session_id"] = self.session_id
            if self.metadata:
                header["metadata"] = self.metadata
            lines.append(json.dumps(header))

        if self.user_input is not None:
            lines.append(json.dumps({"user_input": self.user_input}))

        for call in self.calls:
            d = call.model_dump(exclude_none=True)
            if d.get("timestamp"):
                d["timestamp"] = d["timestamp"].isoformat()
            lines.append(json.dumps(d))

        if self.agent_output is not None:
            lines.append(json.dumps({"agent_output": self.agent_output}))

        filepath.write_text("\n".join(lines) + "\n")
