"""Compute and persist session summary statistics (including costs).

The summary JSON is written to ``<session_dir>/session_summary.json`` at the
end of every spens run so the viewer and CLI can display a fast recap without
re-parsing all traces.

Prices come from the Portkey dataset via :mod:`spens.pricing`; the table
below is the fallback for models that dataset does not list.
"""

from __future__ import annotations

import contextlib
import json
from datetime import datetime
from pathlib import Path
from typing import Any

from spens import pricing, ui
from spens.pricing import Price, PricingTable
from spens.sessions import (
    TraceIndex,
    is_chat_request,
    load_audit_session,
    load_nono_session_meta,
    load_request_log,
    load_trace_index,
    read_json,
)

# Message part types that represent a tool invocation across APIs.
TOOLISH_TYPES = ("tool_use", "tool_call", "custom_tool_call", "function_call")

# Cost multipliers applied to the *input* price for cached tokens, used only
# when the pricing data does not give cached tokens their own rate: a cache
# read is billed at roughly a tenth of the input rate, and writing to the
# cache at roughly 1.25x.
CACHE_READ_MULTIPLIER = 0.1
CACHE_WRITE_MULTIPLIER = 1.25

# Fallback per-1M-token pricing (input / output) in USD for models the
# Portkey dataset does not list -- it can lag a brand-new model id by days,
# and a hand-checked price beats costing a session at $0.  Keys are
# normalised model-id prefixes, matched longest-first, so
# ``claude-opus-4-20250514`` matches ``claude-opus-4`` while
# ``gpt-4o-mini`` matches ``gpt-4o-mini`` rather than ``gpt-4o``.  Prices are
# best-effort estimates -- providers change them, and an unlisted model is
# costed at $0 rather than guessed.
MODEL_PRICING: dict[str, tuple[float, float]] = {
    # Anthropic -- current models
    "claude-fable-5-1": (10.00, 50.00),
    "claude-fable-5": (10.00, 50.00),
    "claude-mythos-5-1": (10.00, 50.00),
    "claude-opus-5": (5.00, 25.00),
    "claude-opus-4-8": (5.00, 25.00),
    "claude-opus-4-7": (5.00, 25.00),
    "claude-opus-4-6": (5.00, 25.00),
    "claude-sonnet-5": (2.00, 10.00),
    "claude-sonnet-4-6": (3.00, 15.00),
    "claude-haiku-4-5": (1.00, 5.00),
    # Anthropic -- older models
    "claude-opus-4": (15.00, 75.00),
    "claude-sonnet-4": (3.00, 15.00),
    "claude-3-5-sonnet": (3.00, 15.00),
    "claude-3-5-haiku": (0.80, 4.00),
    "claude-3-opus": (15.00, 75.00),
    "claude-3-sonnet": (3.00, 15.00),
    "claude-3-haiku": (0.80, 4.00),
    # OpenAI
    "gpt-4o": (2.50, 10.00),
    "gpt-4o-mini": (0.15, 0.60),
    "o3-mini": (1.10, 4.40),
    "o1": (15.00, 60.00),
    "o1-mini": (1.10, 4.40),
    "gpt-4-turbo": (10.00, 30.00),
    "gpt-4": (30.00, 60.00),
    "gpt-3.5-turbo": (0.50, 1.50),
}

# Returned for a model with no pricing entry, so an unpriced model shows as
# $0 instead of being silently charged at some other model's rate.
_NO_PRICE = Price(0.0, 0.0)


def _static_price_for(model: str | None) -> Price:
    """Return the built-in fallback price for *model*.

    Uses the same match rules as the dataset lookup (exact, then longest
    prefix, then longest contained key) so ``gpt-4o-mini`` is not billed at
    the ``gpt-4o`` rate and vendor-prefixed ids like
    ``anthropic/claude-sonnet-4`` still resolve.  Cached tokens have no
    separate rate here, so the caller estimates them from the input price.
    """
    for candidate in pricing.model_candidates(model):
        key = pricing.best_match(MODEL_PRICING, candidate)
        if key:
            input_price, output_price = MODEL_PRICING[key]
            return Price(input_per_1m=input_price, output_per_1m=output_price)
    return _NO_PRICE


def _price_for(
    model: str | None,
    context_tokens: int = 0,
    table: PricingTable | None = None,
) -> Price:
    """Return the per-1M-token prices to cost *model* with.

    The Portkey dataset wins; :data:`MODEL_PRICING` covers what it does not
    list.  ``context_tokens`` (input tokens including cache reads) selects
    the band for models priced by context size, such as Gemini.  Passing
    ``table`` explicitly skips the dataset load, which is how the viewer
    costs many sessions from one table.
    """
    if table is None:
        table = pricing.load_table()
    price = table.lookup(model, context_tokens)
    if price is not None:
        return price
    return _static_price_for(model)


