"""
Screening-session accumulator.

Accumulates one SessionEntry per analyzed clip, persists each to a temp session
dir (session.json + analysis.md + key-frame PNGs), and on finish builds a
ReportResult (via ReportAgent) + a PDF (via PdfReportAgent).

Pure orchestration — no Gradio imports. Disk writes tolerate failure with a
logged warning and never block scoring.
"""
from __future__ import annotations

import json
import logging
import os
import tempfile
import uuid
from dataclasses import dataclass, replace

from formscout.rubric import score_test
from formscout.types import MovementResult, ReportResult, SessionEntry

logger = logging.getLogger(__name__)

# Maps each test to the BiomechFeatures.timing key holding its governing frame.
TIMING_KEY = {
    "deep_squat": "deepest_frame",
    "hurdle_step": "peak_step_frame",
    "inline_lunge": "deepest_lunge_frame",
    "shoulder_mobility": "measure_frame",
    "active_slr": "peak_raise_frame",
    "trunk_stability_pushup": "max_sag_frame",
    "rotary_stability": "peak_extension_frame",
}


@dataclass
class Session:
    """Mutable session: an id, its temp dir, and accumulated entries."""
    session_id: str
    session_dir: str
    entries: list  # list[SessionEntry]


def new_session() -> Session:
    sid = uuid.uuid4().hex[:12]
    base = os.path.join(tempfile.gettempdir(), "formscout_sessions", sid)
    try:
        os.makedirs(os.path.join(base, "keyframes"), exist_ok=True)
    except Exception as e:
        logger.warning("session dir create failed: %s", e)
    return Session(session_id=sid, session_dir=base, entries=[])


def governing_frame_index(features) -> int | None:
    """Return the governing frame index for this test, or None."""
    key = TIMING_KEY.get(features.test_name)
    if key is None:
        return None
    idx = features.timing.get(key)
    return int(idx) if isinstance(idx, (int, float)) else None


def worst_compensation_caption(judge, features) -> str:
    """Short caption naming the worst compensation for the key-frame still."""
    if judge and getattr(judge, "compensation_tags", None):
        return ", ".join(judge.compensation_tags)
    failed = [k.replace("_", " ") for k, v in features.alignments.items() if v is False]
    return ("compensation: " + ", ".join(failed)) if failed else "key position"


def add_analysis(session, *, ingest, pose2d, features, judge, test_name, side,
                 draw_trails: bool = False) -> SessionEntry:
    """Build a SessionEntry from a completed analysis, render its key-frame,
    persist the session, append, and return the entry."""
    movement = MovementResult(test_name=test_name, side=side, confidence=1.0)
    rubric = score_test(features)

    needs_human = bool((judge and judge.needs_human) or rubric.needs_human)
    if needs_human:
        score = None
    elif judge and judge.score is not None:
        score = judge.score
    else:
        score = rubric.score

    keyframe_path = None
    idx = governing_frame_index(features)
    if idx is not None and 0 <= idx < len(pose2d.keypoints):
        from formscout.agents.visualizer import PoseVisualizer
        caption = (f"{test_name.replace('_', ' ').title()} "
                   f"({side}) — {worst_compensation_caption(judge, features)}")
        layers = {"skeleton", "trails"} if draw_trails else {"skeleton"}
        out_png = os.path.join(session.session_dir, "keyframes", f"{test_name}_{side}.png")
        try:
            keyframe_path = PoseVisualizer().render_frame(ingest, pose2d, idx, layers, caption, out_png)
        except Exception as e:
            logger.warning("keyframe render failed: %s", e)

    measurements = {}
    measurements.update(features.angles)
    measurements.update(features.alignments)

    entry = SessionEntry(
        test_name=test_name, side=side, score=score, needs_human=needs_human,
        rationale=(judge.rationale if judge else rubric.rationale),
        compensation_tags=list(judge.compensation_tags) if judge else [],
        corrective_hint=(judge.corrective_hint if judge else ""),
        measurements=measurements,
        confidence=(judge.confidence if judge else rubric.confidence),
        view=features.view,
        keyframe_path=keyframe_path,
        movement=movement, features=features, rubric_score=rubric, judge=judge,
    )
    session.entries.append(entry)
    _persist(session)
    return entry


def finish_session(session) -> tuple[ReportResult | None, str | None]:
    """Build the composite report + PDF. Returns (report, pdf_path).
    Returns (None, None) for an empty session."""
    if not session.entries:
        return None, None

    from formscout.agents.report import ReportAgent
    report_inputs = [{
        "movement": e.movement, "features": e.features,
        "rubric_score": e.rubric_score, "judge": e.judge, "side": e.side,
    } for e in session.entries]
    report = ReportAgent().run(report_inputs)

    pdf_path = None
    try:
        from formscout.agents.pdf_report import PdfReportAgent
        pdf_path = PdfReportAgent().run(report, session.entries, session.session_dir)
    except Exception as e:
        logger.warning("pdf generation failed: %s", e)

    report = replace(report, pdf_path=pdf_path)
    return report, pdf_path


# ── Persistence ───────────────────────────────────────────────────────────────

def _jsonable(d: dict) -> dict:
    out = {}
    for k, v in d.items():
        if isinstance(v, float):
            out[k] = round(v, 2)
        elif isinstance(v, (int, str, bool)) or v is None:
            out[k] = v
        else:
            out[k] = str(v)
    return out


def _entry_display(e: SessionEntry) -> dict:
    return {
        "test_name": e.test_name, "side": e.side, "score": e.score,
        "needs_human": e.needs_human, "rationale": e.rationale,
        "compensation_tags": list(e.compensation_tags), "corrective_hint": e.corrective_hint,
        "measurements": _jsonable(e.measurements), "confidence": round(e.confidence, 2),
        "view": e.view, "keyframe_path": e.keyframe_path,
    }


def _render_markdown(session: Session) -> str:
    lines = ["# FormScout — Session Log", ""]
    for e in session.entries:
        title = e.test_name.replace("_", " ").title()
        if e.side in ("left", "right"):
            title += f" ({e.side})"
        score = "Clinician review required" if e.needs_human else f"{e.score}/3"
        lines.append(f"## {title} — {score}")
        lines.append(e.rationale or "")
        if e.compensation_tags:
            lines.append(f"- Compensations: {', '.join(e.compensation_tags)}")
        if e.corrective_hint:
            lines.append(f"- Corrective: {e.corrective_hint}")
        if e.keyframe_path:
            lines.append(f"- Key frame: `{e.keyframe_path}`")
        lines.append("")
    return "\n".join(lines)


def _persist(session: Session) -> None:
    try:
        with open(os.path.join(session.session_dir, "session.json"), "w") as f:
            json.dump([_entry_display(e) for e in session.entries], f, indent=2)
        with open(os.path.join(session.session_dir, "analysis.md"), "w") as f:
            f.write(_render_markdown(session))
    except Exception as e:
        logger.warning("session persist failed: %s", e)