"""MCP server for LLM trace interaction.

This server exposes tools for searching, viewing, and analyzing LLM traces.
Configure it with your API key and base URL to connect to your traces backend.

Usage:
    # Run directly
    mcp-traces

    # Or with uvx
    uvx mcp-traces

Environment variables:
    TRACES_API_KEY: Your API key for authentication (required)
    TRACES_BASE_URL: Base URL of the traces API (default: https://api.021.dev)
"""

from __future__ import annotations

import json
import os
from datetime import datetime
from typing import Annotated

from mcp.server.fastmcp import FastMCP

from mcp_traces.client import TracesClient
from mcp_traces.schemas import SpanStatus


# Initialize FastMCP server
mcp = FastMCP(
    "mcp-traces",
    instructions="""You have access to an LLM tracing system. Use these tools to:
- Search and list traces from AI agent executions
- Inspect individual traces to see the full conversation and tool calls
- Analyze span details to understand LLM inputs/outputs
- Get analytics on performance, errors, and token usage

When debugging agent issues:
1. Start with list_traces to find relevant traces (filter by status=ERROR for failures)
2. Use get_trace to see the full execution flow
3. Use get_span to inspect specific LLM calls or tool executions
4. Use get_analytics to identify patterns across many traces
""",
)


def _get_client() -> TracesClient:
    """Get configured traces client from environment."""
    api_key = os.environ.get("TRACES_API_KEY")
    if not api_key:
        raise ValueError(
            "TRACES_API_KEY environment variable is required. "
            "Get your API key from the dashboard."
        )

    base_url = os.environ.get("TRACES_BASE_URL", "https://api.021labs.ai")
    return TracesClient(base_url=base_url, api_key=api_key)


def _format_trace_summary(trace: dict) -> str:
    """Format a trace summary for display."""
    status_icon = {"SUCCESS": "[ok]", "ERROR": "[err]", "RUNNING": "[...]"}.get(
        trace.get("status", ""), ""
    )
    duration = trace.get("duration_ms")
    duration_str = f"{duration}ms" if duration else "running"

    line = (
        f"{status_icon} {trace['id'][:12]}  {trace['name'][:40]:<40}  "
        f"{duration_str:>10}  spans:{trace.get('span_count', 0)}"
    )

    # Add tags if present
    tags = trace.get("tags", [])
    if tags:
        line += f"  [{', '.join(tags[:3])}]"
        if len(tags) > 3:
            line += f" +{len(tags) - 3}"

    return line


def _format_span_tree(spans: list[dict], parent_id: str | None = None, indent: int = 0) -> str:
    """Format spans as an indented tree."""
    lines = []
    for span in spans:
        if span.get("parent_id") == parent_id:
            prefix = "  " * indent
            status_icon = {"SUCCESS": "[ok]", "ERROR": "[err]", "RUNNING": "[...]"}.get(
                span.get("status", ""), ""
            )
            duration = span.get("duration_ms")
            duration_str = f"{duration}ms" if duration else ""

            type_str = span.get("type", "")[:4]
            model_str = f" ({span['model']})" if span.get("model") else ""

            lines.append(
                f"{prefix}{status_icon} [{type_str}] {span['name']}{model_str} {duration_str}"
            )
            lines.append(_format_span_tree(spans, span["id"], indent + 1))

    return "\n".join(filter(None, lines))


@mcp.tool()
async def list_traces(
    status: Annotated[
        str | None,
        "Filter by status: SUCCESS, ERROR, or RUNNING",
    ] = None,
    search: Annotated[
        str | None,
        "Full-text search query to find traces by name or content",
    ] = None,
    tags: Annotated[
        str | None,
        "Filter by tags (comma-separated). Traces must have ALL specified tags. Example: 'production,critical'",
    ] = None,
    start_time: Annotated[
        str | None,
        "Filter traces after this ISO timestamp (e.g., 2024-01-15T00:00:00Z)",
    ] = None,
    end_time: Annotated[
        str | None,
        "Filter traces before this ISO timestamp",
    ] = None,
    session_id: Annotated[
        str | None,
        "Filter by session ID to see all traces in a session",
    ] = None,
    user_id: Annotated[
        str | None,
        "Filter by user ID",
    ] = None,
    project_id: Annotated[
        str | None,
        "Filter by project ID",
    ] = None,
    limit: Annotated[
        int,
        "Maximum number of traces to return (1-100)",
    ] = 20,
) -> str:
    """List traces with optional filters.

    Returns a summary of matching traces including ID, name, status, duration, and span count.
    Use this to find traces to investigate further with get_trace.
    """
    async with _get_client() as client:
        # Parse optional datetime filters
        start_dt = datetime.fromisoformat(start_time.replace("Z", "+00:00")) if start_time else None
        end_dt = datetime.fromisoformat(end_time.replace("Z", "+00:00")) if end_time else None

        # Parse tags
        tags_list = [t.strip() for t in tags.split(",")] if tags else None

        # Validate status
        status_filter: SpanStatus | None = None
        if status:
            status_upper = status.upper()
            if status_upper in ("SUCCESS", "ERROR", "RUNNING"):
                status_filter = status_upper  # type: ignore
            else:
                return f"Invalid status '{status}'. Use SUCCESS, ERROR, or RUNNING."

        response = await client.list_traces(
            status=status_filter,
            search=search,
            tags=tags_list,
            start_time=start_dt,
            end_time=end_dt,
            session_id=session_id,
            user_id=user_id,
            project_id=project_id,
            limit=min(limit, 100),
        )

        if not response.traces:
            return "No traces found matching the filters."

        lines = [
            f"Found {response.pagination.total} traces (showing {len(response.traces)}):",
            "",
        ]
        for trace in response.traces:
            lines.append(_format_trace_summary(trace.model_dump()))

        if response.pagination.has_more:
            lines.append(f"\n... and {response.pagination.total - len(response.traces)} more")

        return "\n".join(lines)


