"""
MovementClassifierAgent — identifies which FMS test is in the clip.

Input:  IngestResult (keyframes), Pose2DResult (skeleton context)
Output: MovementResult(test_name, side, confidence)
Failure: returns MovementResult(test_name="unknown") — pipeline stops and asks for manual override.
Model:  Qwen3-VL-8B-Instruct via llama.cpp (8B params, Apache-2.0).
Gated:  No.
"""
from __future__ import annotations

import logging
from pathlib import Path

from formscout import config
from formscout.types import IngestResult, Pose2DResult, MovementResult
from formscout.serving.llama_cpp import LlamaCppClient

logger = logging.getLogger(__name__)

_PROMPT_PATH = Path(__file__).parent / "prompts" / "c1_classifier.md"


class MovementClassifierAgent:
    """Classifies which FMS test is being performed via VLM or manual override."""

    def __init__(self):
        self._client = LlamaCppClient(port=config.LLAMA_CPP_PORT_VLM)
        self._system_prompt = _PROMPT_PATH.read_text(encoding="utf-8")

    def run(
        self,
        ingest: IngestResult,
        pose2d: Pose2DResult | None = None,
        manual_override: str | None = None,
    ) -> MovementResult:
        """
        Classify the movement. If manual_override is provided, use it directly.
        Otherwise, use VLM inference on keyframes.
        """
        if manual_override and manual_override != "unknown":
            return MovementResult(
                test_name=manual_override, side="na",
                confidence=1.0, notes="manual override",
            )

        if not self._client.available:
            return MovementResult(
                test_name="unknown", side="na", confidence=0.0,
                notes="VLM server unavailable — use manual override",
            )

        # Select keyframes for classification (3 evenly spaced)
        n = len(ingest.frames)
        indices = [0, n // 2, n - 1] if n >= 3 else list(range(n))
        images = self._encode_frames(ingest.frames, indices)

        prompt = f"{self._system_prompt}\n\nClassify this movement from the keyframes shown."
        result = self._client.complete(prompt, images=images, max_tokens=256, temperature=0.1)

        return self._parse_response(result)

    def _encode_frames(self, frames: list, indices: list[int]) -> list[str]:
        """Encode selected frames as base64 JPEG for the VLM."""
        import cv2
        import base64

        encoded = []
        for idx in indices:
            if idx < len(frames):
                _, buf = cv2.imencode(".jpg", frames[idx], [cv2.IMWRITE_JPEG_QUALITY, 80])
                encoded.append(base64.b64encode(buf.tobytes()).decode())
        return encoded

    def _parse_response(self, result: dict) -> MovementResult:
        """Parse VLM JSON response into MovementResult."""
        if "error" in result:
            return MovementResult(
                test_name="unknown", side="na", confidence=0.0,
                notes=f"VLM error: {result['error']}",
            )

        test = result.get("test", "unknown")
        side = result.get("side", "na")
        confidence = float(result.get("confidence", 0.0))
        reason = result.get("reason", "")

        valid_tests = {
            "deep_squat", "hurdle_step", "inline_lunge",
            "shoulder_mobility", "active_slr",
            "trunk_stability_pushup", "rotary_stability", "unknown",
        }
        if test not in valid_tests:
            test = "unknown"

        if side not in ("left", "right", "na"):
            side = "na"

        return MovementResult(
            test_name=test, side=side,
            confidence=confidence, notes=reason,
        )