"""Tests for codex.transcript -- Codex JSONL transcript parsing."""
from __future__ import annotations

import json

import pytest

from staso.codex.transcript import (
    LLMMessage,
    count_lines,
    find_transcript,
    parse_turn_llm_data,
)


def _write_jsonl(path, entries: list[dict]) -> None:
    lines = [json.dumps(e) for e in entries]
    path.write_text("\n".join(lines) + "\n")


def _make_session_file(sessions_dir, session_id: str, entries: list[dict] | None = None):
    """Create a Codex session transcript file in the expected directory structure."""
    day_dir = sessions_dir / "2026" / "04" / "02"
    day_dir.mkdir(parents=True, exist_ok=True)
    path = day_dir / f"rollout-2026-04-02T21-00-00-{session_id}.jsonl"
    if entries:
        _write_jsonl(path, entries)
    else:
        path.write_text("")
    return path


SAMPLE_TURN = [
    {
        "timestamp": "2026-04-02T15:42:40.279Z",
        "type": "event_msg",
        "payload": {"type": "task_started", "turn_id": "turn-1"},
    },
    {
        "timestamp": "2026-04-02T15:42:40.279Z",
        "type": "turn_context",
        "payload": {"turn_id": "turn-1", "model": "gpt-5.4"},
    },
    {
        "timestamp": "2026-04-02T15:42:40.541Z",
        "type": "response_item",
        "payload": {
            "type": "message",
            "role": "user",
            "content": [{"type": "input_text", "text": "Hello"}],
        },
    },
    {
        "timestamp": "2026-04-02T15:42:41.636Z",
        "type": "response_item",
        "payload": {
            "type": "message",
            "role": "assistant",
            "content": [{"type": "output_text", "text": "Hi there!"}],
        },
    },
    {
        "timestamp": "2026-04-02T15:42:41.658Z",
        "type": "event_msg",
        "payload": {
            "type": "token_count",
            "info": {
                "total_token_usage": {
                    "input_tokens": 12002,
                    "cached_input_tokens": 9728,
                    "output_tokens": 6,
                    "reasoning_output_tokens": 0,
                    "total_tokens": 12008,
                },
                "last_token_usage": {
                    "input_tokens": 12002,
                    "cached_input_tokens": 9728,
                    "output_tokens": 6,
                    "reasoning_output_tokens": 0,
                    "total_tokens": 12008,
                },
            },
        },
    },
    {
        "timestamp": "2026-04-02T15:42:42.360Z",
        "type": "event_msg",
        "payload": {
            "type": "task_complete",
            "turn_id": "turn-1",
            "last_agent_message": "Hi there!",
        },
    },
]

TURN_2 = [
    {
        "timestamp": "2026-04-02T15:45:00.000Z",
        "type": "event_msg",
        "payload": {"type": "task_started", "turn_id": "turn-2"},
    },
    {
        "timestamp": "2026-04-02T15:45:00.100Z",
        "type": "turn_context",
        "payload": {"turn_id": "turn-2", "model": "o3-pro"},
    },
    {
        "timestamp": "2026-04-02T15:45:01.000Z",
        "type": "response_item",
        "payload": {
            "type": "message",
            "role": "assistant",
            "content": [{"type": "output_text", "text": "Done."}],
        },
    },
    {
        "timestamp": "2026-04-02T15:45:01.500Z",
        "type": "event_msg",
        "payload": {
            "type": "token_count",
            "info": {
                "last_token_usage": {
                    "input_tokens": 5000,
                    "cached_input_tokens": 4000,
                    "output_tokens": 200,
                    "reasoning_output_tokens": 50,
                    "total_tokens": 5200,
                },
            },
        },
    },
    {
        "timestamp": "2026-04-02T15:45:02.000Z",
        "type": "event_msg",
        "payload": {"type": "task_complete", "turn_id": "turn-2"},
    },
]


