"""Nemotron coach reasoning wrapper with rule-based fallback."""

from __future__ import annotations

import json
import os
import urllib.request
from typing import Any

from .prompts import SYSTEM_PROMPT, build_coach_prompt
from .schemas import CoachResponse


class NemotronCoach:
    """Small replaceable coach-reasoning module.

    V1 keeps this optional. Set NEMOTRON_API_URL, NEMOTRON_API_KEY, and
    NEMOTRON_MODEL to call a hosted Nemotron-compatible chat endpoint. Without
    those variables, the deterministic fallback keeps the Space fully working.
    """

    def __init__(self) -> None:
        self.api_url = os.getenv("NEMOTRON_API_URL", "").strip()
        self.api_key = os.getenv("NEMOTRON_API_KEY", "").strip()
        self.model = os.getenv("NEMOTRON_MODEL", "nvidia/nemotron-nano-4b-v1")

    def reason(self, practice_state: dict[str, Any], history: list[dict[str, Any]]) -> dict[str, Any]:
        if self.api_url and self.api_key:
            try:
                return self._call_nemotron(practice_state, history)
            except Exception as error:  # noqa: BLE001 - fallback is intentional.
                fallback = self._fallback(practice_state)
                fallback["diagnosis"] += f" Nemotron fallback used: {error.__class__.__name__}."
                return fallback
        return self._fallback(practice_state)

    def _call_nemotron(self, practice_state: dict[str, Any], history: list[dict[str, Any]]) -> dict[str, Any]:
        payload = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": build_coach_prompt(practice_state, history)},
            ],
            "temperature": 0.2,
            "max_tokens": 220,
        }
        request = urllib.request.Request(
            self.api_url,
            data=json.dumps(payload).encode("utf-8"),
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
            },
            method="POST",
        )
        with urllib.request.urlopen(request, timeout=10) as response:
            raw = json.loads(response.read().decode("utf-8"))

        content = raw.get("choices", [{}])[0].get("message", {}).get("content", "{}")
        parsed = self._parse_model_json(content)
        parsed["model_source"] = "nemotron"
        return self._normalize_response(parsed, practice_state)

    def _parse_model_json(self, content: str) -> dict[str, Any]:
        """Parse strict JSON, tolerating fenced responses from hosted endpoints."""
        content = content.strip()
        if content.startswith("```"):
            content = content.strip("`")
            if content.startswith("json"):
                content = content[4:].strip()
        start = content.find("{")
        end = content.rfind("}")
        if start >= 0 and end > start:
            content = content[start : end + 1]
        return json.loads(content)

    def _normalize_response(self, response: dict[str, Any], practice_state: dict[str, Any]) -> dict[str, Any]:
        fallback = self._fallback(practice_state)
        normalized = {**fallback, **response}
        normalized["overlay_cues"] = list(normalized.get("overlay_cues") or [])
        normalized["agent_trace"] = list(normalized.get("agent_trace") or fallback["agent_trace"])
        normalized["success_trigger"] = bool(practice_state.get("success_trigger", normalized.get("success_trigger", False)))
        return normalized

    def _fallback(self, practice_state: dict[str, Any]) -> dict[str, Any]:
        state = practice_state.get("state", "idle")
        active_step = practice_state.get("active_step", "start")
        success = bool(practice_state.get("success_trigger", False))
        face = practice_state.get("face", {}) or {}
        audio = practice_state.get("audio", {}) or {}

        messages = {
            "idle": "Press Start Practice when you are ready.",
            "no_face": "Please face the camera and center your mouth in the frame.",
            "mouth_too_open": "Make the opening smaller and relax your jaw.",
            "not_rounded": "Pull both lip corners inward, like saying 'u'.",
            "asymmetric_mouth": "Keep both lip corners level and centered.",
            "mouth_ready_no_airflow": "Good mouth shape. Now blow gently.",
            "airflow_no_tone": "You are close. Make the air stream narrower.",
            "stable_whistle": "Great! Hold that whistle steady.",
        }
        drills = {
            "idle": "start_session",
            "no_face": "center_face",
            "mouth_too_open": "small_opening_hold",
            "not_rounded": "lip_rounding",
            "asymmetric_mouth": "corner_leveling",
            "mouth_ready_no_airflow": "gentle_airflow",
            "airflow_no_tone": "narrow_air_stream",
            "stable_whistle": "steady_tone_hold",
        }
        cues = {
            "mouth_too_open": [{"type": "mouth_opening", "direction": "smaller"}],
            "not_rounded": [{"type": "lip_corners", "direction": "inward"}],
            "asymmetric_mouth": [{"type": "symmetry", "direction": "level"}],
            "airflow_no_tone": [{"type": "air_stream", "direction": "narrower"}],
            "stable_whistle": [{"type": "success", "direction": "hold"}],
        }
        trace = [
            {
                "step": "Observe",
                "action": "read_state",
                "detail": (
                    f"mouth {round(float(face.get('mouth_shape_score', 0)) * 100)}%, "
                    f"airflow {round(float(audio.get('airflow_score', 0)) * 100)}%, "
                    f"tone {round(float(audio.get('pitch_stability_score', 0)) * 100)}%"
                ),
            },
            {
                "step": "Diagnose",
                "action": state,
                "detail": f"Current blocker is {state}.",
            },
            {
                "step": "Plan",
                "action": drills.get(state, "observe"),
                "detail": f"Focus on one micro-drill: {drills.get(state, 'observe')}.",
            },
            {
                "step": "Act",
                "action": "coach_message",
                "detail": messages.get(state, "Keep adjusting one small thing at a time."),
            },
        ]
        response = CoachResponse(
            coach_message=messages.get(state, "Keep adjusting one small thing at a time."),
            diagnosis=f"Rule fallback classified state as {state}.",
            overlay_cues=cues.get(state, []),
            active_step=active_step,
            agent_trace=trace,
            next_drill=drills.get(state, "observe"),
            difficulty_adjustment="easier" if state in {"no_face", "mouth_too_open"} else "hold",
            success_trigger=success,
            model_source="rule_fallback",
        )
        return response.to_dict()