"""Tests for spens.summarizer (session stats, token accounting and costing).

The cost and tool-call cases here are regression tests: each one pins a
specific way the aggregation used to over- or under-count.
"""

from __future__ import annotations

import json
from pathlib import Path

from spens import pricing
from spens.pricing import Price, build_table
from spens.sessions import load_trace_index
from spens.summarizer import (
    CACHE_READ_MULTIPLIER,
    CACHE_WRITE_MULTIPLIER,
    _aggregate_chat_stats,
    _count_tool_calls,
    _parse_iso,
    _price_for,
    _split_calls,
    compute_session_summary,
    format_summary_for_cli,
    read_summary,
    write_summary,
)


def _write_traces(session_dir: Path, rows: list[dict]) -> None:
    traces = session_dir / "traces"
    traces.mkdir(parents=True, exist_ok=True)
    (traces / "captured.jsonl").write_text(
        "".join(json.dumps(r) + "\n" for r in rows), encoding="utf-8"
    )


def _index(tmp_path: Path, rows: list[dict]):
    _write_traces(tmp_path, rows)
    return load_trace_index(tmp_path)


def _call(model: str, **usage) -> dict:
    return {"model": model, "usage": usage or None}


# -- _price_for (fallback table) ---------------------------------------------
#
# The unit suite runs with the Portkey dataset switched off (see
# tests/unit/conftest.py), so these exercise the built-in fallback table;
# the dataset's own resolution lives in test_pricing.py.


def _rates(model, context_tokens: int = 0, table=None) -> tuple[float, float]:
    price = _price_for(model, context_tokens, table)
    return price.input_per_1m, price.output_per_1m


def test_price_for_prefers_the_longest_matching_prefix() -> None:
    """A shorter prefix must not capture a more specific model id."""
    assert _rates("gpt-4o-mini") == (0.15, 0.60)
    assert _rates("gpt-4o") == (2.50, 10.00)
    assert _rates("o1-mini") == (1.10, 4.40)
    assert _rates("o1") == (15.00, 60.00)


def test_price_for_matches_dated_and_current_model_ids() -> None:
    assert _rates("claude-opus-4-20250514") == (15.00, 75.00)
    assert _rates("claude-opus-4-8") == (5.00, 25.00)
    assert _rates("claude-opus-5") == (5.00, 25.00)
    assert _rates("claude-sonnet-5") == (2.00, 10.00)
    assert _rates("claude-haiku-4-5") == (1.00, 5.00)


def test_price_for_handles_vendor_prefixed_ids() -> None:
    """OpenRouter-style ids are not prefixes, so a substring match applies."""
    assert _rates("anthropic/claude-sonnet-4") == (3.00, 15.00)


def test_price_for_unknown_model_is_free_not_mispriced() -> None:
    assert _rates("some-future-model") == (0.0, 0.0)
    assert _rates("") == (0.0, 0.0)
    assert _rates(None) == (0.0, 0.0)


def test_price_for_normalises_separators() -> None:
    assert _rates("GPT-4O-MINI") == (0.15, 0.60)
    assert _rates("gpt_4o_mini") == (0.15, 0.60)


def test_the_fallback_table_has_no_cache_rates() -> None:
    """Which is what tells the aggregation to estimate them from input."""
    price = _price_for("claude-opus-4")
    assert price.cache_read_per_1m is None
    assert price.cache_write_per_1m is None


# -- _price_for (Portkey dataset) --------------------------------------------


def test_dataset_price_wins_over_the_fallback_table() -> None:
    """A stale hand-maintained entry must not override fetched pricing."""
    table = build_table([("anthropic", {"claude-opus-4": Price(7.0, 21.0)})])
    assert _rates("claude-opus-4", 0, table) == (7.0, 21.0)


def test_fallback_table_covers_models_the_dataset_omits() -> None:
    table = build_table([("anthropic", {"claude-opus-4": Price(7.0, 21.0)})])
    assert _rates("gpt-4o", 0, table) == (2.50, 10.00)


# -- cost aggregation --------------------------------------------------------


