formscout/session.py

"""
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)

    laban, flexion, chart_paths = _run_analysis(session, pose2d, test_name, side, idx)

    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,
        laban=laban, flexion=flexion, chart_paths=chart_paths,
    )
    session.entries.append(entry)
    _persist(session)
    return entry


def _run_analysis(session, pose2d, test_name, side, governing_idx):
    """Compute Laban Effort, relevant-joint flexion, and per-test charts.

    Returns (laban, flexion, chart_paths); any failure degrades to (None, None, {})
    without blocking the entry."""
    try:
        from formscout.analysis import charts as C
        from formscout.analysis.laban import compute_laban
        from formscout.analysis.relevant_joints import primary_angle, relevant_joints
        from formscout.analysis.timeseries import angle_series, relevant_flexion_at
    except Exception as e:
        logger.warning("analysis engine unavailable: %s", e)
        return None, None, {}

    try:
        laban = compute_laban(pose2d, test_name, pose2d.fps)
        gidx = governing_idx if governing_idx is not None else 0
        flexion = relevant_flexion_at(pose2d, test_name, gidx)
        series = angle_series(pose2d, test_name)

        cdir = os.path.join(session.session_dir, "charts")
        os.makedirs(cdir, exist_ok=True)
        tag = f"{test_name}_{side}"
        produced = {
            "angle": C.angle_over_time(series, primary_angle(test_name), governing_idx,
                                       os.path.join(cdir, f"{tag}_angle.png"),
                                       title=f"{test_name.replace('_', ' ').title()} — angle over time"),
            "velocity": C.velocity_profile(pose2d.keypoints, pose2d.fps, relevant_joints(test_name),
                                           os.path.join(cdir, f"{tag}_vel.png")),
            "radar": C.laban_radar(laban["effort"], os.path.join(cdir, f"{tag}_radar.png")),
            "flexion": C.flexion_bars(flexion, os.path.join(cdir, f"{tag}_flex.png")),
        }
        chart_paths = {k: p for k, p in produced.items() if p}
        return laban, flexion, chart_paths
    except Exception as e:
        logger.warning("analysis/charts failed: %s", e)
        return None, None, {}


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)

    scoresheet_path = None
    try:
        from formscout.agents.scoresheet import ScoreSheetAgent
        scoresheet_path = ScoreSheetAgent().run(report, session.session_dir)
    except Exception as e:
        logger.warning("scoresheet generation failed: %s", e)

    report = replace(report, pdf_path=pdf_path, scoresheet_path=scoresheet_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,
        "laban": e.laban, "flexion": _jsonable_flexion(e.flexion), "chart_paths": e.chart_paths,
    }


def _jsonable_flexion(flexion: dict | None) -> dict:
    if not flexion:
        return {}
    return {k: {"deg": round(v["deg"], 1), "openness": v["openness"]} for k, v in flexion.items()}


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(f"*view: {e.view} · confidence: {e.confidence:.0%}*")
        lines.append("")
        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}")

        # Relevant-joint flexion (degrees + open/closed)
        if e.flexion:
            lines.append("\n### Relevant joint flexion (key frame)")
            lines.append("| Joint angle | Degrees | State |")
            lines.append("|---|---|---|")
            for name, v in e.flexion.items():
                lines.append(f"| {name.replace('_', ' ')} | {v['deg']:.1f}° | {v['openness']} |")

        # Laban Effort
        if e.laban:
            eff, lab = e.laban.get("effort", {}), e.laban.get("labels", {})
            lines.append("\n### Laban Effort (kinematic estimate)")
            lines.append("| Factor | Value | Quality |")
            lines.append("|---|---|---|")
            for k in ("space", "weight", "time", "flow"):
                lines.append(f"| {k.title()} | {eff.get(k, 0):.2f} | {lab.get(k, '')} |")
            if e.laban.get("body_emphasis"):
                emph = ", ".join(f"{n} ({s})" for n, s in e.laban["body_emphasis"])
                lines.append(f"\n- **Body emphasis:** {emph}")
            if e.laban.get("leading_joint"):
                lines.append(f"- **Leading joint:** {e.laban['leading_joint']}")
            lines.append(f"- *{e.laban.get('notes', '')}*")

        # Full measurement dump
        if e.measurements:
            lines.append("\n### All measurements")
            lines.append("| Measurement | Value |")
            lines.append("|---|---|")
            for k, v in e.measurements.items():
                val = f"{v:.2f}" if isinstance(v, float) else str(v)
                lines.append(f"| {k.replace('_', ' ')} | {val} |")
        if e.features and e.features.symmetry_delta is not None:
            lines.append(f"\n- **L/R symmetry delta:** {e.features.symmetry_delta:.1f}")

        # Artifacts
        if e.keyframe_path:
            lines.append(f"\n- Key frame: `{e.keyframe_path}`")
        for kind, path in (e.chart_paths or {}).items():
            lines.append(f"- Chart ({kind}): `{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)