"""Load and process JSON trace files — Python port of the Next.js viewer logic."""

from __future__ import annotations

import json
import os
import re
from collections import defaultdict
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import Any

from minitrail._pricing import compute_cost, cost_per_1m

# ---------------------------------------------------------------------------
# Data classes
# ---------------------------------------------------------------------------

MODEL_COLORS = [
    "#8b5cf6",  # violet
    "#f59e0b",  # amber
    "#10b981",  # emerald
    "#ef4444",  # red
    "#3b82f6",  # blue
    "#ec4899",  # pink
    "#14b8a6",  # teal
    "#f97316",  # orange
]


@dataclass
class TraceSummary:
    filename: str
    trace_id: str
    timestamp: str | None
    span_count: int
    duration_ms: float | None
    has_errors: bool
    total_input_tokens: int
    total_output_tokens: int
    total_cost: float


@dataclass
class SpanNode:
    span: dict
    children: list[SpanNode] = field(default_factory=list)
    depth: int = 0
    start_ms: float = 0
    end_ms: float = 0
    duration_ms: float = 0


@dataclass
class ModelStats:
    input_tokens: int = 0
    output_tokens: int = 0
    calls: int = 0


@dataclass
class NodeStats:
    by_model: dict[str, ModelStats] = field(default_factory=dict)
    total_cost: float = 0.0


@dataclass
class ReassembledMessage:
    role: str
    content: list[dict]  # [{type, text?} | {type, url?}]


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _parse_ts(ts: str | None) -> float:
    """Parse ISO 8601 timestamp to epoch milliseconds."""
    if not ts:
        return 0
    try:
        dt = datetime.fromisoformat(ts.replace("Z", "+00:00"))
        return dt.timestamp() * 1000
    except (ValueError, AttributeError):
        return 0


def format_duration(ms: float) -> str:
    if ms < 0 or ms != ms:  # NaN check
        return "\u2014"
    if ms < 1000:
        return f"{round(ms)}ms"
    if ms < 60000:
        return f"{ms / 1000:.1f}s"
    return f"{ms / 60000:.1f}m"


def format_cost(usd: float) -> str:
    if usd < 0.01:
        return f"${usd:.4f}"
    return f"${usd:.2f}"


def fmt_tokens(n: int) -> str:
    if n >= 1_000_000:
        return f"{n / 1_000_000:.1f}M"
    if n >= 1_000:
        return f"{n / 1_000:.1f}k"
    return str(n)


# ---------------------------------------------------------------------------
# Attribute helpers — support both OpenInference (llm.*) and GenAI (gen_ai.*)
# ---------------------------------------------------------------------------

def _get_model(attrs: dict) -> str | None:
    """Extract model name from OpenInference or GenAI attributes."""
    return attrs.get("llm.model_name") or attrs.get("gen_ai.request.model") or None


def _get_tokens(attrs: dict) -> tuple[int, int]:
    """Extract (input, output) token counts."""
    inp = int(
        attrs.get("llm.token_count.prompt")
        or attrs.get("gen_ai.usage.input_tokens")
        or attrs.get("gen_ai.usage.prompt_tokens")
        or 0
    )
    out = int(
        attrs.get("llm.token_count.completion")
        or attrs.get("gen_ai.usage.output_tokens")
        or attrs.get("gen_ai.usage.completion_tokens")
        or 0
    )
    return inp, out


def _parse_genai_content(raw: str) -> str:
    """Parse GenAI event content — may be a JSON array of {text: ...} blocks."""
    if not raw:
        return ""
    try:
        parsed = json.loads(raw)
        if isinstance(parsed, list):
            parts = []
            for item in parsed:
                if isinstance(item, dict) and "text" in item:
                    parts.append(item["text"])
                elif isinstance(item, str):
                    parts.append(item)
            return "\n".join(parts) if parts else raw
    except (json.JSONDecodeError, TypeError):
        pass
    return raw


