"""Tests for the JSONL parser — scan_jsonl_entries, scan_jsonl_turns, _estimate_cost."""
from __future__ import annotations

import json
import os
import tempfile
import unittest
from datetime import datetime, timedelta, timezone
from pathlib import Path
from unittest.mock import patch

# We need to set MYBURN_DATA_DIR before importing server
_tmpdir = tempfile.mkdtemp()
os.environ["MYBURN_DATA_DIR"] = _tmpdir

from myburn.server import (
    CLAUDE_DIR,
    _estimate_cost,
    _scan_jsonl_entries,
    _scan_jsonl_turns,
)


def _ts(minutes_ago: int = 0) -> str:
    """Return an ISO timestamp N minutes ago (recent enough to pass cutoff)."""
    dt = datetime.now(timezone.utc) - timedelta(minutes=minutes_ago)
    return dt.strftime("%Y-%m-%dT%H:%M:%S.000Z")


def _make_entry(
    msg_id: str = "msg_001",
    model: str = "claude-sonnet-4-6",
    request_id: str = "req_001",
    input_tokens: int = 100,
    output_tokens: int = 50,
    cache_read: int = 0,
    cache_create: int = 0,
    timestamp: str | None = None,
    entry_type: str = "assistant",
    uuid: str = "",
    content: list | None = None,
) -> dict:
    """Build a minimal JSONL entry dict (Claude Code format)."""
    entry = {
        "type": entry_type,
        "timestamp": timestamp or _ts(5),
        "requestId": request_id,
        "uuid": uuid or f"uuid-{msg_id}",
        "message": {
            "id": msg_id,
            "model": model,
            "usage": {
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "cache_read_input_tokens": cache_read,
                "cache_creation_input_tokens": cache_create,
            },
            "content": content or [{"type": "text", "text": "hello"}],
        },
    }
    return entry


def _make_user_entry(timestamp: str | None = None, permission_mode: str | None = None) -> dict:
    """Build a minimal user entry."""
    entry = {
        "type": "user",
        "timestamp": timestamp or _ts(6),
        "message": {"content": [{"type": "text", "text": "test"}]},
    }
    if permission_mode:
        entry["permissionMode"] = permission_mode
    return entry


def _make_turn_duration(parent_uuid: str, duration_ms: int) -> dict:
    """Build a system/turn_duration entry."""
    return {
        "type": "system",
        "subtype": "turn_duration",
        "parentUuid": parent_uuid,
        "durationMs": duration_ms,
        "timestamp": _ts(4),
    }


class _TempJSONLDir:
    """Context manager that creates a fake CLAUDE_DIR with JSONL files.

    Usage:
        with _TempJSONLDir() as d:
            d.write_session("my-project", "sess1", [entry1, entry2])
            # Now _scan_jsonl_entries() will find the files
    """

    def __init__(self):
        self._tmpdir = tempfile.mkdtemp()
        self._root = Path(self._tmpdir)

    def __enter__(self):
        self._patcher = patch("myburn.server.CLAUDE_DIR", self._root)
        self._patcher.start()
        return self

    def __exit__(self, *args):
        self._patcher.stop()

    def write_session(
        self,
        project: str,
        session_id: str,
        entries: list[dict],
        partial_last: bool = False,
        subagent: bool = False,
    ):
        """Write entries to a JSONL file. If partial_last=True, last line has no trailing newline."""
        if subagent:
            path = self._root / project / session_id / "subagent" / "agent.jsonl"
        else:
            path = self._root / project / f"{session_id}.jsonl"
        path.parent.mkdir(parents=True, exist_ok=True)
        with open(path, "w") as fh:
            for i, entry in enumerate(entries):
                line = json.dumps(entry)
                if partial_last and i == len(entries) - 1:
                    fh.write(line)  # no trailing newline
                else:
                    fh.write(line + "\n")