class TestFindTranscript:
    def test_finds_matching_file(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.transcript.Path.home", lambda: tmp_path)
        sessions = tmp_path / ".codex" / "sessions"
        _make_session_file(sessions, "abc-123")

        result = find_transcript("abc-123")
        assert result is not None
        assert "abc-123" in str(result)

    def test_returns_none_when_not_found(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.transcript.Path.home", lambda: tmp_path)
        sessions = tmp_path / ".codex" / "sessions"
        sessions.mkdir(parents=True)

        assert find_transcript("nonexistent") is None

    def test_returns_none_when_sessions_dir_missing(self, tmp_path, monkeypatch):
        monkeypatch.setattr("staso.codex.transcript.Path.home", lambda: tmp_path)
        assert find_transcript("abc-123") is None


class TestCountLines:
    def test_counts_lines(self, tmp_path):
        path = tmp_path / "test.jsonl"
        path.write_text("line1\nline2\nline3\n")
        assert count_lines(str(path)) == 3

    def test_returns_zero_for_missing_file(self):
        assert count_lines("/nonexistent/file.jsonl") == 0

    def test_returns_zero_for_empty_file(self, tmp_path):
        path = tmp_path / "empty.jsonl"
        path.write_text("")
        assert count_lines(str(path)) == 0


class TestParseTurnLlmData:
    def test_extracts_single_turn(self, tmp_path):
        path = tmp_path / "transcript.jsonl"
        _write_jsonl(path, SAMPLE_TURN)

        messages, total_turns = parse_turn_llm_data(str(path), 0)
        assert len(messages) == 1
        assert total_turns == 1
        msg = messages[0]
        assert msg.model == "gpt-5.4"
        assert msg.input_tokens == 12002
        assert msg.output_tokens == 6
        assert msg.cached_input_tokens == 9728
        assert msg.reasoning_output_tokens == 0
        assert msg.total_tokens == 12008

    def test_extracts_assistant_content(self, tmp_path):
        path = tmp_path / "transcript.jsonl"
        _write_jsonl(path, SAMPLE_TURN)

        messages, _ = parse_turn_llm_data(str(path), 0)
        msg = messages[0]
        assert len(msg.content) == 1
        assert msg.content[0]["type"] == "output_text"
        assert msg.content[0]["text"] == "Hi there!"

    def test_skips_already_processed_turns(self, tmp_path):
        path = tmp_path / "transcript.jsonl"
        _write_jsonl(path, SAMPLE_TURN)

        # Already processed 1 turn — should find nothing new
        messages, total = parse_turn_llm_data(str(path), 1)
        assert len(messages) == 0
        assert total == 1

    def test_two_turns_parse_second_only(self, tmp_path):
        path = tmp_path / "transcript.jsonl"
        _write_jsonl(path, SAMPLE_TURN + TURN_2)

        # Already processed turn 1 — should get only turn 2
        messages, total = parse_turn_llm_data(str(path), 1)
        assert len(messages) == 1
        assert total == 2
        msg = messages[0]
        assert msg.model == "o3-pro"
        assert msg.input_tokens == 5000
        assert msg.output_tokens == 200
        assert msg.cached_input_tokens == 4000
        assert msg.reasoning_output_tokens == 50

    def test_two_turns_parse_both(self, tmp_path):
        path = tmp_path / "transcript.jsonl"
        _write_jsonl(path, SAMPLE_TURN + TURN_2)

        # 0 processed — should get both turns
        messages, total = parse_turn_llm_data(str(path), 0)
        assert len(messages) == 2
        assert total == 2
        assert messages[0].model == "gpt-5.4"
        assert messages[1].model == "o3-pro"

    def test_missing_file_returns_empty(self):
        messages, total = parse_turn_llm_data("/nonexistent/file.jsonl", 0)
        assert messages == []
        assert total == 0

    def test_no_token_data_produces_no_message(self, tmp_path):
        path = tmp_path / "transcript.jsonl"
        entries = [
            {"timestamp": "t1", "type": "event_msg", "payload": {"type": "task_started"}},
            {"timestamp": "t2", "type": "turn_context", "payload": {"model": "gpt-5.4"}},
            {"timestamp": "t3", "type": "event_msg", "payload": {"type": "task_complete"}},
        ]
        _write_jsonl(path, entries)

        messages, total = parse_turn_llm_data(str(path), 0)
        assert len(messages) == 0
        assert total == 1  # turn was seen, just no LLM data

    def test_uses_last_token_usage_over_total(self, tmp_path):
        """When both last_token_usage and total_token_usage exist, prefer last_token_usage."""
        path = tmp_path / "transcript.jsonl"
        entries = [
            {"timestamp": "t1", "type": "event_msg", "payload": {"type": "task_started"}},
            {"timestamp": "t2", "type": "turn_context", "payload": {"model": "gpt-5.4"}},
            {
                "timestamp": "t3",
                "type": "event_msg",
                "payload": {
                    "type": "token_count",
                    "info": {
                        "total_token_usage": {
                            "input_tokens": 99999,
                            "cached_input_tokens": 0,
                            "output_tokens": 99999,
                            "reasoning_output_tokens": 0,
                        },
                        "last_token_usage": {
                            "input_tokens": 100,
                            "cached_input_tokens": 50,
                            "output_tokens": 20,
                            "reasoning_output_tokens": 5,
                        },
                    },
                },
            },
            {"timestamp": "t4", "type": "event_msg", "payload": {"type": "task_complete"}},
        ]
        _write_jsonl(path, entries)

        messages, _ = parse_turn_llm_data(str(path), 0)
        assert len(messages) == 1
        assert messages[0].input_tokens == 100
        assert messages[0].output_tokens == 20


class TestLLMMessage:
    def test_total_tokens(self):
        msg = LLMMessage(
            model="gpt-5.4",
            input_tokens=100,
            output_tokens=50,
            cached_input_tokens=0,
            reasoning_output_tokens=0,
            content=[],
            timestamp="",
        )
        assert msg.total_tokens == 150
