"""Pure-function analytics over the turn log.

Each function takes a path to the JSONL log + optional date filters and
returns a structured aggregate. No I/O beyond reading the log (and, for
lecture-session aggregation, the doc catalog).
"""
from __future__ import annotations
import json
import re
from collections import Counter, defaultdict
from pathlib import Path
from typing import Any

from scripts.lib.chat.pricing import (
    PRICE_SONNET_CACHE_READ,
    PRICE_SONNET_CACHE_WRITE,
    PRICE_SONNET_CACHE_WRITE_1H,
    PRICE_SONNET_IN,
    PRICE_SONNET_OUT,
)
from scripts.lib.chat_eval.turn_log import iter_turns, iter_turns_compat


_LECTURE_DATE_RE = re.compile(r"lecture_problem_(\d{4}-\d{2}-\d{2})_")


def _build_doc_to_lecture_dates(catalog_path: Path) -> dict[str, set[str]]:
    """Return {doc_id: {YYYY-MM-DD, ...}} mapping each catalog record to
    the lecture session date(s) it's associated with.

    Only `lecture_problem` docs map. Earlier versions also chained
    `syntopicon_entry` records via `references[].lecture_problem_doc_id`,
    but that transitive mapping fanned a single demand-question turn out
    to many game-theory / linear-pricing dates whose syntopicon entries
    happened to reference back to the same primitive — the section
    became noise (same example query repeated across rows, topics
    unrelated to the listed lecture). Direct-only keeps it honest: a
    date appears only when the bot actually retrieved that lecture's
    own problem set.
    """
    doc_to_dates: dict[str, set[str]] = {}
    if not Path(catalog_path).is_file():
        return doc_to_dates
    for line in Path(catalog_path).read_text(encoding="utf-8").splitlines():
        if not line.strip():
            continue
        rec = json.loads(line)
        if rec.get("kind") != "lecture_problem":
            continue
        doc_id = rec.get("doc_id", "")
        m = _LECTURE_DATE_RE.match(doc_id)
        if m:
            doc_to_dates[doc_id] = {m.group(1)}
    return doc_to_dates


def chat_sessions(
    path: Path,
    *,
    since: str | None = None,
) -> list[dict[str, Any]]:
    """Return one row per chat conversation, in start-time order.

    A "chat session" is a contiguous run of turns sharing the same
    session_id where turn_index_in_session increments monotonically
    from 0. When the index resets to 0 (Gradio session reopened, or
    server restarted), a new chat begins.

    Each row carries:
      - session_id: the session identifier (often "local-default")
      - started_at: ISO timestamp of the opening turn
      - opener: the first user_message of the chat (the framing question)
      - turn_count: total turns in this chat
      - topics: top syntopicon topics retrieved across all turns
      - intents: count by intent across the chat
      - last_user_message: the final student turn (often shows where
        the conversation ended)
    """
    sessions: list[dict[str, Any]] = []
    current: dict[str, Any] | None = None
    for turn in iter_turns(path, since=since):
        sid = turn.get("session_id", "")
        idx = turn.get("turn_index_in_session", 0)
        # New chat when session_id changes OR turn_index resets to 0
        if (current is None
                or current["session_id"] != sid
                or (idx == 0 and current["turn_count"] > 0)):
            current = {
                "session_id": sid,
                "started_at": turn.get("ts", ""),
                "opener": turn.get("user_message", ""),
                "turn_count": 0,
                "topics": Counter(),
                "intents": Counter(),
                "last_user_message": "",
            }
            sessions.append(current)
        current["turn_count"] += 1
        intent = turn.get("intent", "?")
        current["intents"][intent] += 1
        current["last_user_message"] = turn.get("user_message", "")
        for r in turn.get("retrieved", []):
            if r.get("kind") == "syntopicon_entry":
                topic = r.get("topic", "")
                if topic:
                    current["topics"][topic] += 1
    # Convert Counters → dicts for clean serialization
    for s in sessions:
        s["topics"] = dict(s["topics"])
        s["intents"] = dict(s["intents"])
    return sessions