def reassemble_genai_messages(events: list[dict]) -> tuple[list[ReassembledMessage], list[ReassembledMessage]]:
    """Extract input/output messages from GenAI semantic convention events."""
    input_msgs: list[ReassembledMessage] = []
    output_msgs: list[ReassembledMessage] = []

    for event in events:
        event_name = event.get("name", "")
        event_attrs = event.get("attributes", {})
        content_str = event_attrs.get("content", "")

        if event_name == "gen_ai.user.message":
            text = _parse_genai_content(content_str)
            if text:
                input_msgs.append(ReassembledMessage(role="user", content=[{"type": "text", "text": text}]))
        elif event_name == "gen_ai.system.message":
            text = _parse_genai_content(content_str)
            if text:
                input_msgs.append(ReassembledMessage(role="system", content=[{"type": "text", "text": text}]))
        elif event_name == "gen_ai.choice":
            text = event_attrs.get("message", "")
            text = _parse_genai_content(text) or str(text)
            if text:
                output_msgs.append(ReassembledMessage(role="assistant", content=[{"type": "text", "text": text}]))

    return input_msgs, output_msgs


# ---------------------------------------------------------------------------
# Trace loading
# ---------------------------------------------------------------------------

def _get_json_dir(logs_dir: str) -> str:
    json_dir = os.path.join(logs_dir, "json")
    return json_dir if os.path.isdir(json_dir) else logs_dir


def load_traces(logs_dir: str) -> list[TraceSummary]:
    """Load all trace summaries from the logs directory."""
    json_dir = _get_json_dir(logs_dir)
    if not os.path.isdir(json_dir):
        return []

    files = sorted(
        [f for f in os.listdir(json_dir) if f.endswith(".json")],
        reverse=True,
    )

    traces: list[TraceSummary] = []
    for filename in files:
        filepath = os.path.join(json_dir, filename)
        try:
            with open(filepath) as f:
                data = json.load(f)
        except (json.JSONDecodeError, OSError):
            continue

        spans: list[dict] = data.get("spans", [])
        min_start = float("inf")
        max_end = 0.0
        error_count = 0
        total_input = 0
        total_output = 0
        total_cost = 0.0

        for span in spans:
            start = _parse_ts(span.get("start_time"))
            end = _parse_ts(span.get("end_time"))
            if start < min_start:
                min_start = start
            if end > max_end:
                max_end = end
            if span.get("status", {}).get("status_code") == "ERROR":
                error_count += 1
            attrs = span.get("attributes", {})
            inp, out = _get_tokens(attrs)
            total_input += inp
            total_output += out
            model = _get_model(attrs)
            if model:
                total_cost += compute_cost(str(model), inp, out)

        duration_ms = max_end - min_start
        timestamp = (
            datetime.fromtimestamp(min_start / 1000, tz=timezone.utc).isoformat()
            if min_start != float("inf")
            else None
        )

        traces.append(TraceSummary(
            filename=filename,
            trace_id=data.get("trace_id", ""),
            timestamp=timestamp,
            span_count=len(spans),
            duration_ms=duration_ms if min_start != float("inf") else None,
            has_errors=error_count > 0,
            total_input_tokens=total_input,
            total_output_tokens=total_output,
            total_cost=total_cost,
        ))

    return traces


def load_trace(logs_dir: str, trace_id: str) -> dict | None:
    """Load a single trace by ID."""
    json_dir = _get_json_dir(logs_dir)
    if not os.path.isdir(json_dir):
        return None

    files = [f for f in os.listdir(json_dir) if f.endswith(".json")]
    match = next((f for f in files if trace_id in f), None)
    if not match:
        return None

    with open(os.path.join(json_dir, match)) as f:
        return json.load(f)


# ---------------------------------------------------------------------------
# Span tree
# ---------------------------------------------------------------------------

def build_span_tree(spans: list[dict]) -> list[SpanNode]:
    """Build parent-child tree from flat span list."""
    node_map: dict[str, SpanNode] = {}

    for span in spans:
        span_id = span.get("context", {}).get("span_id", "")
        start = _parse_ts(span.get("start_time"))
        end = _parse_ts(span.get("end_time"))
        node_map[span_id] = SpanNode(
            span=span,
            start_ms=start,
            end_ms=end,
            duration_ms=end - start,
        )

    roots: list[SpanNode] = []
    for span in spans:
        span_id = span.get("context", {}).get("span_id", "")
        parent_id = span.get("parent_id")
        node = node_map[span_id]
        if parent_id and parent_id in node_map:
            node_map[parent_id].children.append(node)
        else:
            roots.append(node)

    def sort_children(node: SpanNode) -> None:
        node.children.sort(key=lambda n: n.start_ms)
        for child in node.children:
            sort_children(child)

    roots.sort(key=lambda n: n.start_ms)
    for root in roots:
        sort_children(root)

    return roots