def test_cache_cost_is_linear_in_the_number_of_calls() -> None:
    """Regression: cached tokens were costed against the running total, so
    N identical calls cost O(N^2) instead of N x the per-call price."""
    in_price = _price_for("claude-opus-4").input_per_1m
    per_call = 1_000_000 * in_price * CACHE_READ_MULTIPLIER / 1_000_000
    for n in (1, 2, 3, 10, 50):
        calls = [_call("claude-opus-4", cache_read_input_tokens=1_000_000)] * n
        stats = _aggregate_chat_stats(calls)
        assert stats["estimated_cost_usd"] == round(per_call * n, 6), n


def test_cache_write_uses_its_own_multiplier() -> None:
    in_price = _price_for("claude-opus-4").input_per_1m
    stats = _aggregate_chat_stats([_call("claude-opus-4", cache_creation_input_tokens=1_000_000)])
    assert stats["estimated_cost_usd"] == round(in_price * CACHE_WRITE_MULTIPLIER, 6)
    assert stats["cache_write"] == 1_000_000


def test_cost_uses_input_and_output_prices() -> None:
    stats = _aggregate_chat_stats(
        [_call("gpt-4o", prompt_tokens=1_000_000, completion_tokens=1_000_000)]
    )
    assert stats["estimated_cost_usd"] == round(2.50 + 10.00, 6)
    assert stats["tokens_in"] == 1_000_000
    assert stats["tokens_out"] == 1_000_000


def test_provider_reported_cost_wins_over_the_pricing_table() -> None:
    stats = _aggregate_chat_stats(
        [_call("gpt-4o", prompt_tokens=1_000_000, cost=0.42)]
    )
    assert stats["estimated_cost_usd"] == 0.42


def test_zero_provider_cost_falls_back_to_the_pricing_table() -> None:
    """Regression: proxies fronting Anthropic inject ``cost: 0`` because they
    cannot price the upstream model; taking that at face value zeroed the
    whole session's estimate so the cost line never appeared in the recap."""
    stats = _aggregate_chat_stats(
        [_call("claude-sonnet-5", prompt_tokens=90, completion_tokens=9,
               cache_read_input_tokens=42_667, cost=0)]
    )
    expected = round(
        (90 * 2.00 + 9 * 10.00 + 42_667 * 2.00 * CACHE_READ_MULTIPLIER) / 1e6, 6
    )
    assert stats["estimated_cost_usd"] == expected
    assert stats["models"]["claude-sonnet-5"]["cost_usd"] == expected


def test_dataset_cache_rates_replace_the_multiplier_estimate() -> None:
    """A published cache rate beats guessing a tenth of the input price."""
    table = build_table(
        [("anthropic", {"m": Price(10.0, 20.0, cache_read_per_1m=0.5, cache_write_per_1m=12.0)})]
    )
    stats = _aggregate_chat_stats(
        [_call("m", cache_read_input_tokens=1_000_000, cache_creation_input_tokens=1_000_000)],
        table,
    )
    assert stats["estimated_cost_usd"] == round(0.5 + 12.0, 6)


def test_multipliers_apply_when_the_dataset_has_no_cache_rate() -> None:
    table = build_table([("anthropic", {"m": Price(10.0, 20.0)})])
    stats = _aggregate_chat_stats([_call("m", cache_read_input_tokens=1_000_000)], table)
    assert stats["estimated_cost_usd"] == round(10.0 * CACHE_READ_MULTIPLIER, 6)


def test_context_band_is_chosen_from_the_calls_own_token_counts() -> None:
    """Gemini bands on the context sent, which includes the cached prefix."""
    table = build_table(
        [("google", {"m-lte-128k": Price(1.0, 1.0), "m-gt-128k": Price(100.0, 100.0)})]
    )
    cheap = _aggregate_chat_stats([_call("m", prompt_tokens=50_000)], table)
    assert cheap["estimated_cost_usd"] == round(50_000 * 1.0 / 1e6, 6)

    # 100k fresh input plus a 100k cached prefix is a 200k context, so the
    # call is billed in the larger band.
    dear = _aggregate_chat_stats(
        [_call("m", prompt_tokens=100_000, cache_read_input_tokens=100_000)], table
    )
    assert dear["estimated_cost_usd"] == round(
        (100_000 * 100.0 + 100_000 * 100.0 * CACHE_READ_MULTIPLIER) / 1e6, 6
    )