def topics_by_session(
    path: Path,
    *,
    catalog_path: Path,
    since: str | None = None,
) -> dict[str, dict[str, Any]]:
    """Aggregate turns by the lecture session(s) their retrieved docs trace to.

    A turn is associated with every lecture date its retrieved docs link to
    (either directly via lecture_problem doc_id or transitively via
    syntopicon_entry references). Turns whose retrieval doesn't link to any
    dated lecture (e.g., only generic slides) are excluded — flagged in the
    `unmapped_turns` count.

    Returns: {
      "sessions": {date: {turn_count, intents, topics, sample_query}},
      "unmapped_turns": int,
    }
    """
    doc_to_dates = _build_doc_to_lecture_dates(catalog_path)

    by_date: dict[str, dict[str, Any]] = defaultdict(
        lambda: {
            "turn_count": 0,
            "intents": Counter(),
            "topics": Counter(),
            "sample_query": "",
            "sample_ts": "",
        }
    )
    unmapped = 0

    for turn in iter_turns(path, since=since):
        turn_dates: set[str] = set()
        topics_this_turn: set[str] = set()
        for r in turn.get("retrieved", []):
            doc_id = r.get("doc_id", "")
            for d in doc_to_dates.get(doc_id, ()):
                turn_dates.add(d)
            if r.get("kind") == "syntopicon_entry":
                topic = r.get("topic", "")
                if topic:
                    topics_this_turn.add(topic)

        if not turn_dates:
            unmapped += 1
            continue

        intent = turn.get("intent", "?")
        sample = turn.get("user_message", "")[:80]
        ts = turn.get("ts", "")
        for d in turn_dates:
            entry = by_date[d]
            entry["turn_count"] += 1
            entry["intents"][intent] += 1
            for t in topics_this_turn:
                entry["topics"][t] += 1
            # Keep the most-recent representative query so the example
            # column actually reflects current student activity on this
            # date's material, not whichever turn happened to land first.
            if ts >= entry["sample_ts"]:
                entry["sample_query"] = sample
                entry["sample_ts"] = ts

    # Convert Counters to dicts for clean serialization
    sessions = {
        d: {
            "turn_count": v["turn_count"],
            "intents": dict(v["intents"]),
            "topics": dict(v["topics"]),
            "sample_query": v["sample_query"],
        }
        for d, v in by_date.items()
    }
    return {"sessions": sessions, "unmapped_turns": unmapped}


def coverage(path: Path, *, since: str | None = None) -> dict[str, int]:
    """Count syntopicon-card retrievals across turns.

    Only counts `matched_via` values where the retrieved doc is itself
    a syntopicon_entry. Other matches (embedding similarity, slide
    matches, escort_for hooks) are excluded so faculty see the
    syntopicon-vocabulary coverage signal cleanly.
    """
    counts: Counter[str] = Counter()
    for turn in iter_turns(path, since=since):
        for r in turn.get("retrieved", []):
            if r.get("kind") != "syntopicon_entry":
                continue
            mv = r.get("matched_via", "")
            if not mv or mv.startswith("escort_for:") or mv == "embedding":
                continue
            counts[mv] += 1
    return dict(counts)