def _count_rollback_changes(session_dir: Path) -> dict[str, Any]:
    """Count file changes from rollback data."""
    rollback_root = session_dir / "nono-audit" / "rollbacks"
    total = created = modified = deleted = workspace = 0
    if not rollback_root.is_dir():
        return {
            "total": 0, "created": 0, "modified": 0, "deleted": 0,
            "workspace": 0, "available": False,
        }
    for rb_session in sorted(rollback_root.iterdir()):
        if not rb_session.is_dir():
            continue
        changes_dir = rb_session / "changes"
        if not changes_dir.is_dir():
            continue
        for change_file in sorted(changes_dir.glob("*.json")):
            raw_changes = read_json(change_file) or []
            for change in raw_changes:
                total += 1
                ct = change.get("change_type", "")
                if ct == "Created":
                    created += 1
                elif ct == "Modified":
                    modified += 1
                elif ct == "Deleted":
                    deleted += 1
                if change.get("path", "").startswith("/workspace/"):
                    workspace += 1
    return {
        "total": total, "created": created, "modified": modified,
        "deleted": deleted, "workspace": workspace,
        "available": total > 0,
    }


def _usage_from_chunks(chunks: list[dict[str, Any]]) -> tuple[dict[str, Any] | None, str | None]:
    """Extract ``(usage, model)`` from a streamed response's chunks.

    Anthropic splits usage across two events -- ``message_start`` carries the
    input and cache counts, ``message_delta`` the output counts -- so those
    must be *merged*.  Dispatching on the event type first matters: treating
    a ``message_delta``'s usage as a whole-usage replacement (as an
    OpenAI-shaped chunk's would be) drops the input and cache totals, which
    are the bulk of a cached agent session's cost.
    """
    usage: dict[str, Any] | None = None
    model: str | None = None
    for chunk in chunks:
        content = chunk.get("content")
        if not isinstance(content, dict):
            continue
        ctype = content.get("type")

        if ctype == "message_start":
            message = content.get("message") or {}
            if message.get("model"):
                model = message["model"]
            if isinstance(message.get("usage"), dict):
                usage = {**(usage or {}), **message["usage"]}
        elif ctype == "message_delta":
            if isinstance(content.get("usage"), dict):
                usage = {**(usage or {}), **content["usage"]}
        elif ctype == "response.completed":
            # OpenAI responses API: the final event carries complete usage.
            response = content.get("response") or {}
            if response.get("model"):
                model = response["model"]
            if isinstance(response.get("usage"), dict):
                usage = response["usage"]
        elif ctype is None:
            # OpenAI chat completions: model on every chunk, usage on the last.
            if content.get("model"):
                model = content["model"]
            if isinstance(content.get("usage"), dict):
                usage = content["usage"]
    return usage, model


def _split_calls(index: TraceIndex) -> dict[str, Any]:
    """Split the trace index into chat calls (with usage) and other HTTP calls."""
    chat_calls: list[dict[str, Any]] = []
    http_calls: list[dict[str, Any]] = []

    for rid in index.ordered_ids:
        req = index.requests[rid]
        url = req.get("url", "")
        body = req.get("body")

        if not is_chat_request(url, body):
            http_calls.append({
                "method": req.get("method", ""),
                "url": url,
                "timestamp": req.get("timestamp", ""),
                "status_code": index.responses.get(rid, {}).get("status_code"),
            })
            continue

        if rid in index.chunks:
            usage, model = _usage_from_chunks(index.chunks[rid])
        elif rid in index.responses:
            rbody = index.responses[rid].get("body") or {}
            usage, model = rbody.get("usage"), rbody.get("model")
        else:
            usage, model = None, None

        chat_calls.append({
            "model": body.get("model") or model or "unknown",
            "usage": usage,
            "timestamp": req.get("timestamp", ""),
            "url": url,
        })

    return {"chat_calls": chat_calls, "http_calls": http_calls}


