"""ECG-FM numerical verifier (Phase 4).

The interpretation LLM produces measurements (heart rate, intervals, axis,
ST deviation) along with its qualitative read. Multimodal LLMs are
demonstrably weak at reading numbers off rendered EKG strips — they
confabulate plausible-looking values. This module runs a numerical model
on the underlying WFDB signal to produce ground-truth measurements, then
the caller reconciles: if the LLM and verifier disagree by more than a
configurable threshold, the verifier wins and the corrected number is
flagged on the case.

This module is intentionally lazy — the model loads on first call so the
service can boot even on a host that hasn't fetched the weights yet. If
torch or the model files are missing, every call returns measurements_skipped
with a clear reason rather than failing hard.

Setup (one-time, on the droplet):

    cd apps/ekg-tutor/interpret
    .venv/bin/pip install torch torchvision wfdb
    .venv/bin/python -c "from huggingface_hub import snapshot_download; \
        snapshot_download('bowang-lab/ECG-FM', local_dir='./models/ecg-fm')"

Then set EKG_TUTOR_ECG_FM_PATH=./models/ecg-fm in /root/secrets.env.
"""

from __future__ import annotations

import logging
import os
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class MeasurementVerification:
    field: str           # "rate" | "intervals.pr" | "intervals.qrs" | "intervals.qtc" | "axis"
    llm_value: float
    verifier_value: float
    unit: str
    delta: float
    threshold: float
    agreement: bool
    correction_applied: bool


@dataclass(frozen=True)
class VerifierResult:
    available: bool
    skipped_reason: str | None
    model_version: str | None
    verifications: list[MeasurementVerification]

    def to_dict(self) -> dict[str, Any]:
        return {
            "available": self.available,
            "skipped_reason": self.skipped_reason,
            "model_version": self.model_version,
            "verifications": [asdict(v) for v in self.verifications],
        }


# Disagreement thresholds. Above these, the verifier overrides the LLM.
THRESHOLDS = {
    "rate": 10.0,        # bpm
    "intervals.pr": 20,  # ms
    "intervals.qrs": 20, # ms
    "intervals.qtc": 30, # ms
    "axis": 20,          # degrees
}


_model: Any = None
_load_error: str | None = None


def _load_model_lazy() -> tuple[Any, str | None]:
    """Returns (model, error_reason). Caches the result so the load only runs
    once per process. On missing torch / weights / etc., returns (None, reason)."""
    global _model, _load_error
    if _model is not None or _load_error is not None:
        return _model, _load_error

    ecg_fm_path = os.environ.get("EKG_TUTOR_ECG_FM_PATH", "").strip()
    if not ecg_fm_path:
        _load_error = "EKG_TUTOR_ECG_FM_PATH not set"
        return None, _load_error
    if not Path(ecg_fm_path).exists():
        _load_error = f"ECG-FM weights directory not found: {ecg_fm_path}"
        return None, _load_error

    try:
        import torch  # noqa: F401
    except ImportError:
        _load_error = "torch is not installed; run: uv pip install torch wfdb"
        return None, _load_error

    try:
        # The real ECG-FM load is two lines from the bowang-lab repo; the
        # exact loader signature depends on which model variant we ship.
        # Phase 4 scaffolding leaves this as a TODO marker so the production
        # wire-up is a focused, reviewable change rather than buried here.
        # When wiring: replace this block with the model loader and weight
        # init, then return the loaded model.
        logger.warning(
            "ECG-FM model loader not implemented — Phase 4 scaffolding only. "
            "Wire up bowang-lab/ECG-FM in pipelines/ecg_fm_verify.py to enable verification."
        )
        _load_error = "ECG-FM loader not implemented (Phase 4 scaffolding only)"
        return None, _load_error
    except Exception as e:
        _load_error = f"ECG-FM load failed: {e}"
        return None, _load_error


def verify_measurements(
    signal_path: str | Path,
    llm_measurements: dict[str, float],
) -> VerifierResult:
    """Runs ECG-FM on the WFDB signal and reconciles with LLM measurements.

    `llm_measurements` is a flat dict of the LLM's numeric reads, keyed by
    the same field strings used in THRESHOLDS. Missing fields are ignored.
    """
    model, err = _load_model_lazy()
    if model is None:
        return VerifierResult(
            available=False,
            skipped_reason=err or "unavailable",
            model_version=None,
            verifications=[],
        )

    # Real ECG-FM call lives here once the loader above is wired. The
    # contract is:
    #   1. Load the WFDB record from signal_path (12-lead × N samples).
    #   2. Forward-pass through ECG-FM to produce the measurement set.
    #   3. Compare each field against the LLM's value; tag agreement.
    #   4. Return per-field verifications with correction_applied=true
    #      when |delta| > threshold.
    # The caller (interpret service) is responsible for actually overwriting
    # the canonical interpretation's measurement fields and tagging the
    # correction in the response.
    raise NotImplementedError(
        "ECG-FM forward pass not implemented; Phase 4 scaffolding only. "
        "Loader returned non-None unexpectedly — investigate."
    )


def is_available() -> bool:
    """Cheap check for /health endpoints. Doesn't trigger the lazy load."""
    return os.environ.get("EKG_TUTOR_ECG_FM_PATH", "").strip() != ""