"""FastAPI application — HTML pages + JSON API for trace viewing."""

from __future__ import annotations

import json
import os
from dataclasses import asdict
from pathlib import Path

from fastapi import FastAPI, Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

from minitrail._server.trace_loader import (
    _get_model,
    _get_tokens,
    aggregate_stats,
    aggregate_stats_for_roots,
    assign_model_colors,
    build_span_tree,
    compute_span_cost,
    extract_crewai_agent_models,
    flatten_tree,
    fmt_tokens,
    format_cost,
    format_duration,
    is_llm_span,
    load_trace,
    load_traces,
    parse_crewai_output,
    parse_smolagents_input_value,
    reassemble_genai_messages,
    reassemble_messages,
)

_HERE = Path(__file__).resolve().parent


def create_app(logs_dir: str = "logs") -> FastAPI:
    """Create the FastAPI application configured to read traces from *logs_dir*."""
    app = FastAPI(title="minitrail viewer")

    templates = Jinja2Templates(directory=str(_HERE / "templates"))

    # Jinja2 globals for templates
    templates.env.globals["format_cost"] = format_cost
    templates.env.globals["format_duration"] = format_duration
    templates.env.globals["fmt_tokens"] = fmt_tokens
    templates.env.globals["is_llm_span"] = is_llm_span
    templates.env.globals["json_dumps"] = lambda v: json.dumps(v, default=str)

    # Static files
    app.mount("/static", StaticFiles(directory=str(_HERE / "static")), name="static")

    # ----- HTML routes -----

    @app.get("/", response_class=HTMLResponse)
    async def home(request: Request):
        traces = load_traces(logs_dir)
        return templates.TemplateResponse("home.html", {
            "request": request,
            "traces": traces,
        })

    @app.get("/traces/{trace_id}", response_class=HTMLResponse)
    async def trace_detail(request: Request, trace_id: str):
        data = load_trace(logs_dir, trace_id)
        if not data:
            return HTMLResponse("<h1>Trace not found</h1>", status_code=404)

        spans = data.get("spans", [])
        tree = build_span_tree(spans)
        flat = flatten_tree(tree)
        trace_stats = aggregate_stats_for_roots(tree)
        model_colors = assign_model_colors(spans)

        # CrewAI: extract agent_role → model mapping from "Crew Created" span
        crewai_models = extract_crewai_agent_models(spans)
        # Add CrewAI models to the color map
        for m in crewai_models.values():
            if m and m not in model_colors:
                idx = len(model_colors)
                from minitrail._server.trace_loader import MODEL_COLORS
                model_colors[m] = MODEL_COLORS[idx % len(MODEL_COLORS)]

        # Compute global trace timing
        trace_start = min((n.start_ms for n in flat), default=0)
        trace_end = max((n.end_ms for n in flat), default=0)
        trace_duration = trace_end - trace_start

        # Pre-compute stats for every node (keyed by span_id)
        stats_map: dict[str, dict] = {}

        def fill_stats(nodes: list) -> None:
            for n in nodes:
                sid = n.span.get("context", {}).get("span_id", "")
                s = aggregate_stats(n)
                stats_map[sid] = {
                    "by_model": {
                        m: {
                            "input_tokens": ms.input_tokens,
                            "output_tokens": ms.output_tokens,
                            "calls": ms.calls,
                        }
                        for m, ms in s.by_model.items()
                    },
                    "total_cost": s.total_cost,
                }
                fill_stats(n.children)

        fill_stats(tree)

        # Build per-span data for template
        span_rows = []
        for node in flat:
            span = node.span
            span_id = span.get("context", {}).get("span_id", "")
            attrs = span.get("attributes", {})
            model_name = _get_model(attrs)
            kind = attrs.get("openinference.span.kind", "")
            has_children = len(node.children) > 0
            llm = is_llm_span(span)
            node_stats = stats_map.get(span_id, {})
            has_cost = bool(node_stats.get("by_model"))

            # CrewAI: resolve model from agent role for AGENT spans
            crewai_model = None
            if kind == "AGENT" and crewai_models:
                # Span name is like "Researcher._execute_core"
                agent_role = span.get("name", "").split(".")[0]
                crewai_model = crewai_models.get(agent_role)
                if not model_name and crewai_model:
                    model_name = crewai_model

            offset_pct = (
                ((node.start_ms - trace_start) / trace_duration) * 100
                if trace_duration > 0
                else 0
            )
            width_pct = (
                max((node.duration_ms / trace_duration) * 100, 0.5)
                if trace_duration > 0
                else 100
            )

            # Bar color
            is_error = span.get("status", {}).get("status_code") == "ERROR"
            bar_color = ""
            if is_error:
                bar_color = ""  # Use CSS class
            elif model_name and model_name in model_colors:
                bar_color = model_colors[model_name]

            # Tokens summary for the row
            if has_children and has_cost:
                sum_in = sum(ms["input_tokens"] for ms in node_stats["by_model"].values())
                sum_out = sum(ms["output_tokens"] for ms in node_stats["by_model"].values())
                tokens_label = f"\u03A3 {fmt_tokens(sum_in)} in / {fmt_tokens(sum_out)} out"
                cost_label = format_cost(node_stats["total_cost"])
                show_tokens = True
            elif llm:
                inp, out = _get_tokens(attrs)
                tokens_label = f"{fmt_tokens(inp)} in / {fmt_tokens(out)} out"
                cost_label = format_cost(compute_span_cost(span))
                show_tokens = True
            else:
                tokens_label = ""
                cost_label = ""
                show_tokens = False

            # Pre-compute detail data
            input_messages = reassemble_messages(attrs, "llm.input_messages")
            output_messages = reassemble_messages(attrs, "llm.output_messages")

            # Fall back to GenAI event-based messages
            used_genai_events = False
            if not input_messages and not output_messages:
                genai_in, genai_out = reassemble_genai_messages(span.get("events", []))
                input_messages = genai_in
                output_messages = genai_out
                if genai_in or genai_out:
                    used_genai_events = True

            # Fall back to smolagents input.value parsing
            if not input_messages and not output_messages:
                smol_msgs = parse_smolagents_input_value(attrs.get("input.value", ""))
                if smol_msgs:
                    input_messages = smol_msgs

            # Fall back to CrewAI output.value messages
            crewai_raw_output = None
            if not input_messages and not output_messages and kind == "AGENT":
                crewai_in, crewai_out, crewai_raw_output = parse_crewai_output(attrs)
                input_messages = crewai_in
                output_messages = crewai_out

            # LLM metadata keys (both llm.* and gen_ai.* scalars)
            meta_keys = [
                k for k in attrs
                if (k.startswith("llm.") or k.startswith("gen_ai."))
                and not k.startswith("llm.input_messages.")
                and not k.startswith("llm.output_messages.")
            ]
            # For CrewAI AGENT spans, add model as a meta pill
            if crewai_model and not meta_keys:
                meta_keys = ["_crewai_model"]
                attrs["_crewai_model"] = crewai_model

            # Other attributes (not llm.*, gen_ai.*, input.*, output.*, crewai internal)
            hidden_prefixes = [
                "llm.", "gen_ai.", "input.value", "input.mime_type",
                "output.value", "output.mime_type", "_crewai_model",
            ]
            other_keys = [k for k in attrs if not any(k.startswith(p) for p in hidden_prefixes)]

            span_rows.append({
                "span_id": span_id,
                "name": span.get("name", ""),
                "depth": node.depth,
                "has_children": has_children,
                "is_llm": llm,
                "is_error": is_error,
                "bar_color": bar_color,
                "offset_pct": offset_pct,
                "width_pct": width_pct,
                "duration": format_duration(node.duration_ms),
                "show_tokens": show_tokens,
                "tokens_label": tokens_label,
                "cost_label": cost_label,
                "is_aggregate": has_children and has_cost,
                # Detail data
                "start_time": span.get("start_time", ""),
                "end_time": span.get("end_time", ""),
                "status_code": span.get("status", {}).get("status_code", "UNSET"),
                "input_messages": [
                    {"role": m.role, "content": m.content} for m in input_messages
                ],
                "output_messages": [
                    {"role": m.role, "content": m.content} for m in output_messages
                ],
                "meta_keys": [
                    {"label": k.split(".")[-1], "value": _stringify_attr(attrs[k])}
                    for k in meta_keys
                ],
                "other_keys": [
                    {"key": k, "value": _stringify_attr(attrs[k])}
                    for k in other_keys
                ],
                "has_input_value": "input.value" in attrs and not input_messages,
                "input_value": str(attrs.get("input.value", "")),
                "has_output_value": "output.value" in attrs and not output_messages,
                "output_value": crewai_raw_output or str(attrs.get("output.value", "")),
                "events": [
                    e for e in span.get("events", [])
                    if not used_genai_events
                    or e.get("name", "") not in (
                        "gen_ai.user.message", "gen_ai.system.message", "gen_ai.choice",
                    )
                ],
                "raw_attrs": attrs,
            })

        # Trace-level cost summary for template
        cost_rows = []
        for model, ms in trace_stats.by_model.items():
            cost = compute_span_cost({"attributes": {
                "llm.model_name": model,
                "llm.token_count.prompt": ms.input_tokens,
                "llm.token_count.completion": ms.output_tokens,
            }})
            cost_rows.append({
                "model": model,
                "color": model_colors.get(model, ""),
                "calls": ms.calls,
                "input_tokens": f"{ms.input_tokens:,}",
                "output_tokens": f"{ms.output_tokens:,}",
                "cost": format_cost(cost),
            })

        return templates.TemplateResponse("trace.html", {
            "request": request,
            "trace_id": data.get("trace_id", trace_id),
            "span_count": len(spans),
            "trace_duration": format_duration(trace_duration),
            "total_cost": format_cost(trace_stats.total_cost),
            "cost_rows": cost_rows,
            "span_rows": span_rows,
            "span_rows_json": json.dumps(span_rows, default=str),
        })

    # ----- JSON API routes -----

    @app.get("/api/traces")
    async def api_traces():
        traces = load_traces(logs_dir)
        return JSONResponse([asdict(t) for t in traces])

    @app.get("/api/traces/{trace_id}")
    async def api_trace(trace_id: str):
        data = load_trace(logs_dir, trace_id)
        if not data:
            return JSONResponse({"error": "Trace not found"}, status_code=404)
        return JSONResponse(data)

    return app


def _stringify_attr(val: object) -> str:
    if isinstance(val, str):
        return val
    return json.dumps(val, indent=2, default=str)