def doc_utilization(path_or_turns, *, since: str | None = None) -> dict[str, dict[str, Any]]:
    """For each doc_id ever retrieved, return retrieved_count, used_count, util_pct.

    A doc is 'used' if it appears in any judge.docs_actually_used entry.
    Accepts either a Path (reads turns.jsonl) or a pre-loaded list of turn
    dicts (avoids re-reading the file when a caller already has it).
    """
    retrieved: Counter[str] = Counter()
    used: Counter[str] = Counter()
    for turn in iter_turns_compat(path_or_turns, since=since):
        seen_in_turn = set()
        for r in turn.get("retrieved", []):
            doc_id = r["doc_id"]
            if doc_id not in seen_in_turn:
                retrieved[doc_id] += 1
                seen_in_turn.add(doc_id)
        judge = turn.get("judge")
        if judge:
            for d in judge.get("docs_actually_used", []):
                used[d.get("doc_id", "")] += 1

    out: dict[str, dict[str, Any]] = {}
    for doc_id, rc in retrieved.items():
        uc = used.get(doc_id, 0)
        out[doc_id] = {
            "retrieved_count": rc,
            "used_count": uc,
            "util_pct": round(100.0 * uc / rc, 1) if rc else 0.0,
        }
    return out


def misconception_frequency(path: Path, *, since: str | None = None) -> dict[str, dict[str, int]]:
    """For each misconception id, return four counts:

    - count: total flag occurrences
    - in_query: subset where the student's query exhibited the misconception
    - handled: subset where the response addressed it
    - missed: in_query AND NOT handled — the most actionable failure mode
    """
    counts: dict[str, dict[str, int]] = defaultdict(
        lambda: {"count": 0, "in_query": 0, "handled": 0, "missed": 0}
    )
    for turn in iter_turns(path, since=since):
        judge = turn.get("judge")
        if not judge:
            continue
        for flag in judge.get("misconception_flags", []):
            mid = flag.get("id", "")
            if not mid:
                continue
            in_q = bool(flag.get("in_query"))
            handled = bool(flag.get("handled_in_response"))
            counts[mid]["count"] += 1
            if in_q:
                counts[mid]["in_query"] += 1
            if handled:
                counts[mid]["handled"] += 1
            if in_q and not handled:
                counts[mid]["missed"] += 1
    return dict(counts)


def voice_failures(path: Path, *, since: str | None = None,
                   threshold: int = 2) -> list[dict]:
    """Return judged turns where voice score <= threshold, sorted oldest first."""
    out = []
    for turn in iter_turns(path, since=since):
        judge = turn.get("judge")
        if not judge:
            continue
        if judge.get("scores", {}).get("voice", 5) <= threshold:
            out.append({
                "turn_id": turn["turn_id"],
                "ts": turn.get("ts", ""),
                "user_message": turn.get("user_message", ""),
                "response": turn.get("response", ""),
                "voice_score": judge["scores"]["voice"],
                "diagnosis": judge.get("diagnosis", ""),
            })
    return out


def timing_density(path: Path, *, since: str | None = None) -> dict[str, dict[str, int]]:
    """Histogram of turn count by weekday and hour."""
    grid: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
    for turn in iter_turns(path, since=since):
        wd = turn.get("weekday", "?")
        hour = str(turn.get("hour_utc", "?"))
        grid[wd][hour] += 1
        grid[wd]["total"] = grid[wd].get("total", 0) + 1
    return {k: dict(v) for k, v in grid.items()}


def query_frequency(path: Path, *, since: str | None = None,
                    top_n: int = 20) -> list[dict[str, Any]]:
    """Top-N raw student queries by frequency.

    Uses `user_message` (what the student typed), not `planned_rewrite`
    (the Haiku-rewritten version). Faculty want to see what students
    actually asked.
    """
    counts: Counter[str] = Counter()
    for turn in iter_turns(path, since=since):
        q = turn.get("user_message", "")
        if q:
            counts[q] += 1
    return [{"query": q, "count": c} for q, c in counts.most_common(top_n)]


def score_summary(path: Path, *, since: str | None = None) -> dict[str, float]:
    """Mean score per rubric dimension across judged turns."""
    sums: dict[str, float] = defaultdict(float)
    counts: dict[str, int] = defaultdict(int)
    for turn in iter_turns(path, since=since):
        judge = turn.get("judge")
        if not judge:
            continue
        for k, v in judge.get("scores", {}).items():
            if isinstance(v, (int, float)):
                sums[k] += v
                counts[k] += 1
    return {k: round(sums[k] / counts[k], 2) for k in sums if counts[k] > 0}


