"""
Codex transcript parser.

Parses JSONL transcript files generated by Codex CLI.
"""

from __future__ import annotations

import json
import logging
from datetime import UTC, datetime
from typing import Any

from gobby.sessions.transcripts.base import ParsedMessage, TokenUsage

logger = logging.getLogger(__name__)


class CodexTranscriptParser:
    """
    Parses JSONL transcript files from Codex.

    Implements the TranscriptParser protocol for Codex's transcript format.
    """

    def __init__(self, logger_instance: logging.Logger | None = None):
        self.logger = logger_instance or logger

    def extract_last_messages(
        self, turns: list[dict[str, Any]], num_pairs: int = 2
    ) -> list[dict[str, Any]]:
        messages: list[dict[str, str]] = []
        for turn in reversed(turns):
            role = turn.get("role")
            content = turn.get("content")

            if role in ["user", "assistant", "system"]:
                messages.insert(0, {"role": role, "content": str(content)})
                if len(messages) >= num_pairs * 2:
                    break
        return messages

    def extract_turns_since_clear(
        self, turns: list[dict[str, Any]], max_turns: int = 50
    ) -> list[dict[str, Any]]:
        # Codex likely uses a new session or clear command
        # For now, default to tail
        return turns[-max_turns:] if len(turns) > max_turns else turns

    def is_session_boundary(self, turn: dict[str, Any]) -> bool:
        return False

    def parse_line(self, line: str, index: int) -> ParsedMessage | None:
        if not line.strip():
            return None

        try:
            data = json.loads(line)
        except json.JSONDecodeError:
            self.logger.warning(f"Invalid JSON at line {index}")
            return None

        # Ensure data is a dict (JSON could be a string, number, etc.)
        if not isinstance(data, dict):
            self.logger.debug(f"Skipping non-object JSON at line {index}")
            return None

        timestamp = datetime.now(UTC)
        if "timestamp" in data:
            try:
                timestamp = datetime.fromisoformat(data["timestamp"].replace("Z", "+00:00"))
            except ValueError:
                pass

        # Assume simple schema for Codex: {"role": "user", "content": "..."}
        role = data.get("role")
        content = data.get("content", "")

        if not role:
            return None

        content_type = "text"
        tool_name = None

        # Check for tool use if Codex supports it in this format
        if role == "tool" or "tool_calls" in data:
            pass  # Placeholder for tool logic

        return ParsedMessage(
            index=index,
            role=role,
            content=str(content),
            content_type=content_type,
            tool_name=tool_name,
            tool_input=None,
            tool_result=None,
            timestamp=timestamp,
            raw_json=data,
            usage=self._extract_usage(data),
        )

    def _extract_usage(self, data: dict[str, Any]) -> TokenUsage | None:
        """Extract token usage from Codex message data."""
        # Codex CLI typically logs these at top level or in usage field
        usage_data = data if "input_tokens" in data else data.get("usage")

        if not usage_data:
            return None

        # Use explicit presence checks to handle 0 correctly
        input_tokens = (
            usage_data["input_tokens"]
            if "input_tokens" in usage_data
            else usage_data.get("inputTokens", 0)
        )
        output_tokens = (
            usage_data["output_tokens"]
            if "output_tokens" in usage_data
            else usage_data.get("outputTokens", 0)
        )
        cache_read_tokens = (
            usage_data["cached_tokens"]
            if "cached_tokens" in usage_data
            else usage_data.get("cachedTokens", 0)
        )
        # Total cost often provided directly
        total_cost_usd = (
            usage_data["cost"] if "cost" in usage_data else usage_data.get("total_cost")
        )

        return TokenUsage(
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            cache_read_tokens=cache_read_tokens,
            total_cost_usd=total_cost_usd,
        )

    def parse_lines(self, lines: list[str], start_index: int = 0) -> list[ParsedMessage]:
        parsed_messages = []
        current_index = start_index

        for line in lines:
            message = self.parse_line(line, current_index)
            if message:
                parsed_messages.append(message)
                current_index += 1

        return parsed_messages
