"""Parse Claude Code JSONL transcript files for LLM span data."""
from __future__ import annotations

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


@dataclass
class AssistantMessage:
    """LLM span data extracted from a single assistant message in the transcript."""

    message_id: str
    model: str
    content: list[dict[str, Any]]
    input_tokens: int
    output_tokens: int
    cache_read_tokens: int
    cache_creation_tokens: int
    timestamp: str

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


def count_lines(transcript_path: str) -> int:
    """Return the total number of lines in the transcript file.

    Uses the same raw-line counting as parse_assistant_messages so the two
    functions can be used together without an off-by-one when blank lines exist.
    """
    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_assistant_messages(
    transcript_path: str, start_line: int
) -> tuple[list[AssistantMessage], int]:
    """Parse assistant messages from the transcript after *start_line* (1-indexed).

    Returns (messages, total_lines_read).
    Only lines beyond start_line are processed; existing lines are skipped.
    """
    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[AssistantMessage] = []
    last_line = start_line

    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

        if entry.get("type") != "assistant":
            continue

        msg = entry.get("message", {})
        if not isinstance(msg, dict):
            continue

        usage = msg.get("usage") or {}
        messages.append(
            AssistantMessage(
                message_id=msg.get("id", ""),
                model=msg.get("model", ""),
                content=msg.get("content") or [],
                input_tokens=usage.get("input_tokens", 0),
                output_tokens=usage.get("output_tokens", 0),
                cache_read_tokens=usage.get("cache_read_input_tokens", 0),
                cache_creation_tokens=usage.get("cache_creation_input_tokens", 0),
                timestamp=entry.get("timestamp", ""),
            )
        )

    return messages, last_line