def score_summary_with_stats(
    path_or_turns,
    *,
    since: str | None = None,
    until: str | None = None,
) -> dict[str, dict[str, float]]:
    """Per-dimension {mean, std, n} over judged turns in [since, until).

    Both bounds are ISO-8601 strings; left-inclusive, right-exclusive.
    `since=None` means no lower bound; `until=None` means no upper bound.
    Sample standard deviation (Bessel correction); std is 0.0 when n < 2.

    Accepts either a Path (reads turns.jsonl) or a pre-loaded list of turn
    dicts (avoids re-reading the file when a caller already has it).
    """
    sums: dict[str, float] = defaultdict(float)
    sumsq: dict[str, float] = defaultdict(float)
    counts: dict[str, int] = defaultdict(int)
    for turn in iter_turns_compat(path_or_turns, since=since):
        if until is not None and turn.get("ts", "") >= until:
            continue
        judge = turn.get("judge")
        if not judge:
            continue
        for k, v in judge.get("scores", {}).items():
            if isinstance(v, (int, float)):
                sums[k] += v
                sumsq[k] += v * v
                counts[k] += 1
    out: dict[str, dict[str, float]] = {}
    for k, n in counts.items():
        mean = sums[k] / n
        if n >= 2:
            var = (sumsq[k] - n * mean * mean) / (n - 1)
            std = var ** 0.5 if var > 0 else 0.0
        else:
            std = 0.0
        out[k] = {"mean": round(mean, 4), "std": round(std, 4), "n": n}
    return out


# Backward-compatible private names for tests and older imports.
_PRICE_SONNET_IN = PRICE_SONNET_IN
_PRICE_SONNET_OUT = PRICE_SONNET_OUT
_PRICE_SONNET_CACHE_WRITE = PRICE_SONNET_CACHE_WRITE
_PRICE_SONNET_CACHE_READ = PRICE_SONNET_CACHE_READ
_PRICE_SONNET_CACHE_WRITE_1H = PRICE_SONNET_CACHE_WRITE_1H


def _turn_cost(turn: dict) -> float:
    """Compute one turn's chat-tutor API cost.

    Prefers actual token counts from the Anthropic response (input_tokens,
    output_tokens, cache_creation_input_tokens, cache_read_input_tokens)
    when present; falls back to len()//4 estimates from turn_record's
    total_input_tokens / response_tokens fields.
    """
    actual_in = turn.get("actual_input_tokens", 0) or 0
    actual_out = turn.get("actual_output_tokens", 0) or 0
    cache_write = turn.get("cache_creation_input_tokens", 0) or 0
    cache_read = turn.get("cache_read_input_tokens", 0) or 0

    if actual_in or actual_out or cache_write or cache_read:
        return (
            actual_in * _PRICE_SONNET_IN
            + actual_out * _PRICE_SONNET_OUT
            + cache_write * _PRICE_SONNET_CACHE_WRITE
            + cache_read * _PRICE_SONNET_CACHE_READ
        )
    # Fallback to estimate.
    est_in = turn.get("total_input_tokens", 0) or 0
    est_out = turn.get("response_tokens", 0) or 0
    return est_in * _PRICE_SONNET_IN + est_out * _PRICE_SONNET_OUT


