"""Trajectory tracking — traces, spans, and execution turn parsing.

Every user interaction creates a trace that tracks the full lifecycle:
input -> decision -> dispatch -> execution (every Claude Code turn) -> result -> notification.
"""
from __future__ import annotations

import json
import logging
import secrets
from datetime import UTC, datetime

logger = logging.getLogger(__name__)

# Cost-per-token assumptions for Claude Sonnet 4 (used if not provided)
_INPUT_COST_PER_TOKEN = 3.0 / 1_000_000   # $3 per 1M input
_OUTPUT_COST_PER_TOKEN = 15.0 / 1_000_000  # $15 per 1M output
_CACHE_READ_COST_PER_TOKEN = 0.30 / 1_000_000  # $0.30 per 1M cache read


def _generate_id(prefix: str = "tr") -> str:
    return f"{prefix}-{secrets.token_hex(8)}"


def _now_iso() -> str:
    return datetime.now(UTC).isoformat()


class TrajectoryParser:
    """Parses stream-json lines into trace spans and execution turns."""

    def __init__(self) -> None:
        # Per-task tracking of the current execution span and turn counter
        self._task_spans: dict[str, str] = {}   # task_id -> span_id
        self._task_turns: dict[str, int] = {}    # task_id -> turn_number

    async def create_trace(
        self,
        store,
        source: str,
        source_id: str | None = None,
        user_id: str | None = None,
        project: str | None = None,
        input_text: str | None = None,
    ) -> str:
        """Create a new trace. Returns trace_id."""
        trace_id = _generate_id("tr")
        await store.create_trace(
            trace_id=trace_id,
            source=source,
            source_id=source_id,
            user_id=user_id,
            project=project,
            input_text=input_text,
            created_at=_now_iso(),
        )
        return trace_id

    async def add_span(
        self,
        store,
        trace_id: str,
        kind: str,
        task_id: str | None = None,
        session_id: str | None = None,
        log_path: str | None = None,
        parent_span_id: str | None = None,
        metadata: dict | None = None,
    ) -> str:
        """Add a span to a trace. Returns span_id."""
        span_id = _generate_id("sp")
        await store.create_span(
            span_id=span_id,
            trace_id=trace_id,
            kind=kind,
            started_at=_now_iso(),
            parent_span_id=parent_span_id,
            task_id=task_id,
            session_id=session_id,
            log_path=log_path,
            metadata=metadata,
        )
        # Track execution spans for ingest_log_line
        if kind == "execution" and task_id:
            self._task_spans[task_id] = span_id
            self._task_turns[task_id] = 0
        return span_id

    async def ingest_log_line(self, store, task_id: str, line: str) -> None:
        """Parse one stream-json line and store as execution_turn.

        Called from spoke_append_log as lines arrive.
        Must be lightweight — called on every streamed line.
        """
        line = line.strip()
        if not line:
            return

        try:
            data = json.loads(line)
        except (json.JSONDecodeError, ValueError):
            return  # skip non-JSON lines

        msg_type = data.get("type")
        if msg_type not in ("assistant", "user", "system", "result"):
            return

        # Find or create the execution span for this task
        span_id = self._task_spans.get(task_id)
        if not span_id:
            # Try to look up from DB
            span_info = await store.get_execution_span_for_task(task_id)
            if span_info:
                span_id = span_info["span_id"]
                self._task_spans[task_id] = span_id
                self._task_turns[task_id] = 0
            else:
                return  # no execution span registered yet

        turn_num = self._task_turns.get(task_id, 0) + 1
        self._task_turns[task_id] = turn_num

        # Extract fields based on message type
        role = msg_type
        content_type = None
        tool_name = None
        input_tokens = 0
        output_tokens = 0
        cache_read_tokens = 0
        cost_usd = 0.0
        content_preview = None
        timestamp = data.get("timestamp") or _now_iso()

        if msg_type == "assistant":
            message = data.get("message", {})
            usage = message.get("usage") or data.get("usage", {})
            input_tokens = usage.get("input_tokens", 0)
            output_tokens = usage.get("output_tokens", 0)
            cache_read_tokens = usage.get("cache_read_input_tokens", 0)
            cost_usd = (
                input_tokens * _INPUT_COST_PER_TOKEN
                + output_tokens * _OUTPUT_COST_PER_TOKEN
                + cache_read_tokens * _CACHE_READ_COST_PER_TOKEN
            )

            # Determine content type from first content block
            content_blocks = message.get("content", [])
            if content_blocks:
                first = content_blocks[0] if isinstance(content_blocks, list) else {}
                if isinstance(first, dict):
                    content_type = first.get("type", "text")
                    if content_type == "tool_use":
                        tool_name = first.get("name")
                    if content_type == "text":
                        text = first.get("text", "")
                        content_preview = text[:200] if text else None
                    elif content_type == "thinking":
                        text = first.get("thinking", "")
                        content_preview = text[:200] if text else None
                elif isinstance(first, str):
                    content_type = "text"
                    content_preview = first[:200]

        elif msg_type == "user":
            content_blocks = data.get("message", {}).get("content", [])
            if content_blocks and isinstance(content_blocks, list):
                first = content_blocks[0] if content_blocks else {}
                if isinstance(first, dict):
                    content_type = first.get("type", "text")
                    if content_type == "tool_result":
                        content_preview = str(first.get("content", ""))[:200]

        elif msg_type == "result":
            role = "result"
            content_type = "result"
            result_data = data.get("result", data)
            cost_info = result_data.get("cost", {})
            input_tokens = cost_info.get("input_tokens", 0)
            output_tokens = cost_info.get("output_tokens", 0)
            cost_usd = cost_info.get("cost_usd", 0.0)
            content_preview = str(result_data.get("summary", ""))[:200]

        await store.insert_execution_turn(
            span_id=span_id,
            turn_number=turn_num,
            role=role,
            content_type=content_type,
            tool_name=tool_name,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cache_read_tokens=cache_read_tokens,
            cost_usd=cost_usd,
            content_preview=content_preview,
            timestamp=timestamp,
        )

    async def finalize_trace(
        self,
        store,
        trace_id: str,
        status: str,
        total_input_tokens: int = 0,
        total_output_tokens: int = 0,
        total_cost_usd: float = 0.0,
    ) -> None:
        """Mark trace as completed with final stats."""
        await store.finalize_trace(
            trace_id=trace_id,
            status=status,
            completed_at=_now_iso(),
            total_input_tokens=total_input_tokens,
            total_output_tokens=total_output_tokens,
            total_cost_usd=total_cost_usd,
        )


# Module-level singleton used by instrumentation points
trajectory_parser = TrajectoryParser()
