"""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).

Additionally, we extract **non-shell tool calls** (apply_patch, update_plan,
view_image, spawn_agent, etc.) from the transcript.  Codex only fires
PreToolUse/PostToolUse hooks for shell/exec tools; all other tool calls are
only recorded in the transcript JSONL.

NOTE: Codex writes transcript lines asynchronously with respect to hooks.
By the time a hook process starts up, more lines may have been written than
expected.  For this reason we parse by **turn count** (number of task_complete
events), not by line offset.
"""
from __future__ import annotations

import json
import logging
from dataclasses import dataclass, field
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


@dataclass
class ToolCallRecord:
    """A non-shell tool call extracted from the transcript."""

    call_id: str
    name: str
    input_data: dict[str, Any] | str = ""
    output_data: dict[str, Any] | str = ""
    timestamp: str = ""


# Shell tool names that already fire hooks — skip these during transcript extraction
_SHELL_TOOLS = frozenset(("exec_command", "write_stdin", "shell", "shell_command"))


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, turns_already_processed: int
) -> tuple[list[LLMMessage], int]:
    """Parse LLM usage data for turns not yet processed.

    Args:
        transcript_path: Path to the Codex session JSONL transcript.
        turns_already_processed: Number of turns (task_complete events)
            already sent to the backend. We skip these and return data
            for any new turns only.

    Returns:
        (new_messages, total_turns_seen) — the caller should store
        ``total_turns_seen`` as the new ``turns_already_processed``
        for the next invocation.
    """
    path = Path(transcript_path)
    if not path.exists():
        return [], turns_already_processed

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

    all_messages: list[LLMMessage] = []
    total_turns = 0

    # 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 line in raw_lines:
        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":
                total_turns += 1
                # Only collect data for turns we haven't processed yet
                if total_turns > turns_already_processed and last_usage:
                    all_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

    # Handle incomplete turn at end of file (Stop hook fires before
    # Codex writes task_complete to the transcript).
    if last_usage and total_turns >= turns_already_processed:
        total_turns += 1
        if total_turns > turns_already_processed:
            all_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,
                )
            )

    return all_messages, total_turns


def parse_turn_tool_calls(
    transcript_path: str,
    turns_already_processed: int,
    hook_call_ids: set[str] | None = None,
) -> list[ToolCallRecord]:
    """Extract non-shell tool calls from the transcript for the current turn.

    Codex only fires PreToolUse/PostToolUse hooks for shell/exec tools.
    Other tools (apply_patch, update_plan, view_image, spawn_agent, etc.)
    are only recorded in the transcript.  This function scans for those
    tool calls so they can be emitted as spans during the Stop handler.

    Args:
        transcript_path: Path to the Codex session JSONL transcript.
        turns_already_processed: Number of turns already sent. We only
            return tool calls from turns after this.
        hook_call_ids: Set of call_ids already captured via hooks.
            These are skipped to avoid duplicates.

    Returns:
        List of ToolCallRecord for non-shell tools in the current turn.
    """
    path = Path(transcript_path)
    if not path.exists():
        return []

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

    hook_ids = hook_call_ids or set()
    records: list[ToolCallRecord] = []
    turn_count = 0
    in_relevant_turn = False

    # Track function_call entries by call_id for matching with outputs
    pending: dict[str, ToolCallRecord] = {}

    for line in raw_lines:
        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 == "event_msg":
            msg_type = payload.get("type", "")
            if msg_type == "task_complete":
                turn_count += 1
                in_relevant_turn = turn_count > turns_already_processed
            elif msg_type == "task_started":
                in_relevant_turn = turn_count >= turns_already_processed

        elif entry_type == "response_item" and in_relevant_turn:
            ptype = payload.get("type", "")

            if ptype == "function_call":
                call_id = payload.get("call_id", "")
                name = payload.get("name", "")
                # Skip shell tools (already captured via hooks) and known hook calls
                if name in _SHELL_TOOLS or call_id in hook_ids:
                    continue
                args_raw = payload.get("arguments", "{}")
                try:
                    args = json.loads(args_raw) if isinstance(args_raw, str) else args_raw
                except json.JSONDecodeError:
                    args = args_raw
                rec = ToolCallRecord(
                    call_id=call_id,
                    name=name,
                    input_data=args,
                    timestamp=timestamp,
                )
                pending[call_id] = rec

            elif ptype == "function_call_output":
                call_id = payload.get("call_id", "")
                if call_id in pending:
                    pending[call_id].output_data = payload.get("output", "")
                    records.append(pending.pop(call_id))

        # Also handle custom_tool_call (freeform tools like apply_patch)
        if entry_type == "response_item" and in_relevant_turn:
            ptype = payload.get("type", "")
            if ptype == "custom_tool_call":
                call_id = payload.get("call_id", "")
                name = payload.get("name", "")
                if call_id in hook_ids:
                    continue
                rec = ToolCallRecord(
                    call_id=call_id,
                    name=name,
                    input_data=payload.get("input", ""),
                    output_data=payload.get("output", ""),
                    timestamp=timestamp,
                )
                records.append(rec)

    # Flush any remaining pending calls (output may come after task_complete)
    records.extend(pending.values())

    return records