def _aggregate_chat_stats(
    chat_calls: list[dict[str, Any]], table: PricingTable | None = None
) -> dict[str, Any]:
    """Aggregate token usage and compute estimated cost.

    ``table`` is the pricing data to cost against; it is loaded once by the
    caller so a viewer listing N sessions does not resolve prices N times.
    """
    tokens_in = tokens_out = total_tokens = 0
    cache_read_total = cache_write_total = 0
    total_cost = 0.0
    models: dict[str, Any] = {}

    for call in chat_calls:
        usage = call.get("usage") or {}
        model = call.get("model", "unknown")

        prompt_tokens = usage.get("prompt_tokens") or usage.get("input_tokens") or 0
        completion_tokens = usage.get("completion_tokens") or usage.get("output_tokens") or 0
        cache_read = usage.get("cache_read_input_tokens") or 0
        cache_write = usage.get("cache_creation_input_tokens") or 0

        # Cost is computed per call from *this* call's counts.  Accumulating
        # first and costing the running totals would charge every earlier
        # call's cached tokens again on each subsequent call.
        # A provider-reported cost is only trusted when it is a positive
        # number: proxies fronting Anthropic often inject ``cost: 0`` (they
        # cannot price the upstream model themselves), and taking that at
        # face value silently zeroes the whole session's estimate.
        provided_cost = usage.get("cost")
        if provided_cost is None:
            provided_cost = usage.get("estimated_cost")
        if isinstance(provided_cost, (int, float)) and provided_cost > 0:
            call_cost = float(provided_cost)
        else:
            # Gemini and friends price by context band, so the band is chosen
            # from what this call actually sent: fresh input plus the cached
            # prefix it read back.
            price = _price_for(model, prompt_tokens + cache_read, table)
            # Cached tokens are billed at their own published rate when the
            # dataset gives one, and estimated off the input rate when it
            # does not.
            cache_read_price = price.cache_read_per_1m
            if cache_read_price is None:
                cache_read_price = price.input_per_1m * CACHE_READ_MULTIPLIER
            cache_write_price = price.cache_write_per_1m
            if cache_write_price is None:
                cache_write_price = price.input_per_1m * CACHE_WRITE_MULTIPLIER
            call_cost = (
                prompt_tokens * price.input_per_1m
                + completion_tokens * price.output_per_1m
                + cache_read * cache_read_price
                + cache_write * cache_write_price
            ) / 1_000_000

        tokens_in += prompt_tokens
        tokens_out += completion_tokens
        cache_read_total += cache_read
        cache_write_total += cache_write
        total_cost += call_cost
        # Providers reporting a total are taken at their word for the
        # in/out pair, but the cache dimensions are always added on top:
        # they are Anthropic-style counts held *outside* the input figure,
        # and an OpenAI-compatible proxy's ``total_tokens`` is computed as
        # prompt + completion only, which would otherwise drop the bulk of
        # a cached agent session's tokens from the total.
        total_tokens += (
            usage.get("total_tokens") or (prompt_tokens + completion_tokens)
        ) + cache_read + cache_write

        stats = models.setdefault(
            model,
            {"calls": 0, "tokens_in": 0, "tokens_out": 0, "cost_usd": 0.0},
        )
        stats["calls"] += 1
        stats["tokens_in"] += prompt_tokens
        stats["tokens_out"] += completion_tokens
        stats["cost_usd"] = round(stats["cost_usd"] + call_cost, 6)

    turns = len(chat_calls)
    return {
        "turns": turns,
        "calls": turns,
        "tokens_in": tokens_in,
        "tokens_out": tokens_out,
        "total_tokens": total_tokens,
        "cache_read": cache_read_total,
        "cache_write": cache_write_total,
        "estimated_cost_usd": round(total_cost, 6),
        "models": models,
    }


def _count_tool_uses(content: Any) -> int:
    """Count tool-invocation blocks in a response ``content``/``output`` value."""
    if isinstance(content, dict):
        return 1 if content.get("type") in TOOLISH_TYPES else 0
    if isinstance(content, list):
        return sum(
            1
            for item in content
            if isinstance(item, dict) and item.get("type") in TOOLISH_TYPES
        )
    return 0


def _count_tool_calls(index: TraceIndex) -> int:
    """Count tool calls the model *made*, from the response side of each turn.

    Counting request bodies instead would inflate the total quadratically:
    every request replays the full conversation so far, so each earlier
    tool call would be counted again on every later turn.
    """
    count = 0

    for rid, response in index.responses.items():
        if rid in index.chunks:
            continue  # streamed; counted from the chunks below
        body = response.get("body")
        if not isinstance(body, dict):
            continue
        # Anthropic: content blocks.  OpenAI responses API: output items.
        count += _count_tool_uses(body.get("content"))
        count += _count_tool_uses(body.get("output"))
        # OpenAI chat completions: tool_calls on the assistant message.
        for choice in body.get("choices") or []:
            if isinstance(choice, dict):
                message = choice.get("message") or {}
                count += len(message.get("tool_calls") or [])

    for chunks in index.chunks.values():
        openai_indices: set[Any] = set()
        for chunk in chunks:
            content = chunk.get("content")
            if not isinstance(content, dict):
                continue
            ctype = content.get("type")
            if ctype == "content_block_start":
                # Anthropic streaming.
                count += _count_tool_uses(content.get("content_block"))
            elif ctype == "response.output_item.added":
                # OpenAI responses API streaming.
                count += _count_tool_uses(content.get("item"))
            # OpenAI chat-completions streaming: one logical call per index,
            # spread across many argument deltas.
            for choice in content.get("choices") or []:
                if not isinstance(choice, dict):
                    continue
                for call in (choice.get("delta") or {}).get("tool_calls") or []:
                    if isinstance(call, dict):
                        openai_indices.add(call.get("index", 0))
        count += len(openai_indices)

    return count