def flatten_tree(roots: list[SpanNode], depth: int = 0) -> list[SpanNode]:
    """Flatten tree into a list with depth assigned."""
    result: list[SpanNode] = []
    for node in roots:
        node.depth = depth
        result.append(node)
        result.extend(flatten_tree(node.children, depth + 1))
    return result


# ---------------------------------------------------------------------------
# Stats aggregation
# ---------------------------------------------------------------------------

def aggregate_stats(node: SpanNode) -> NodeStats:
    """Recursively aggregate token/cost stats for a subtree."""
    by_model: dict[str, ModelStats] = {}

    def walk(n: SpanNode) -> None:
        attrs = n.span.get("attributes", {})
        model = _get_model(attrs)
        inp, out = _get_tokens(attrs)

        if model and (inp or out):
            if model not in by_model:
                by_model[model] = ModelStats()
            by_model[model].input_tokens += inp
            by_model[model].output_tokens += out
            by_model[model].calls += 1

        for child in n.children:
            walk(child)

    walk(node)

    total_cost = 0.0
    for model, stats in by_model.items():
        total_cost += compute_cost(model, stats.input_tokens, stats.output_tokens)

    return NodeStats(by_model=by_model, total_cost=total_cost)


def aggregate_stats_for_roots(roots: list[SpanNode]) -> NodeStats:
    """Aggregate stats across all root nodes."""
    merged: dict[str, ModelStats] = {}
    total_cost = 0.0
    for root in roots:
        s = aggregate_stats(root)
        for model, ms in s.by_model.items():
            if model not in merged:
                merged[model] = ModelStats()
            merged[model].input_tokens += ms.input_tokens
            merged[model].output_tokens += ms.output_tokens
            merged[model].calls += ms.calls
        total_cost += s.total_cost
    return NodeStats(by_model=merged, total_cost=total_cost)


# ---------------------------------------------------------------------------
# Reassemble OpenInference flat attributes into chat messages
# ---------------------------------------------------------------------------

def reassemble_messages(attrs: dict, prefix: str) -> list[ReassembledMessage]:
    """Reconstruct structured chat messages from flat OpenInference keys."""
    dot = prefix + "."
    entries = [(k[len(dot):], v) for k, v in attrs.items() if k.startswith(dot)]
    if not entries:
        return []

    by_idx: dict[int, dict[str, Any]] = defaultdict(dict)
    for suffix, val in entries:
        parts = suffix.split(".", 1)
        if len(parts) < 2:
            continue
        try:
            idx = int(parts[0])
        except ValueError:
            continue
        by_idx[idx][parts[1]] = val

    messages: list[ReassembledMessage] = []
    for idx in sorted(by_idx):
        fields = by_idx[idx]
        role = str(fields.get("message.role", "unknown"))
        blocks: list[dict] = []

        # Simple string content
        simple = fields.get("message.content")
        if isinstance(simple, str) and simple:
            blocks.append({"type": "text", "text": simple})

        # Multimodal content blocks
        content_fields: dict[int, dict[str, Any]] = defaultdict(dict)
        for k, v in fields.items():
            m = re.match(r"^message\.contents\.(\d+)\.message_content\.(.+)$", k)
            if m:
                content_fields[int(m.group(1))][m.group(2)] = v

        for cidx in sorted(content_fields):
            cf = content_fields[cidx]
            ctype = cf.get("type")
            if ctype == "text":
                t = cf.get("text")
                if isinstance(t, str):
                    blocks.append({"type": "text", "text": t})
            elif ctype == "image":
                url = cf.get("image.image.url")
                if isinstance(url, str):
                    blocks.append({"type": "image", "url": url})

        messages.append(ReassembledMessage(role=role, content=blocks))

    return messages


# ---------------------------------------------------------------------------
# Model color mapping
# ---------------------------------------------------------------------------