@mcp.tool()
async def get_trace(
    trace_id: Annotated[str, "The trace ID to fetch"],
    include_io: Annotated[
        bool,
        "Include full input/output data for spans (slower but more detailed)",
    ] = False,
) -> str:
    """Get detailed information about a specific trace.

    Shows the trace metadata and a tree view of all spans (LLM calls, tool executions, etc.).
    Set include_io=True to see the full input/output for each span.
    """
    async with _get_client() as client:
        if include_io:
            trace = await client.get_trace(trace_id)
        else:
            trace = await client.get_trace_summary(trace_id)

        # Format header
        status_icon = {"SUCCESS": "[ok]", "ERROR": "[err]", "RUNNING": "[...]"}.get(
            trace.status, ""
        )
        duration_str = f"{trace.duration_ms}ms" if trace.duration_ms else "running"

        lines = [
            f"Trace: {trace.name}",
            f"ID: {trace.id}",
            f"Status: {status_icon} {trace.status}",
            f"Duration: {duration_str}",
            f"Spans: {trace.span_count} total, {trace.llm_count} LLM, {trace.tool_call_count} tools",
            f"Errors: {trace.error_count}",
        ]

        if trace.session_id:
            lines.append(f"Session: {trace.session_id}")
        if trace.total_tokens:
            lines.append(f"Tokens: {trace.total_tokens:,}")
        if trace.tags:
            lines.append(f"Tags: {', '.join(trace.tags)}")

        lines.append("")
        lines.append("Span tree:")
        lines.append("-" * 60)

        # Build span tree
        spans_data = [s.model_dump() for s in trace.spans]
        lines.append(_format_span_tree(spans_data))

        # If include_io, also show span details
        if include_io and trace.spans:
            lines.append("")
            lines.append("Span details:")
            lines.append("-" * 60)
            for span in trace.spans:
                if span.input or span.output:
                    lines.append(f"\n[{span.type}] {span.name} ({span.id[:8]})")
                    if span.input:
                        input_str = json.dumps(span.input, indent=2, default=str)
                        if len(input_str) > 500:
                            input_str = input_str[:500] + "..."
                        lines.append(f"  Input: {input_str}")
                    if span.output:
                        output_str = json.dumps(span.output, indent=2, default=str)
                        if len(output_str) > 500:
                            output_str = output_str[:500] + "..."
                        lines.append(f"  Output: {output_str}")

        return "\n".join(lines)


@mcp.tool()
async def get_span(
    span_id: Annotated[str, "The span ID to fetch"],
) -> str:
    """Get full details for a specific span including input/output.

    Use this to inspect the exact input sent to an LLM or tool and its response.
    Useful for debugging why a specific step failed or produced unexpected results.
    """
    async with _get_client() as client:
        span = await client.get_span(span_id)

        status_icon = {"SUCCESS": "[ok]", "ERROR": "[err]", "RUNNING": "[...]"}.get(
            span.status, ""
        )
        duration_str = f"{span.duration_ms}ms" if span.duration_ms else "running"

        lines = [
            f"Span: {span.name}",
            f"ID: {span.id}",
            f"Type: {span.type}",
            f"Status: {status_icon} {span.status}",
            f"Duration: {duration_str}",
        ]

        if span.model:
            lines.append(f"Model: {span.model}")
        if span.token_usage:
            lines.append(
                f"Tokens: {span.token_usage.prompt_tokens} prompt + "
                f"{span.token_usage.completion_tokens} completion = "
                f"{span.token_usage.total_tokens} total"
            )
        if span.parent_id:
            lines.append(f"Parent: {span.parent_id}")
        if span.tags:
            lines.append(f"Tags: {', '.join(span.tags)}")

        lines.append("")

        if span.input:
            lines.append("Input:")
            lines.append("-" * 40)
            lines.append(json.dumps(span.input, indent=2, default=str))
            lines.append("")

        if span.output:
            lines.append("Output:")
            lines.append("-" * 40)
            lines.append(json.dumps(span.output, indent=2, default=str))
            lines.append("")

        if span.error:
            lines.append("Error:")
            lines.append("-" * 40)
            lines.append(json.dumps(span.error, indent=2, default=str))

        if span.metadata:
            lines.append("")
            lines.append("Metadata:")
            lines.append("-" * 40)
            lines.append(json.dumps(span.metadata, indent=2, default=str))

        return "\n".join(lines)