# ===========================================================================
# _estimate_cost
# ===========================================================================
class TestEstimateCost(unittest.TestCase):
    def test_known_model(self):
        usage = {"input_tokens": 1000, "output_tokens": 500}
        cost = _estimate_cost("claude-sonnet-4-6", usage)
        # sonnet: $3/MTok input, $15/MTok output
        expected = (1000 * 3.0 + 500 * 15.0) / 1_000_000
        self.assertAlmostEqual(cost, expected)

    def test_unknown_model_uses_default(self):
        usage = {"input_tokens": 1000, "output_tokens": 500}
        cost = _estimate_cost("some-future-model-2027", usage)
        # default: $5/MTok input, $25/MTok output
        expected = (1000 * 5.0 + 500 * 25.0) / 1_000_000
        self.assertAlmostEqual(cost, expected)

    def test_cache_tokens(self):
        usage = {
            "input_tokens": 0,
            "output_tokens": 0,
            "cache_read_input_tokens": 10000,
            "cache_creation_input_tokens": 5000,
        }
        cost = _estimate_cost("claude-haiku-4-5-20251001", usage)
        # haiku: cache_read=$0.10/MTok, cache_create=$1.25/MTok
        expected = (10000 * 0.10 + 5000 * 1.25) / 1_000_000
        self.assertAlmostEqual(cost, expected)

    def test_zero_tokens(self):
        self.assertEqual(_estimate_cost("claude-sonnet-4-6", {}), 0.0)

    def test_missing_fields_treated_as_zero(self):
        usage = {"input_tokens": 1000}
        cost = _estimate_cost("claude-sonnet-4-6", usage)
        expected = (1000 * 3.0) / 1_000_000
        self.assertAlmostEqual(cost, expected)


# ===========================================================================
# _scan_jsonl_entries
# ===========================================================================
class TestScanJSONLEntries(unittest.TestCase):
    def test_basic_scan(self):
        """A single valid entry is parsed correctly."""
        entry = _make_entry(input_tokens=1000, output_tokens=500)
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0]["input_tokens"], 1000)
        self.assertEqual(results[0]["output_tokens"], 500)
        self.assertEqual(results[0]["project"], "proj1")
        self.assertEqual(results[0]["session_id"], "sess1")

    def test_dedup_same_msg_id(self):
        """Multiple entries with same msg_id + requestId → only first kept."""
        e1 = _make_entry(msg_id="msg_001", request_id="req_001", input_tokens=100)
        e2 = _make_entry(msg_id="msg_001", request_id="req_001", input_tokens=200)
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [e1, e2])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0]["input_tokens"], 100)  # first wins

    def test_dedup_different_request_id(self):
        """Same msg_id but different requestId → treated as distinct (FIX-16)."""
        e1 = _make_entry(msg_id="msg_001", request_id="req_001")
        e2 = _make_entry(msg_id="msg_001", request_id="req_002")
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [e1, e2])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 2)

    def test_dedup_no_request_id_fallback(self):
        """Entries with same msg_id but empty requestId → deduped by msg_id alone (FIX-16)."""
        e1 = _make_entry(msg_id="msg_001", request_id="", input_tokens=100)
        e2 = _make_entry(msg_id="msg_001", request_id="", input_tokens=200)
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [e1, e2])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 1)

    def test_dedup_cross_file(self):
        """Dedup works across files in same project."""
        e1 = _make_entry(msg_id="msg_001", request_id="req_001")
        e2 = _make_entry(msg_id="msg_001", request_id="req_001")
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [e1])
            d.write_session("proj1", "sess2", [e2])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 1)

    def test_dedup_cross_project(self):
        """Dedup works across different projects (subagent files)."""
        e1 = _make_entry(msg_id="msg_001", request_id="req_001")
        e2 = _make_entry(msg_id="msg_001", request_id="req_001")
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [e1])
            d.write_session("proj2", "sess2", [e2])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 1)

    def test_filter_non_assistant(self):
        """User and system entries are skipped."""
        e1 = _make_entry(entry_type="user")
        e2 = _make_entry(entry_type="system")
        e3 = _make_entry(entry_type="assistant", msg_id="msg_002")
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [e1, e2, e3])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 1)

    def test_filter_unknown_model(self):
        """Entries with model='unknown' or starting with '<' are skipped."""
        e1 = _make_entry(model="unknown", msg_id="msg_001")
        e2 = _make_entry(model="<model>", msg_id="msg_002")
        e3 = _make_entry(model="claude-sonnet-4-6", msg_id="msg_003")
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [e1, e2, e3])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0]["model"], "claude-sonnet-4-6")

    def test_filter_no_usage(self):
        """Entries without usage data are skipped."""
        entry = _make_entry()
        entry["message"]["usage"] = {}
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 0)

    def test_partial_line_skipped(self):
        """Last line without trailing newline is skipped (FIX-13)."""
        e1 = _make_entry(msg_id="msg_001")
        e2 = _make_entry(msg_id="msg_002")
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [e1, e2], partial_last=True)
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 1)  # only e1

    def test_malformed_json_skipped(self):
        """Invalid JSON lines are silently skipped."""
        entry = _make_entry()
        with _TempJSONLDir() as d:
            path = d._root / "proj1" / "sess1.jsonl"
            path.parent.mkdir(parents=True, exist_ok=True)
            with open(path, "w") as fh:
                fh.write("NOT VALID JSON\n")
                fh.write(json.dumps(entry) + "\n")
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 1)

    def test_empty_lines_skipped(self):
        """Empty lines don't cause errors."""
        entry = _make_entry()
        with _TempJSONLDir() as d:
            path = d._root / "proj1" / "sess1.jsonl"
            path.parent.mkdir(parents=True, exist_ok=True)
            with open(path, "w") as fh:
                fh.write("\n")
                fh.write(json.dumps(entry) + "\n")
                fh.write("\n")
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 1)

    def test_cache_tokens_in_output(self):
        """Cache tokens are captured in the output dict."""
        entry = _make_entry(cache_read=5000, cache_create=2000)
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(results[0]["cache_read_tokens"], 5000)
        self.assertEqual(results[0]["cache_creation_tokens"], 2000)

    def test_cost_calculation(self):
        """Cost is computed correctly for the entry."""
        entry = _make_entry(
            model="claude-sonnet-4-6",
            input_tokens=1_000_000,
            output_tokens=500_000,
        )
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry])
            results = _scan_jsonl_entries(full=True)
        # sonnet: $3/MTok input, $15/MTok output
        expected = (1_000_000 * 3.0 + 500_000 * 15.0) / 1_000_000
        self.assertAlmostEqual(results[0]["cost"], expected)

    def test_multiple_sessions_multiple_projects(self):
        """Entries from different sessions/projects get correct metadata."""
        e1 = _make_entry(msg_id="msg_001")
        e2 = _make_entry(msg_id="msg_002")
        with _TempJSONLDir() as d:
            d.write_session("project-a", "sess-abc", [e1])
            d.write_session("project-b", "sess-xyz", [e2])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 2)
        projects = {r["project"] for r in results}
        self.assertEqual(projects, {"project-a", "project-b"})

    def test_no_timestamp_skipped(self):
        """Entries without timestamp are skipped."""
        entry = _make_entry()
        entry["timestamp"] = ""
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 0)

    def test_no_msg_id_not_deduped(self):
        """Entries without message.id are not subject to dedup (each one counts)."""
        e1 = _make_entry(msg_id="", request_id="", input_tokens=100)
        e2 = _make_entry(msg_id="", request_id="", input_tokens=200)
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [e1, e2])
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(len(results), 2)

    def test_no_claude_dir(self):
        """Returns empty list if CLAUDE_DIR doesn't exist."""
        with patch("myburn.server.CLAUDE_DIR", Path("/nonexistent/path")):
            results = _scan_jsonl_entries(full=True)
        self.assertEqual(results, [])