def assign_model_colors(spans: list[dict]) -> dict[str, str]:
    """Assign a stable color to each model found in the spans."""
    models: set[str] = set()
    for span in spans:
        m = _get_model(span.get("attributes", {}))
        if m:
            models.add(m)
    sorted_models = sorted(models)
    return {
        m: MODEL_COLORS[i % len(MODEL_COLORS)]
        for i, m in enumerate(sorted_models)
    }


def is_llm_span(span: dict) -> bool:
    """Check if a span is an LLM span."""
    name = span.get("name", "").lower()
    attrs = span.get("attributes", {})
    return (
        "llm" in name
        or "chat" in name
        or any(k.startswith("llm.") or k.startswith("gen_ai.") for k in attrs)
    )


def compute_span_cost(span: dict) -> float:
    """Compute cost for a single span."""
    attrs = span.get("attributes", {})
    model = _get_model(attrs)
    if not model:
        return 0.0
    inp, out = _get_tokens(attrs)
    return compute_cost(model, inp, out)


# ---------------------------------------------------------------------------
# smolagents-specific helpers
# ---------------------------------------------------------------------------

def parse_smolagents_input_value(raw: str) -> list[ReassembledMessage]:
    """Parse smolagents input.value to extract chat messages.

    smolagents stores input.value as a JSON dict with a "messages" key containing
    Python repr strings of ChatMessage objects, e.g.:
        ChatMessage(role=<MessageRole.SYSTEM: 'system'>, content=[{'type': 'text', 'text': '...'}], ...)
    """
    if not raw or not isinstance(raw, str):
        return []
    try:
        parsed = json.loads(raw)
    except (json.JSONDecodeError, TypeError):
        return []
    if not isinstance(parsed, dict) or "messages" not in parsed:
        return []

    messages: list[ReassembledMessage] = []
    for msg_str in parsed["messages"]:
        if not isinstance(msg_str, str):
            continue
        # Extract role from ChatMessage repr
        role_match = re.search(r"role=<MessageRole\.\w+:\s*'(\w+)'>", msg_str)
        if not role_match:
            continue
        role = role_match.group(1)

        # Extract content list from repr and eval it
        content_match = re.search(r"content=(\[.+?\])(?:,\s*tool_calls)", msg_str, re.DOTALL)
        text = ""
        if content_match:
            try:
                import ast
                content_list = ast.literal_eval(content_match.group(1))
                text_parts = [
                    item["text"] for item in content_list
                    if isinstance(item, dict) and "text" in item
                ]
                text = "\n".join(text_parts)
            except Exception:
                text = msg_str
        else:
            text = msg_str

        if text:
            messages.append(ReassembledMessage(role=role, content=[{"type": "text", "text": text}]))

    return messages


# ---------------------------------------------------------------------------
# CrewAI-specific helpers
# ---------------------------------------------------------------------------

def extract_crewai_agent_models(spans: list[dict]) -> dict[str, str]:
    """Extract agent_role → model mapping from CrewAI 'Crew Created' span."""
    for span in spans:
        if span.get("name") == "Crew Created":
            raw = span.get("attributes", {}).get("crew_agents", "")
            try:
                agents = json.loads(raw)
                return {a["role"]: a["llm"] for a in agents if "role" in a and "llm" in a}
            except (json.JSONDecodeError, TypeError):
                pass
    return {}


def parse_crewai_output(attrs: dict) -> tuple[list[ReassembledMessage], list[ReassembledMessage], str | None]:
    """Parse CrewAI AGENT span output.value to extract messages and raw output.

    Returns (input_messages, output_messages, raw_output).
    """
    raw_val = attrs.get("output.value", "")
    if not raw_val or not isinstance(raw_val, str):
        return [], [], None

    try:
        parsed = json.loads(raw_val)
    except (json.JSONDecodeError, TypeError):
        return [], [], None

    if not isinstance(parsed, dict) or "messages" not in parsed:
        return [], [], None

    input_msgs: list[ReassembledMessage] = []
    output_msgs: list[ReassembledMessage] = []

    for msg in parsed["messages"]:
        role = msg.get("role", "unknown")
        content = msg.get("content", "")
        rm = ReassembledMessage(role=role, content=[{"type": "text", "text": content}])
        if role in ("assistant", "ai"):
            output_msgs.append(rm)
        else:
            input_msgs.append(rm)

    raw_output = parsed.get("raw")
    return input_msgs, output_msgs, raw_output
