"""Readers for the on-disk artifacts of a spens session.

A session directory (``<spens_dir>/sessions/<id>``) is written by several
independent producers -- the mitmproxy addon (``traces/``), the DNS forwarder
(``traces/dns_log.jsonl``) and nono (``nono-audit/``).  Both consumers of that
data, :mod:`spens.summarizer` (cheap aggregate stats) and :mod:`spens.viewer`
(full display parsing), need the same loading and grouping step first, so it
lives here once instead of being forked per consumer.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

# Directory (relative to the workspace root, by default) holding all spens
# session state.  Overridable per run with ``--spens-dir``.
SPENS_DIR_NAME = ".spens"

# A completion endpoint across the APIs spens captures: OpenAI chat
# completions, Anthropic messages, OpenAI responses.  Anchored with ``\b`` so
# lookalike paths (``/messages_beta``, ``/responses_archive``) do not match.
_CHAT_URL_RE = re.compile(r"/(chat/completions|messages|responses)\b")


# ---------------------------------------------------------------------------
# Raw file readers
# ---------------------------------------------------------------------------


def read_json(path: Path) -> dict[str, Any] | None:
    """Read a JSON file, returning None if it is missing or malformed."""
    try:
        with open(path, encoding="utf-8") as fh:
            return json.load(fh)
    except (OSError, json.JSONDecodeError):
        return None


def read_ndjson(path: Path) -> list[dict[str, Any]]:
    """Read a newline-delimited JSON file, skipping blank/malformed lines.

    Trace and audit files are appended to by a live process, so a truncated
    final line is normal and must not fail the whole read.
    """
    rows: list[dict[str, Any]] = []
    if not path.exists():
        return rows
    with open(path, encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if not line:
                continue
            try:
                rows.append(json.loads(line))
            except json.JSONDecodeError:
                continue
    return rows


# ---------------------------------------------------------------------------
# nono audit artifacts
# ---------------------------------------------------------------------------


def load_nono_session_meta(session_dir: Path) -> dict[str, Any] | None:
    """Load the nono session metadata (``nono-audit/sessions/<id>.json``)."""
    sessions_root = session_dir / "nono-audit" / "sessions"
    if not sessions_root.is_dir():
        return None
    for json_file in sorted(sessions_root.glob("*.json")):
        data = read_json(json_file)
        if data:
            return data
    return None


def load_audit_session(
    session_dir: Path,
) -> tuple[dict[str, Any] | None, list[dict[str, Any]]]:
    """Return ``(session.json, audit-events)`` from the nono audit directory."""
    audit_root = session_dir / "nono-audit" / "audit"
    if not audit_root.is_dir():
        return None, []
    for sub in sorted(audit_root.iterdir()):
        if not sub.is_dir():
            continue
        sess = read_json(sub / "session.json")
        events = read_ndjson(sub / "audit-events.ndjson")
        if sess is not None or events:
            return sess, events
    return None, []


def load_request_log(session_dir: Path) -> list[dict[str, Any]]:
    """Load the mitmproxy request log (method, URL, status -- no bodies)."""
    return read_ndjson(session_dir / "traces" / "request_log.jsonl")


# ---------------------------------------------------------------------------
# Trace index
# ---------------------------------------------------------------------------


def is_chat_url(url: str) -> bool:
    """Return whether ``url`` looks like an LLM completion endpoint."""
    return bool(_CHAT_URL_RE.search(url) or "chat/completions" in url)


def is_chat_request(url: str, body: Any) -> bool:
    """Return whether a captured request is an LLM completion call.

    Both the endpoint *and* the body shape must agree: a request to a chat
    URL with no ``messages``/``input`` is some other API call (e.g. a model
    listing) and belongs in the plain HTTP bucket.
    """
    return (
        is_chat_url(url)
        and isinstance(body, dict)
        and ("messages" in body or "input" in body)
    )


@dataclass
class TraceIndex:
    """Captured trace rows grouped by request id.

    ``ordered_ids`` preserves capture order (timestamp-sorted) so consumers
    can walk the session chronologically. A streamed response appears in
    ``chunks``; a buffered one in ``responses``; ``metas`` carries the
    end-of-stream latency/status row that streamed responses have instead.
    """

    requests: dict[str, dict[str, Any]] = field(default_factory=dict)
    responses: dict[str, dict[str, Any]] = field(default_factory=dict)
    chunks: dict[str, list[dict[str, Any]]] = field(default_factory=dict)
    metas: dict[str, dict[str, Any]] = field(default_factory=dict)
    ordered_ids: list[str] = field(default_factory=list)


def load_trace_index(session_dir: Path) -> TraceIndex:
    """Read every trace file for a session and group the rows by request id."""
    index = TraceIndex()
    traces_dir = session_dir / "traces"
    if not traces_dir.is_dir():
        return index

    rows: list[dict[str, Any]] = []
    for trace_file in sorted(traces_dir.glob("*.jsonl")):
        rows.extend(read_ndjson(trace_file))
    rows.sort(key=lambda r: r.get("timestamp", ""))

    for row in rows:
        rtype = row.get("type")
        if rtype == "request":
            rid = row.get("id", "")
            index.requests[rid] = row
            index.ordered_ids.append(rid)
        elif rtype == "response":
            index.responses[row.get("request_id", "")] = row
        elif rtype == "response_chunk":
            index.chunks.setdefault(row.get("request_id", ""), []).append(row)
        elif rtype == "response_meta":
            index.metas[row.get("request_id", "")] = row

    return index