def test_non_numeric_provider_cost_is_ignored() -> None:
    stats = _aggregate_chat_stats(
        [_call("gpt-4o", prompt_tokens=1_000_000, cost="n/a")]
    )
    assert stats["estimated_cost_usd"] == round(2.50, 6)


def test_per_model_breakdown_records_cost() -> None:
    """Regression: models[*]['cost_usd'] was initialised and never written."""
    stats = _aggregate_chat_stats([
        _call("gpt-4o", prompt_tokens=1_000_000),
        _call("gpt-4o", prompt_tokens=1_000_000),
        _call("claude-opus-4", prompt_tokens=1_000_000),
    ])
    assert stats["models"]["gpt-4o"]["calls"] == 2
    assert stats["models"]["gpt-4o"]["cost_usd"] == round(2.50 * 2, 6)
    assert stats["models"]["claude-opus-4"]["cost_usd"] == round(15.00, 6)
    assert stats["estimated_cost_usd"] == round(5.00 + 15.00, 6)


def test_unpriced_model_costs_nothing_but_still_counts_tokens() -> None:
    stats = _aggregate_chat_stats(
        [_call("future-model-9", prompt_tokens=1000, completion_tokens=500)]
    )
    assert stats["estimated_cost_usd"] == 0.0
    assert stats["tokens_in"] == 1000
    assert stats["tokens_out"] == 500


def test_calls_without_usage_are_counted_as_turns() -> None:
    stats = _aggregate_chat_stats([_call("gpt-4o"), _call("gpt-4o")])
    assert stats["turns"] == 2
    assert stats["calls"] == 2
    assert stats["tokens_in"] == 0
    assert stats["estimated_cost_usd"] == 0.0


def test_empty_session_aggregates_to_zero() -> None:
    stats = _aggregate_chat_stats([])
    assert stats["turns"] == 0
    assert stats["total_tokens"] == 0
    assert stats["models"] == {}


# -- token totals ------------------------------------------------------------


def test_total_tokens_sums_across_mixed_providers() -> None:
    """Regression: one provider's total_tokens became the whole session's
    total, dropping every call that does not report one."""
    stats = _aggregate_chat_stats([
        _call("gpt-4o", prompt_tokens=100, completion_tokens=50, total_tokens=150),
        _call("claude-opus-4", input_tokens=1000, output_tokens=500),
    ])
    assert stats["tokens_in"] == 1100
    assert stats["tokens_out"] == 550
    assert stats["total_tokens"] == 1650


def test_total_tokens_includes_anthropic_cache_dimensions() -> None:
    """Anthropic reports cached tokens outside input_tokens."""
    stats = _aggregate_chat_stats([
        _call(
            "claude-opus-4",
            input_tokens=10,
            output_tokens=5,
            cache_read_input_tokens=100,
            cache_creation_input_tokens=20,
        )
    ])
    assert stats["total_tokens"] == 135
    assert stats["cache_read"] == 100
    assert stats["cache_write"] == 20


def test_total_tokens_includes_cache_even_when_provider_reports_a_total() -> None:
    """Regression: an OpenAI-compatible proxy fronting Anthropic reports
    ``total_tokens`` as prompt + completion only, while also passing the
    Anthropic cache dimensions; trusting its total dropped the cache from
    the session total (2,475 shown next to 1.1M cache-read tokens)."""
    stats = _aggregate_chat_stats([
        _call(
            "claude-sonnet-5",
            prompt_tokens=90,
            completion_tokens=9,
            total_tokens=99,
            cache_read_input_tokens=42_667,
            cache_creation_input_tokens=1_855,
        )
    ])
    assert stats["total_tokens"] == 99 + 42_667 + 1_855
    assert stats["tokens_in"] == 90
    assert stats["tokens_out"] == 9


# -- tool-call counting ------------------------------------------------------