def compute_session_summary(session_dir: Path) -> dict[str, Any]:
    """Read all session artifacts and return a summary dict."""
    session_dir = Path(session_dir)
    nono_meta = load_nono_session_meta(session_dir) or {}
    audit_sess, audit_events = load_audit_session(session_dir)
    audit_sess = audit_sess or {}

    started = nono_meta.get("started") or audit_sess.get("started") or ""
    ended = audit_sess.get("ended", "")
    command = nono_meta.get("command") or audit_sess.get("command") or []
    command_str = " ".join(command) if isinstance(command, list) else str(command)

    duration_seconds: float | None = None
    if started and ended:
        with contextlib.suppress(Exception):
            duration_seconds = (
                _parse_iso(ended) - _parse_iso(started)
            ).total_seconds()

    index = load_trace_index(session_dir)
    traces = _split_calls(index)
    table = pricing.load_table()
    chat_stats = _aggregate_chat_stats(traces["chat_calls"], table)
    chat_stats["tool_calls"] = _count_tool_calls(index)

    reqlog_entries = len(load_request_log(session_dir))
    http_calls = len(traces["http_calls"])

    return {
        "session_id": session_dir.name,
        "name": nono_meta.get("name", ""),
        "command": command_str,
        "started": started,
        "ended": ended,
        "exit_code": nono_meta.get("exit_code", audit_sess.get("exit_code")),
        "status": nono_meta.get("status", ""),
        "duration_seconds": duration_seconds,
        "llm": {
            "turns": chat_stats["turns"],
            "calls": chat_stats["calls"],
            "tool_calls": chat_stats["tool_calls"],
            "tokens_in": chat_stats["tokens_in"],
            "tokens_out": chat_stats["tokens_out"],
            "total_tokens": chat_stats["total_tokens"],
            "cache_read": chat_stats["cache_read"],
            "cache_write": chat_stats["cache_write"],
            "estimated_cost_usd": chat_stats["estimated_cost_usd"],
            "models": chat_stats["models"],
            # Where the prices came from, so a cost figure in an old summary
            # can still be explained (and spotted as stale) later.
            "pricing": {"dataset": "portkey", "providers": table.sources},
        },
        "network": {
            "http_calls": http_calls,
            "request_log_entries": reqlog_entries,
            "total_requests": http_calls + reqlog_entries,
        },
        "files": _count_rollback_changes(session_dir),
        "audit": {
            "event_count": len(audit_events),
        },
    }


def write_summary(
    session_dir: Path, summary: dict[str, Any] | None = None
) -> tuple[Path, dict[str, Any]]:
    """Compute (or write the provided) summary to ``session_summary.json``.

    Returns ``(path, summary)`` so callers can use the computed dict immediately
    without re-parsing.
    """
    session_dir = Path(session_dir)
    if summary is None:
        summary = compute_session_summary(session_dir)
    path = session_dir / "session_summary.json"
    with open(path, "w", encoding="utf-8") as fh:
        json.dump(summary, fh, indent=2, default=str)
    return path, summary


def read_summary(session_dir: Path) -> dict[str, Any] | None:
    """Read a previously-written ``session_summary.json`` if it exists."""
    return read_json(Path(session_dir) / "session_summary.json")


def _parse_iso(ts: str) -> datetime:
    """Parse an ISO-8601 timestamp, tolerating sub-microsecond precision."""
    ts = ts.replace("Z", "+00:00")
    try:
        return datetime.fromisoformat(ts)
    except ValueError:
        # Some producers emit nanoseconds (2026-09-01T14:34:04.292565116+00:00);
        # datetime only accepts microseconds, so truncate the fraction.
        if "." not in ts:
            raise
        prefix, rest = ts.split(".", 1)
        digits = ""
        offset = ""
        for i, ch in enumerate(rest):
            if not ch.isdigit():
                offset = rest[i:]
                break
            digits += ch
        return datetime.fromisoformat(f"{prefix}.{digits[:6].ljust(6, '0')}{offset}")


def format_summary_for_cli(summary: dict[str, Any]) -> str:
    """Return a human-readable recap string for the CLI.

    Renders the same rich panel the tty sink prints at the end of a session,
    but flattened to plain, ANSI-free text (rich emits no escape codes for
    non-terminal streams) so it is safe for pipes, logs and golden tests.
    """
    return ui.render_summary_text(summary)