def cost_summary(path: Path, *, since: str | None = None) -> dict:
    """Estimate total chat-tutor API spend across logged turns.

    Returns a dict with:
      - total_cost_usd: sum across all turns
      - mean_cost_usd: per-turn average
      - max_cost_usd: most expensive turn
      - n_turns: count
      - max_turn_id: id of most expensive turn (for drill-down)
      - by_intent: {intent: {n, total, mean}} breakdown

    Prefers actual token counts from the Anthropic response when available
    (including cache write/read pricing); falls back to len(text)//4
    estimates. Costs the chat tutor only — judge cost is computed separately
    if needed.
    """
    items = []
    by_intent_totals: dict[str, list[float]] = defaultdict(list)
    for turn in iter_turns(path, since=since):
        cost = _turn_cost(turn)
        items.append((turn.get("turn_id", ""), cost))
        by_intent_totals[turn.get("intent", "?")].append(cost)
    if not items:
        return {
            "total_cost_usd": 0.0, "mean_cost_usd": 0.0, "max_cost_usd": 0.0,
            "n_turns": 0, "max_turn_id": "", "by_intent": {},
        }
    total = sum(c for _, c in items)
    n = len(items)
    max_id, max_cost = max(items, key=lambda x: x[1])
    by_intent = {
        intent: {
            "n": len(costs),
            "total": round(sum(costs), 4),
            "mean": round(sum(costs)/len(costs), 4),
        }
        for intent, costs in by_intent_totals.items()
    }
    return {
        "total_cost_usd": round(total, 4),
        "mean_cost_usd": round(total / n, 4),
        "max_cost_usd": round(max_cost, 4),
        "n_turns": n,
        "max_turn_id": max_id,
        "by_intent": by_intent,
    }


def cache_summary(path: Path, *, since: str | None = None) -> dict:
    """Aggregate Anthropic cache usage across logged turns.

    Returns:
      - total_uncached_input: tokens billed at full input rate
      - total_cache_writes: tokens billed at 1.25x (one-time cache creation)
      - total_cache_reads: tokens billed at 0.10x (cache hit)
      - cache_hit_rate: cache_reads / (cache_reads + uncached_input + cache_writes)
      - estimated_savings_usd: vs no-cache counterfactual
        ( (cache_writes + cache_reads) * IN_RATE )
        minus
        ( cache_writes * WRITE_RATE + cache_reads * READ_RATE )
      - n_turns_with_data: count of turns with non-zero actual_* fields
    """
    total_uncached = 0
    total_writes = 0
    total_reads = 0
    n_with_data = 0
    for turn in iter_turns(path, since=since):
        actual_in = turn.get("actual_input_tokens", 0) or 0
        cw = turn.get("cache_creation_input_tokens", 0) or 0
        cr = turn.get("cache_read_input_tokens", 0) or 0
        if actual_in or cw or cr:
            n_with_data += 1
            total_uncached += actual_in
            total_writes += cw
            total_reads += cr

    total_cacheable = total_uncached + total_writes + total_reads
    hit_rate = (total_reads / total_cacheable) if total_cacheable else 0.0

    # Counterfactual: if no caching, all of (writes + reads) would have been
    # billed at the full input rate. With caching: writes at 1.25x, reads at 0.10x.
    no_cache_cost = (total_writes + total_reads) * _PRICE_SONNET_IN
    with_cache_cost = (total_writes * _PRICE_SONNET_CACHE_WRITE
                       + total_reads * _PRICE_SONNET_CACHE_READ)
    savings = round(no_cache_cost - with_cache_cost, 4)

    return {
        "total_uncached_input": total_uncached,
        "total_cache_writes": total_writes,
        "total_cache_reads": total_reads,
        "cache_hit_rate": round(hit_rate, 3),
        "estimated_savings_usd": savings,
        "n_turns_with_data": n_with_data,
    }