def test_tool_calls_are_not_inflated_by_replayed_history(tmp_path: Path) -> None:
    """Regression: each request replays the whole conversation, so counting
    request bodies grew the total quadratically with the turn count."""
    rows: list[dict] = []
    history: list[dict] = []
    for turn in range(3):
        history.append(
            {"role": "assistant", "content": [{"type": "tool_use", "id": f"t{turn}", "name": "bash"}]}
        )
        rows.append({
            "type": "request",
            "id": f"r{turn}",
            "timestamp": f"2026-01-01T00:00:0{turn}Z",
            "url": "https://api.anthropic.com/v1/messages",
            "body": {"messages": list(history)},
        })
        rows.append({
            "type": "response",
            "request_id": f"r{turn}",
            "timestamp": f"2026-01-01T00:00:0{turn}Z",
            "body": {"content": [{"type": "tool_use", "id": f"t{turn}", "name": "bash"}]},
        })
    assert _count_tool_calls(_index(tmp_path, rows)) == 3


def test_tool_calls_from_buffered_anthropic_response(tmp_path: Path) -> None:
    rows = [
        {"type": "request", "id": "r1", "timestamp": "t0", "url": "u"},
        {"type": "response", "request_id": "r1", "timestamp": "t1", "body": {"content": [
            {"type": "text", "text": "thinking about it"},
            {"type": "tool_use", "id": "a", "name": "read"},
            {"type": "tool_use", "id": "b", "name": "write"},
        ]}},
    ]
    assert _count_tool_calls(_index(tmp_path, rows)) == 2


def test_tool_calls_from_buffered_openai_response(tmp_path: Path) -> None:
    rows = [
        {"type": "request", "id": "r1", "timestamp": "t0", "url": "u"},
        {"type": "response", "request_id": "r1", "timestamp": "t1", "body": {"choices": [
            {"message": {"tool_calls": [{"id": "1"}, {"id": "2"}]}}
        ]}},
    ]
    assert _count_tool_calls(_index(tmp_path, rows)) == 2


def test_tool_calls_from_buffered_responses_api(tmp_path: Path) -> None:
    rows = [
        {"type": "request", "id": "r1", "timestamp": "t0", "url": "u"},
        {"type": "response", "request_id": "r1", "timestamp": "t1", "body": {"output": [
            {"type": "message"},
            {"type": "function_call", "name": "f"},
            {"type": "custom_tool_call", "name": "g"},
        ]}},
    ]
    assert _count_tool_calls(_index(tmp_path, rows)) == 2


def test_tool_calls_from_streamed_anthropic_response(tmp_path: Path) -> None:
    rows = [
        {"type": "request", "id": "r1", "timestamp": "t0", "url": "u"},
        {"type": "response_chunk", "request_id": "r1", "timestamp": "t1", "content": {
            "type": "content_block_start", "index": 0,
            "content_block": {"type": "text", "text": ""}}},
        {"type": "response_chunk", "request_id": "r1", "timestamp": "t2", "content": {
            "type": "content_block_start", "index": 1,
            "content_block": {"type": "tool_use", "id": "a", "name": "read"}}},
    ]
    assert _count_tool_calls(_index(tmp_path, rows)) == 1


def test_streamed_openai_argument_deltas_count_as_one_call(tmp_path: Path) -> None:
    """Arguments arrive as many deltas for the same tool-call index."""
    rows = [{"type": "request", "id": "r1", "timestamp": "t0", "url": "u"}]
    for i, fragment in enumerate(['{"pa', 'th": "', 'a.txt"}']):
        rows.append({"type": "response_chunk", "request_id": "r1", "timestamp": f"t{i + 1}",
                     "content": {"choices": [{"delta": {"tool_calls": [
                         {"index": 0, "function": {"arguments": fragment}}]}}]}})
    rows.append({"type": "response_chunk", "request_id": "r1", "timestamp": "t9",
                 "content": {"choices": [{"delta": {"tool_calls": [
                     {"index": 1, "function": {"name": "other", "arguments": "{}"}}]}}]}})
    assert _count_tool_calls(_index(tmp_path, rows)) == 2


def test_no_tool_calls_in_a_plain_session(tmp_path: Path) -> None:
    rows = [
        {"type": "request", "id": "r1", "timestamp": "t0", "url": "u"},
        {"type": "response", "request_id": "r1", "timestamp": "t1",
         "body": {"content": [{"type": "text", "text": "hello"}]}},
    ]
    assert _count_tool_calls(_index(tmp_path, rows)) == 0


# -- call splitting ----------------------------------------------------------