@mcp.tool()
async def get_analytics(
    start_time: Annotated[
        str | None,
        "Start of time range (ISO timestamp)",
    ] = None,
    end_time: Annotated[
        str | None,
        "End of time range (ISO timestamp)",
    ] = None,
    granularity: Annotated[
        str,
        "Time series granularity: hour, day, or week",
    ] = "day",
    project_id: Annotated[
        str | None,
        "Filter by project ID",
    ] = None,
) -> str:
    """Get analytics summary for traces.

    Shows total traces, spans, error rates, latency percentiles, token usage,
    and per-model breakdowns. Use this to understand overall system health
    and identify performance patterns.
    """
    async with _get_client() as client:
        start_dt = datetime.fromisoformat(start_time.replace("Z", "+00:00")) if start_time else None
        end_dt = datetime.fromisoformat(end_time.replace("Z", "+00:00")) if end_time else None

        analytics = await client.get_analytics(
            start_time=start_dt,
            end_time=end_dt,
            granularity=granularity,
            project_id=project_id,
        )

        lines = [
            "Analytics Summary",
            "=" * 40,
            f"Total traces: {analytics.total_traces:,}",
            f"Total spans: {analytics.total_spans:,}",
            f"Error rate: {analytics.error_rate:.1%}",
            "",
            "Latency:",
            f"  Average: {analytics.avg_duration_ms:.0f}ms",
            f"  P50: {analytics.p50_duration_ms:.0f}ms",
            f"  P95: {analytics.p95_duration_ms:.0f}ms",
            f"  P99: {analytics.p99_duration_ms:.0f}ms",
            "",
            "Token Usage:",
            f"  Prompt: {analytics.token_usage.total_prompt:,}",
            f"  Completion: {analytics.token_usage.total_completion:,}",
            f"  Total: {analytics.token_usage.total:,}",
        ]

        if analytics.token_usage.estimated_cost:
            lines.append(f"  Est. Cost: ${analytics.token_usage.estimated_cost:.2f}")

        if analytics.by_model:
            lines.append("")
            lines.append("By Model:")
            lines.append("-" * 40)
            for model in analytics.by_model:
                avg_ms = f"{model.avg_duration_ms:.0f}ms" if model.avg_duration_ms else "N/A"
                tokens = f"{model.total_tokens:,}" if model.total_tokens else "N/A"
                lines.append(f"  {model.model}: {model.spans:,} calls, {avg_ms} avg, {tokens} tokens")

        if analytics.time_series:
            lines.append("")
            lines.append(f"Time Series ({granularity}):")
            lines.append("-" * 40)
            for point in analytics.time_series[-10:]:  # Last 10 points
                ts = point.timestamp.strftime("%Y-%m-%d %H:%M")
                lines.append(
                    f"  {ts}: {point.traces} traces, {point.spans} spans, {point.errors} errors"
                )

        return "\n".join(lines)


@mcp.tool()
async def search_errors(
    limit: Annotated[
        int,
        "Maximum number of error traces to return",
    ] = 10,
    start_time: Annotated[
        str | None,
        "Filter errors after this ISO timestamp",
    ] = None,
    tags: Annotated[
        str | None,
        "Filter by tags (comma-separated). Example: 'production,critical'",
    ] = None,
    project_id: Annotated[
        str | None,
        "Filter by project ID",
    ] = None,
) -> str:
    """Find recent traces that ended in errors.

    Shortcut for list_traces with status=ERROR. Use this to quickly identify
    and debug failing agent executions.
    """
    async with _get_client() as client:
        start_dt = datetime.fromisoformat(start_time.replace("Z", "+00:00")) if start_time else None
        tags_list = [t.strip() for t in tags.split(",")] if tags else None

        response = await client.list_traces(
            status="ERROR",
            start_time=start_dt,
            tags=tags_list,
            project_id=project_id,
            limit=min(limit, 50),
        )

        if not response.traces:
            return "No error traces found. All systems healthy!"

        lines = [
            f"Found {response.pagination.total} error traces (showing {len(response.traces)}):",
            "",
        ]
        for trace in response.traces:
            lines.append(_format_trace_summary(trace.model_dump()))

        return "\n".join(lines)


def main() -> None:
    """Run the MCP server."""
    mcp.run()


if __name__ == "__main__":
    main()