def cache_recommendation(
    path: Path,
    *,
    since: str | None = None,
    ttl: str = "1h",
    top_n: int = 15,
) -> dict:
    """Recommend which retrieved documents to permanently cache.

    For each doc that's been retrieved at least twice (caching pays only when
    the doc is re-fetched), compute the projected savings of pinning it into a
    cached prefix. Rank by savings descending. Returns the top-N recommended
    docs plus an aggregate "if you cache all of these, total $/turn savings".

    The economics:

    Without caching, N retrievals of a doc with B body tokens cost:
        N * B * IN_RATE
    With a write-once + read-many cache:
        B * WRITE_RATE  +  (N - 1) * B * READ_RATE
    Savings:
        B * [N * IN_RATE - WRITE_RATE - (N - 1) * READ_RATE]
      = B * [(N - 1) * (IN_RATE - READ_RATE) - (WRITE_RATE - IN_RATE)]

    Break-even N depends on TTL:
      - ttl='5min' (ephemeral): write at 1.25x input → break-even at N≈1.3
      - ttl='1h'   (extended):  write at 2.00x input → break-even at N≈2.1

    Args:
      ttl: '5min' or '1h'. '1h' is the right setting for "core curriculum"
        docs you want pinned across student sessions during a class window.
      top_n: how many recommendations to return.

    Returns:
      {
        "ttl": str,
        "n_docs_with_data": int,
        "recommended": [
          {"doc_id", "kind", "topic", "retrieval_count",
           "body_tokens", "savings_per_term_usd", "break_even_at"},
          ...   # ranked by savings desc, max top_n
        ],
        "total_recommended_savings_usd": float,
        "no_data_warning": str | None,  # if no body_chars in log
      }
    """
    write_rate = (_PRICE_SONNET_CACHE_WRITE_1H if ttl == "1h"
                  else _PRICE_SONNET_CACHE_WRITE)
    read_rate = _PRICE_SONNET_CACHE_READ
    in_rate = _PRICE_SONNET_IN

    # Per-doc accumulators: count + median body_chars (we keep first observed
    # since concept_profile is stable across retrievals of the same doc_id).
    by_doc: dict[str, dict[str, Any]] = {}
    for turn in iter_turns(path, since=since):
        for r in turn.get("retrieved", []):
            doc_id = r.get("doc_id", "")
            if not doc_id:
                continue
            entry = by_doc.setdefault(doc_id, {
                "retrieval_count": 0, "body_chars": 0,
                "kind": r.get("kind", ""), "topic": r.get("topic", ""),
            })
            entry["retrieval_count"] += 1
            # Capture body_chars from the first retrieval that has it.
            if not entry["body_chars"]:
                entry["body_chars"] = r.get("body_chars", 0) or 0

    # Filter: must have body_chars > 0 (instrumented by orchestrator) AND
    # retrieved 2+ times (1 retrieval can never benefit from caching).
    candidates = []
    docs_missing_body = 0
    for doc_id, info in by_doc.items():
        n = info["retrieval_count"]
        body_chars = info["body_chars"]
        if n < 2:
            continue
        if body_chars <= 0:
            docs_missing_body += 1
            continue
        body_tokens = body_chars // 4
        # savings = body_tokens × [(N-1) × (IN-READ) − (WRITE-IN)]
        savings = body_tokens * (
            (n - 1) * (in_rate - read_rate) - (write_rate - in_rate)
        )
        if savings <= 0:
            continue
        # Break-even N: the smallest integer N where caching wins.
        break_even = max(2, int((write_rate - in_rate) / (in_rate - read_rate)) + 2)
        candidates.append({
            "doc_id": doc_id,
            "kind": info["kind"],
            "topic": info["topic"],
            "retrieval_count": n,
            "body_tokens": body_tokens,
            "savings_per_term_usd": round(savings, 5),
            "break_even_at": break_even,
        })

    candidates.sort(key=lambda c: -c["savings_per_term_usd"])
    recommended = candidates[:top_n]
    total_savings = round(sum(c["savings_per_term_usd"] for c in recommended), 4)

    no_data_warning = None
    if docs_missing_body > 0 and not candidates:
        no_data_warning = (
            f"{docs_missing_body} retrieved docs lack body_chars data — these "
            "logs predate the cache-recommendation instrumentation. New chat "
            "sessions will populate this analysis."
        )

    return {
        "ttl": ttl,
        "n_docs_with_data": len([d for d in by_doc.values() if d["body_chars"] > 0]),
        "recommended": recommended,
        "total_recommended_savings_usd": total_savings,
        "no_data_warning": no_data_warning,
    }