def test_split_calls_separates_chat_from_plain_http(tmp_path: Path) -> None:
    rows = [
        {"type": "request", "id": "r1", "timestamp": "t0",
         "url": "https://api.anthropic.com/v1/messages",
         "body": {"model": "claude-opus-4", "messages": []}},
        {"type": "response", "request_id": "r1", "timestamp": "t1",
         "body": {"usage": {"input_tokens": 7}}},
        {"type": "request", "id": "r2", "timestamp": "t2", "method": "GET",
         "url": "https://pypi.org/simple/", "body": None},
        {"type": "response", "request_id": "r2", "timestamp": "t3", "status_code": 200},
    ]
    result = _split_calls(_index(tmp_path, rows))
    assert len(result["chat_calls"]) == 1
    assert result["chat_calls"][0]["model"] == "claude-opus-4"
    assert result["chat_calls"][0]["usage"] == {"input_tokens": 7}
    assert len(result["http_calls"]) == 1
    assert result["http_calls"][0]["status_code"] == 200


def test_split_calls_reads_usage_from_streamed_chunks(tmp_path: Path) -> None:
    """Anthropic splits usage across message_start and message_delta."""
    rows = [
        {"type": "request", "id": "r1", "timestamp": "t0",
         "url": "https://api.anthropic.com/v1/messages", "body": {"messages": []}},
        {"type": "response_chunk", "request_id": "r1", "timestamp": "t1", "content": {
            "type": "message_start",
            "message": {"model": "claude-opus-4", "usage": {"input_tokens": 100}}}},
        {"type": "response_chunk", "request_id": "r1", "timestamp": "t2", "content": {
            "type": "message_delta", "usage": {"output_tokens": 20}}},
    ]
    call = _split_calls(_index(tmp_path, rows))["chat_calls"][0]
    assert call["model"] == "claude-opus-4"
    assert call["usage"]["input_tokens"] == 100
    assert call["usage"]["output_tokens"] == 20


# -- timestamps --------------------------------------------------------------


def test_parse_iso_handles_z_suffix_and_nanoseconds() -> None:
    a = _parse_iso("2026-09-01T14:34:04Z")
    b = _parse_iso("2026-09-01T14:34:09.292565116+00:00")
    assert (b - a).total_seconds() == 5.292565


# -- end to end --------------------------------------------------------------


def test_compute_session_summary_end_to_end(tmp_path: Path) -> None:
    session = tmp_path / "abc123"
    _write_traces(session, [
        {"type": "request", "id": "r1", "timestamp": "2026-01-01T00:00:00Z",
         "url": "https://api.anthropic.com/v1/messages",
         "body": {"model": "claude-opus-4", "messages": []}},
        {"type": "response", "request_id": "r1", "timestamp": "2026-01-01T00:00:01Z",
         "body": {"content": [{"type": "tool_use", "id": "a", "name": "bash"}],
                  "usage": {"input_tokens": 1_000_000, "output_tokens": 0}}},
    ])
    (session / "traces" / "request_log.jsonl").write_text(
        '{"method": "GET", "url": "https://pypi.org", "status_code": 200}\n', encoding="utf-8"
    )
    meta = session / "nono-audit" / "sessions"
    meta.mkdir(parents=True)
    (meta / "s.json").write_text(json.dumps({
        "name": "demo", "command": ["codex", "run"], "exit_code": 0,
        "started": "2026-01-01T00:00:00Z",
    }), encoding="utf-8")
    audit = session / "nono-audit" / "audit" / "aa"
    audit.mkdir(parents=True)
    (audit / "session.json").write_text('{"ended": "2026-01-01T00:02:30Z"}', encoding="utf-8")
    (audit / "audit-events.ndjson").write_text('{"e": 1}\n{"e": 2}\n', encoding="utf-8")

    summary = compute_session_summary(session)

    assert summary["session_id"] == "abc123"
    assert summary["name"] == "demo"
    assert summary["command"] == "codex run"
    assert summary["exit_code"] == 0
    assert summary["duration_seconds"] == 150.0
    assert summary["llm"]["calls"] == 1
    assert summary["llm"]["tool_calls"] == 1
    assert summary["llm"]["tokens_in"] == 1_000_000
    assert summary["llm"]["estimated_cost_usd"] == 15.00
    assert summary["llm"]["models"]["claude-opus-4"]["cost_usd"] == 15.00
    assert summary["network"]["request_log_entries"] == 1
    assert summary["audit"]["event_count"] == 2
    assert summary["files"]["available"] is False


