"""Helpers for shaping chat turn telemetry records."""
from __future__ import annotations

import hashlib
import time
import uuid
from datetime import datetime, timezone
from typing import Any

from scripts.lib.runtime_config import DEFAULT_CLAUDE_MODEL, DEFAULT_OFFERING_ID
from scripts.lib.chat.prompt_loader import load_prompt_bundle_id, load_prompt_manifest


def usage_field(llm, key: str) -> int:
    """Return a token-count field from llm.last_usage, or 0 when absent."""
    usage = getattr(llm, "last_usage", None)
    if not isinstance(usage, dict):
        return 0
    val = usage.get(key, 0)
    return int(val) if isinstance(val, (int, float)) else 0


def count_by(retrieved, by_doc: dict, field: str) -> dict[str, int]:
    counts: dict[str, int] = {}
    for r in retrieved:
        v = by_doc.get(r.doc_id, {}).get(field, "")
        counts[v] = counts.get(v, 0) + 1
    return counts


def count_layers(retrieved) -> dict[str, int]:
    counts: dict[str, int] = {}
    for r in retrieved:
        counts[r.layer] = counts.get(r.layer, 0) + 1
    return counts


def prompt_identity() -> tuple[str | None, dict[str, str] | None]:
    try:
        return (
            load_prompt_bundle_id(),
            {m["name"]: m["sha256"] for m in load_prompt_manifest()},
        )
    except Exception:
        return None, None


def build_failure_turn_record(
    *,
    session_id: str,
    history: list,
    user_message: str,
    failure_reason: str,
    t0: float,
    user_login: str = "local",
    mode: str = "tutor",
    prompt_identity_func=prompt_identity,
) -> dict[str, Any]:
    bundle_id, section_hashes = prompt_identity_func()
    now = datetime.now(timezone.utc)
    return {
        "ts": now.strftime("%Y-%m-%dT%H:%M:%SZ"),
        "turn_id": uuid.uuid4().hex,
        "session_id": session_id,
        "user_login": user_login,
        "mode": mode,
        "turn_index_in_session": len(history) // 2,
        "weekday": now.strftime("%a"),
        "hour_utc": now.hour,
        "user_message": user_message,
        "intent": "<error>",
        "retrieved": [],
        "retrieved_count": 0,
        "response": "",
        "elapsed_s": round(time.monotonic() - t0, 3),
        "prompt_bundle_id": bundle_id,
        "prompt_section_hashes": section_hashes,
        "failure_reason": failure_reason,
        "judge": None,
    }


def build_success_turn_record(
    *,
    session_id: str,
    user_login: str,
    mode: str = "tutor",
    history: list,
    user_message: str,
    planned,
    retrieved,
    catalog,
    pack,
    system_prompt: str,
    messages: list[dict],
    llm,
    core,
    core_doc_ids: set[str],
    session_core_enabled: bool,
    store_response: str,
    scan_result,
    elapsed: float,
    prompt_identity_func=prompt_identity,
) -> dict[str, Any]:
    by_doc = {rec["doc_id"]: rec for rec in catalog.records}
    bundle_id, section_hashes = prompt_identity_func()
    now = datetime.now(timezone.utc)
    matched_terms = list(getattr(planned, "matched_terms", []) or [])
    last_usage = getattr(llm, "last_usage", None)
    return {
        "ts": now.strftime("%Y-%m-%dT%H:%M:%SZ"),
        "turn_id": uuid.uuid4().hex,
        "session_id": session_id,
        "user_login": user_login,
        "mode": mode,
        "turn_index_in_session": len(history) // 2,
        "weekday": now.strftime("%a"),
        "hour_utc": now.hour,
        "user_message": user_message,
        "user_message_tokens": len(user_message) // 4,
        "planned_rewrite": planned.rewritten,
        "rewrite_changed": planned.rewritten != user_message,
        "intent": planned.intent,
        "is_follow_up": planned.is_follow_up,
        "matched_terms": matched_terms,
        "matched_terms_count": len(matched_terms),
        "retrieved": [
            {
                "doc_id": r.doc_id,
                "kind": by_doc.get(r.doc_id, {}).get("kind", ""),
                "topic": by_doc.get(r.doc_id, {}).get("topic", ""),
                "score": r.score,
                "layer": r.layer,
                "matched_via": r.matched_via,
                "body_chars": len(by_doc.get(r.doc_id, {}).get("concept_profile", "") or ""),
            }
            for r in retrieved
        ],
        "retrieved_count": len(retrieved),
        "retrieved_kind_distribution": count_by(retrieved, by_doc, "kind"),
        "retrieved_layer_distribution": count_layers(retrieved),
        "pack_topic_count": len(pack.topic),
        "pack_constitution_count": len(pack.course_constitution),
        "pack_module_count": len(pack.module),
        "system_prompt_tokens": len(system_prompt) // 4,
        "prompt_bundle_id": bundle_id,
        "prompt_section_hashes": section_hashes,
        "user_prompt_tokens": sum(len(m.get("content", "")) for m in messages) // 4,
        "history_tokens": sum(len(t.get("content", "")) for t in history) // 4,
        "total_input_tokens": (
            len(system_prompt)
            + (len(core.frozen_prompt_text) if core else 0)
            + sum(len(m.get("content", "")) for m in messages)
        ) // 4,
        "actual_input_tokens": usage_field(llm, "input_tokens"),
        "actual_output_tokens": usage_field(llm, "output_tokens"),
        "cache_creation_input_tokens": usage_field(llm, "cache_creation_input_tokens"),
        "cache_read_input_tokens": usage_field(llm, "cache_read_input_tokens"),
        "response": store_response,
        "response_tokens": len(store_response) // 4,
        "elapsed_s": round(elapsed, 3),
        "pii_violations": len(scan_result.pii_violations),
        "answer_leak_chunks": len(scan_result.answer_leak_chunks),
        "catalog_artifact_version_id": getattr(catalog, "artifact_version_id", ""),
        "embedding_model": catalog.embedding_model,
        "chat_model": getattr(llm, "model", DEFAULT_CLAUDE_MODEL),
        "offering_id": DEFAULT_OFFERING_ID,
        "session_core_enabled": session_core_enabled,
        "session_core_doc_ids": (sorted(core_doc_ids) if core else None),
        "incremental_doc_ids": [r.doc_id for r in retrieved],
        "cache_ttl_fallback": bool(
            isinstance(last_usage, dict) and last_usage.get("cache_ttl_fallback")
        ),
        "session_core_sha256": (
            hashlib.sha256(core.frozen_prompt_text.encode("utf-8")).hexdigest()[:16]
            if core else None
        ),
        "system_prompt_sha256": hashlib.sha256(
            system_prompt.encode("utf-8")
        ).hexdigest()[:16],
        "judge": None,
    }