# ===========================================================================
# _scan_jsonl_turns
# ===========================================================================
class TestScanJSONLTurns(unittest.TestCase):
    def test_basic_turn(self):
        """A single assistant entry produces one turn."""
        entry = _make_entry(input_tokens=1000, output_tokens=500)
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry])
            results = _scan_jsonl_turns(full=True)
        self.assertEqual(len(results), 1)
        self.assertEqual(results[0]["input_tokens"], 1000)
        self.assertEqual(results[0]["output_tokens"], 500)
        self.assertEqual(results[0]["model"], "claude-sonnet-4-6")

    def test_streaming_dedup(self):
        """Multiple entries with same message.id are merged into one turn."""
        ts = _ts(5)
        e1 = _make_entry(
            msg_id="msg_001", timestamp=ts,
            content=[{"type": "thinking", "text": "hmm"}],
        )
        e2 = _make_entry(
            msg_id="msg_001", timestamp=ts,
            content=[{"type": "text", "text": "hello"}],
        )
        e3 = _make_entry(
            msg_id="msg_001", timestamp=ts,
            content=[{"type": "tool_use", "name": "Read", "id": "t1"}],
        )
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [e1, e2, e3])
            results = _scan_jsonl_turns(full=True)
        self.assertEqual(len(results), 1)
        self.assertTrue(results[0]["has_thinking"])
        self.assertEqual(results[0]["tool_count"], 1)
        tools = json.loads(results[0]["tools_used"])
        self.assertEqual(tools, ["Read"])

    def test_tool_extraction(self):
        """Tool names are extracted from content blocks."""
        entry = _make_entry(
            content=[
                {"type": "tool_use", "name": "Read", "id": "t1"},
                {"type": "tool_use", "name": "Edit", "id": "t2"},
                {"type": "text", "text": "done"},
            ]
        )
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry])
            results = _scan_jsonl_turns(full=True)
        self.assertEqual(results[0]["tool_count"], 2)
        tools = json.loads(results[0]["tools_used"])
        self.assertIn("Read", tools)
        self.assertIn("Edit", tools)

    def test_thinking_detection(self):
        """Thinking content blocks set has_thinking flag."""
        entry = _make_entry(
            content=[{"type": "thinking", "text": "let me think"}]
        )
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry])
            results = _scan_jsonl_turns(full=True)
        self.assertTrue(results[0]["has_thinking"])

    def test_no_thinking(self):
        """Turns without thinking blocks have has_thinking=False."""
        entry = _make_entry(content=[{"type": "text", "text": "hello"}])
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry])
            results = _scan_jsonl_turns(full=True)
        self.assertFalse(results[0]["has_thinking"])

    def test_permission_mode_tracking(self):
        """Permission mode is inherited from preceding user entry."""
        ts1 = _ts(10)
        ts2 = _ts(9)
        ts3 = _ts(8)
        user = _make_user_entry(timestamp=ts1, permission_mode="bypassPermissions")
        assistant = _make_entry(msg_id="msg_001", timestamp=ts2)
        # Second assistant with default mode (no new user entry)
        assistant2 = _make_entry(msg_id="msg_002", timestamp=ts3)
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [user, assistant, assistant2])
            results = _scan_jsonl_turns(full=True)
        self.assertEqual(len(results), 2)
        # Both should inherit bypassPermissions since no user entry reset it
        self.assertEqual(results[0]["permission_mode"], "bypassPermissions")
        self.assertEqual(results[1]["permission_mode"], "bypassPermissions")

    def test_default_permission_mode(self):
        """Without user entries, permission mode defaults to 'default'."""
        entry = _make_entry()
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry])
            results = _scan_jsonl_turns(full=True)
        self.assertEqual(results[0]["permission_mode"], "default")

    def test_duration_linkage(self):
        """Duration is linked via parentUuid → last entry's uuid."""
        uuid = "my-uuid-123"
        entry = _make_entry(uuid=uuid)
        dur = _make_turn_duration(parent_uuid=uuid, duration_ms=4200)
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry, dur])
            results = _scan_jsonl_turns(full=True)
        self.assertEqual(results[0]["duration_ms"], 4200)

    def test_no_duration(self):
        """Turn without matching duration entry → None."""
        entry = _make_entry()
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry])
            results = _scan_jsonl_turns(full=True)
        self.assertIsNone(results[0]["duration_ms"])

    def test_subagent_detection(self):
        """Files in subagent directories set is_subagent=True."""
        entry = _make_entry()
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry], subagent=True)
            results = _scan_jsonl_turns(full=True)
        self.assertEqual(len(results), 1)
        self.assertTrue(results[0]["is_subagent"])

    def test_partial_line_in_turns(self):
        """Partial lines (no trailing newline) are skipped in turn scanning (FIX-13)."""
        e1 = _make_entry(msg_id="msg_001")
        e2 = _make_entry(msg_id="msg_002")
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [e1, e2], partial_last=True)
            results = _scan_jsonl_turns(full=True)
        self.assertEqual(len(results), 1)  # only e1

    def test_turn_index_increments(self):
        """Turn index increments per file."""
        ts1 = _ts(10)
        ts2 = _ts(9)
        e1 = _make_entry(msg_id="msg_001", timestamp=ts1)
        e2 = _make_entry(msg_id="msg_002", timestamp=ts2)
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [e1, e2])
            results = _scan_jsonl_turns(full=True)
        self.assertEqual(results[0]["turn_index"], 0)
        self.assertEqual(results[1]["turn_index"], 1)

    def test_cost_in_turns(self):
        """Turns include cost estimation."""
        entry = _make_entry(
            model="claude-sonnet-4-6",
            input_tokens=1_000_000,
            output_tokens=500_000,
        )
        with _TempJSONLDir() as d:
            d.write_session("proj1", "sess1", [entry])
            results = _scan_jsonl_turns(full=True)
        expected = (1_000_000 * 3.0 + 500_000 * 15.0) / 1_000_000
        self.assertAlmostEqual(results[0]["cost"], expected)

    def test_no_claude_dir(self):
        """Returns empty list if CLAUDE_DIR doesn't exist."""
        with patch("myburn.server.CLAUDE_DIR", Path("/nonexistent/path")):
            results = _scan_jsonl_turns(full=True)
        self.assertEqual(results, [])


if __name__ == "__main__":
    unittest.main()