def test_summary_records_where_the_prices_came_from(tmp_path: Path) -> None:
    """So a cost figure in an archived summary stays explainable."""
    summary = compute_session_summary(tmp_path / "s")
    provenance = summary["llm"]["pricing"]
    assert provenance["dataset"] == "portkey"
    assert set(provenance["providers"]) == set(pricing.PROVIDERS)


def test_compute_session_summary_of_an_empty_dir(tmp_path: Path) -> None:
    summary = compute_session_summary(tmp_path / "empty")
    assert summary["llm"]["calls"] == 0
    assert summary["duration_seconds"] is None
    assert summary["files"]["available"] is False


def test_write_then_read_summary_round_trip(tmp_path: Path) -> None:
    session = tmp_path / "s1"
    session.mkdir()
    path, summary = write_summary(session)
    assert path == session / "session_summary.json"
    assert read_summary(session) == summary


def test_read_summary_missing(tmp_path: Path) -> None:
    assert read_summary(tmp_path) is None


def test_format_summary_for_cli_reports_cost_and_models() -> None:
    text = format_summary_for_cli({
        "session_id": "abc",
        "command": "codex run",
        "duration_seconds": 90,
        "exit_code": 0,
        "llm": {
            "turns": 2, "calls": 2, "tool_calls": 3,
            "tokens_in": 1000, "tokens_out": 200, "total_tokens": 1200,
            "cache_read": 500, "cache_write": 10,
            "estimated_cost_usd": 1.2345,
            "models": {"claude-opus-4": {
                "calls": 2, "tokens_in": 1000, "tokens_out": 200, "cost_usd": 1.2345}},
        },
        "network": {"http_calls": 4, "request_log_entries": 9},
        "files": {"available": True, "total": 3, "created": 1, "modified": 1, "deleted": 1,
                  "workspace": 2},
        "audit": {"event_count": 7},
    })
    assert "Session recap" in text
    assert "session    abc" in text
    assert "duration   1m 30s" in text
    assert "exit code  OK" in text
    assert "3 tool calls" in text
    assert "1,000 in · 200 out · 1,200 total" in text
    assert "500 read · 10 write" in text
    assert "cost       $1.2345 (est.)" in text
    # the model breakdown row carries the per-model cost
    assert "claude-opus-4" in text
    assert "$1.2345" in text.split("models")[1]
    assert "3 changes · 1 created · 1 modified · 1 deleted" in text
    assert "2 under /workspace" in text
    assert "4 http calls · 9 request-log entries" in text
    assert "7 events" in text
    # plain rendering: never any ANSI escape codes (pipes/logs/goldens)
    assert "\x1b" not in text


def test_format_summary_for_cli_tolerates_an_empty_summary() -> None:
    text = format_summary_for_cli({})
    assert "Session recap" in text
    assert "0 turns" in text
    assert "\x1b" not in text


def test_streamed_anthropic_usage_keeps_input_and_cache_counts(tmp_path: Path) -> None:
    """Regression: a message_delta's usage was treated as a whole-usage
    replacement, discarding the input and cache counts from message_start --
    the bulk of a cached session's cost."""
    rows = [
        {"type": "request", "id": "r1", "timestamp": "t0",
         "url": "https://api.anthropic.com/v1/messages", "body": {"messages": []}},
        {"type": "response_chunk", "request_id": "r1", "timestamp": "t1", "content": {
            "type": "message_start",
            "message": {"model": "claude-opus-4", "usage": {
                "input_tokens": 100, "cache_read_input_tokens": 5000}}}},
        {"type": "response_chunk", "request_id": "r1", "timestamp": "t2", "content": {
            "type": "message_delta", "usage": {"output_tokens": 20}}},
    ]
    stats = _aggregate_chat_stats(_split_calls(_index(tmp_path, rows))["chat_calls"])
    assert stats["tokens_in"] == 100
    assert stats["tokens_out"] == 20
    assert stats["cache_read"] == 5000
    assert stats["estimated_cost_usd"] > 0
