"""Parse Codex JSONL transcript files for LLM span data.

Codex stores session transcripts at:
    ~/.codex/sessions/YYYY/MM/DD/rollout-<timestamp>-<session_id>.jsonl

Each turn produces a sequence of JSONL entries. The relevant ones are:
    - turn_context   → model name
    - token_count    → per-request token usage (last_token_usage)
    - response_item  → assistant message content
    - event_msg      → task_started / task_complete boundaries

We extract one LLMMessage per turn (Codex makes a single LLM call per turn,
unlike Claude Code which may make multiple within one turn).
"""
from __future__ import annotations

import json
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Any

logger = logging.getLogger("staso.codex")


@dataclass
class LLMMessage:
    """LLM span data extracted from a Codex transcript turn."""

    model: str
    input_tokens: int
    output_tokens: int
    cached_input_tokens: int
    reasoning_output_tokens: int
    content: list[dict[str, Any]]
    timestamp: str

    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens


def find_transcript(session_id: str) -> Path | None:
    """Find the transcript JSONL file for a Codex session.

    Searches ~/.codex/sessions/ for files whose name contains the session_id.
    Returns None if not found.
    """
    sessions_dir = Path.home() / ".codex" / "sessions"
    if not sessions_dir.exists():
        return None
    # Session files are stored as: rollout-<timestamp>-<session_id>.jsonl
    matches = list(sessions_dir.rglob(f"*{session_id}.jsonl"))
    if not matches:
        return None
    # Return the most recently modified match
    return max(matches, key=lambda p: p.stat().st_mtime)


def count_lines(transcript_path: str) -> int:
    """Return the total number of lines in the transcript file."""
    path = Path(transcript_path)
    if not path.exists():
        return 0
    try:
        with open(path, encoding="utf-8") as f:
            return sum(1 for _ in f)
    except OSError:
        return 0


def parse_turn_llm_data(
    transcript_path: str, start_line: int
) -> tuple[list[LLMMessage], int]:
    """Parse LLM usage data from transcript lines after *start_line*.

    Returns (messages, last_line_read).
    Extracts token usage from ``token_count`` events and model from
    ``turn_context`` entries.
    """
    path = Path(transcript_path)
    if not path.exists():
        return [], start_line

    try:
        raw_lines = path.read_text(encoding="utf-8").splitlines()
    except OSError:
        return [], start_line

    messages: list[LLMMessage] = []
    last_line = start_line

    # Accumulate per-turn data as we scan
    current_model = ""
    current_timestamp = ""
    current_content: list[dict[str, Any]] = []
    last_usage: dict[str, Any] | None = None

    for i, line in enumerate(raw_lines):
        line_num = i + 1
        if line_num <= start_line:
            continue
        last_line = line_num
        stripped = line.strip()
        if not stripped:
            continue
        try:
            entry = json.loads(stripped)
        except json.JSONDecodeError:
            continue

        entry_type = entry.get("type", "")
        payload = entry.get("payload", {})
        timestamp = entry.get("timestamp", "")

        if entry_type == "turn_context":
            current_model = payload.get("model", "")
            current_timestamp = timestamp

        elif entry_type == "event_msg":
            msg_type = payload.get("type", "")

            if msg_type == "token_count" and payload.get("info"):
                info = payload["info"]
                usage = info.get("last_token_usage") or info.get("total_token_usage")
                if usage:
                    last_usage = usage

            elif msg_type == "task_complete":
                # Turn complete — emit an LLM message if we have token data
                if last_usage:
                    messages.append(
                        LLMMessage(
                            model=current_model,
                            input_tokens=last_usage.get("input_tokens", 0),
                            output_tokens=last_usage.get("output_tokens", 0),
                            cached_input_tokens=last_usage.get("cached_input_tokens", 0),
                            reasoning_output_tokens=last_usage.get("reasoning_output_tokens", 0),
                            content=current_content,
                            timestamp=current_timestamp or timestamp,
                        )
                    )
                # Reset for next turn
                current_model = ""
                current_timestamp = ""
                current_content = []
                last_usage = None

            elif msg_type == "task_started":
                # New turn starting — reset accumulators
                current_model = ""
                current_timestamp = ""
                current_content = []
                last_usage = None

        elif entry_type == "response_item":
            if payload.get("role") == "assistant":
                content = payload.get("content") or []
                current_content = content

    return messages, last_line
