"""Qwen3-ASR-0.6B transcription plus the selected v8 semantic classifier."""
from __future__ import annotations

from dataclasses import dataclass
import hashlib
import time
import os
from pathlib import Path
import re
from typing import Any

import numpy as np
from safetensors import safe_open
import torch
from torch import nn
from transformers import AutoModel, PreTrainedTokenizerFast
from huggingface_hub import snapshot_download

from .qwen_backend import Qwen3ASRForConditionalGeneration, Qwen3ASRProcessor
from .typed_audio import question_options, typed_answer


MODEL_ID = "Qwen/Qwen3-ASR-0.6B"
REVISION = "5eb144179a02acc5e5ba31e748d22b0cf3e303b0"
SEMANTIC_REPO_ID = "gojiteji/mocovoice-audio-jev-semantic-v8"
SEMANTIC_REVISION = "23779177e4f8abe672cef0d59812e7ee63bb1095"
ASR_VARIANT = "hybrid_channel-alpha050"
ASR_ADAPTATION = Path(__file__).with_name("qwen06_asr_ja_keyword_alpha050.safetensors")
ASR_ADAPTATION_SHA256 = "776d9ff19aff2cce4b5a59dacf56e747c5d2945b06f19970844a875e7f803533"
ACTS = ("backchannel", "confirm", "request", "correct", "cancel", "other")
INTENTS = ("none", "reservation", "order", "change", "cancel")
TASK_LABELS = {"speech_act": ACTS, "intent": INTENTS}


@dataclass(frozen=True)
class ASRTranscription:
    language: str
    text: str


def _apply_asr_adaptation(model: nn.Module, path: Path) -> dict[str, str]:
    """Apply the sealed-test-selected BF16 parameters and reject silent fallback."""
    if not path.is_file():
        raise FileNotFoundError(f"required ASR adaptation is missing: {path.name}")
    fingerprint = hashlib.file_digest(path.open("rb"), "sha256").hexdigest()
    if fingerprint != ASR_ADAPTATION_SHA256:
        raise ValueError("ASR adaptation fingerprint differs from the frozen runtime contract")
    named = dict(model.named_parameters())
    with safe_open(path, framework="pt", device="cpu") as checkpoint:
        metadata = checkpoint.metadata() or {}
        if metadata.get("format") != "qwen06-asr-runtime-weights-v1":
            raise ValueError("unexpected ASR adaptation format")
        if metadata.get("model") != MODEL_ID or metadata.get("revision") != REVISION:
            raise ValueError("ASR adaptation base model differs")
        if metadata.get("candidate") != ASR_VARIANT:
            raise ValueError("ASR adaptation candidate differs")
        keys = list(checkpoint.keys())
        if not keys or any(name not in named for name in keys):
            raise ValueError("ASR adaptation parameter set differs")
        for name in keys:
            value = checkpoint.get_tensor(name)
            parameter = named[name]
            if tuple(value.shape) != tuple(parameter.shape):
                raise ValueError(f"ASR adaptation shape differs: {name}")
            parameter.data.copy_(value.to(parameter.device, dtype=parameter.dtype))
    return {"candidate": ASR_VARIANT, "sha256": fingerprint}


def fixed_semantic_task(question: dict[str, Any]) -> str | None:
    """Return the fixed classifier task only for an exact supported label set."""
    if not isinstance(question, dict) or question.get("type") != "choice":
        return None
    criteria = question.get("criteria")
    if not isinstance(criteria, dict):
        return None
    keys = set(criteria)
    for task, labels in TASK_LABELS.items():
        if keys == set(labels):
            return task
    return None


class QwenASRSemanticPredictor:
    """One ASR decode shared by all supported intent and speech-act questions."""

    def __init__(self, semantic_dir: str | Path, device: str = "cpu"):
        self.device = torch.device(device)
        self.dtype = torch.float32 if self.device.type == "cpu" else torch.bfloat16
        snapshot = snapshot_download(MODEL_ID, revision=REVISION)
        self.processor = Qwen3ASRProcessor.from_pretrained(
            snapshot, local_files_only=True, fix_mistral_regex=True
        )
        self.asr = Qwen3ASRForConditionalGeneration.from_pretrained(
            snapshot, local_files_only=True, dtype=self.dtype, attn_implementation="sdpa"
        ).to(self.device).eval()
        self.asr_adaptation = _apply_asr_adaptation(self.asr, ASR_ADAPTATION)
        self.asr.requires_grad_(False)

        semantic_dir = Path(semantic_dir)
        if not (semantic_dir / "text_encoder/model.safetensors").is_file():
            semantic_dir = Path(snapshot_download(
                SEMANTIC_REPO_ID, revision=SEMANTIC_REVISION,
                token=os.environ.get("AUDIO_JEV_MODEL_TOKEN"),
            ))
        encoder_dir = semantic_dir / "text_encoder"
        self.tokenizer = PreTrainedTokenizerFast(
            tokenizer_file=str(encoder_dir / "tokenizer.json"),
            bos_token="<s>", eos_token="</s>", unk_token="<unk>", pad_token="<pad>",
            sep_token="<sep>", cls_token="<cls>", mask_token="<mask>",
        )
        self.semantic = AutoModel.from_pretrained(
            encoder_dir, local_files_only=True, torch_dtype=torch.float32,
            attn_implementation="sdpa",
        ).to(self.device).eval()
        self.head = nn.Linear(self.semantic.config.hidden_size, len(ACTS) + len(INTENTS)).to(self.device).eval()
        self.head.load_state_dict(torch.load(semantic_dir / "semantic_head.pt", map_location=self.device,
                                             weights_only=True), strict=True)
        self.semantic.requires_grad_(False); self.head.requires_grad_(False)

    def _prompt(self, force_language: str | None = "Japanese") -> str:
        messages = [
            {"role": "system", "content": ""},
            {"role": "user", "content": [{"type": "audio", "audio": ""}]},
        ]
        prompt = self.processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
        if force_language:
            prompt += f"language {force_language}<asr_text>"
        return prompt

    @staticmethod
    def _parse_auto_output(decoded: str) -> ASRTranscription:
        if "<asr_text>" not in decoded:
            return ASRTranscription("", decoded.strip())
        metadata, text = decoded.split("<asr_text>", 1)
        match = re.search(r"(?:^|\n)\s*language\s+([^\s<]+)", metadata, flags=re.IGNORECASE)
        language = match.group(1).strip().capitalize() if match else ""
        if language.casefold() == "none":
            language = ""
        return ASRTranscription(language, text.strip())

    @torch.inference_mode()
    def transcribe_result(self, wave16k: np.ndarray, *, detect_language: bool = False) -> ASRTranscription:
        wave = np.ascontiguousarray(np.clip(np.asarray(wave16k, dtype=np.float32), -1, 1))
        if wave.ndim != 1 or not 4000 <= len(wave) <= 480000 or not np.isfinite(wave).all():
            raise ValueError("Expected finite mono 16-kHz audio between 0.25 and 30 seconds")
        force_language = None if detect_language else "Japanese"
        inputs = self.processor(text=[self._prompt(force_language)], audio=[wave], return_tensors="pt", padding=True)
        inputs = inputs.to(self.device).to(self.dtype)
        generated = self.asr.generate(**inputs, max_new_tokens=192)
        decoded = self.processor.batch_decode(
            generated.sequences[:, inputs["input_ids"].shape[1]:],
            skip_special_tokens=True, clean_up_tokenization_spaces=False,
        )[0].strip()
        if not decoded and not detect_language:
            raise ValueError("Qwen3-ASR returned an empty transcript")
        if detect_language:
            return self._parse_auto_output(decoded)
        return ASRTranscription("Japanese", decoded)

    def transcribe(self, wave16k: np.ndarray) -> str:
        return self.transcribe_result(wave16k).text

    @torch.inference_mode()
    def _probabilities(self, transcript: str, context: str) -> dict[str, dict[str, float]]:
        values = self.tokenizer(
            [context + "[SEP]" + transcript], return_tensors="pt", padding=True,
            truncation=True, max_length=256,
        )
        values.pop("token_type_ids", None)
        logits = self.head(self.semantic(**values.to(self.device)).last_hidden_state[:, 0])[0].float().cpu()
        act = logits[:len(ACTS)].softmax(-1).tolist()
        intent = logits[len(ACTS):].softmax(-1).tolist()
        return {
            "speech_act": dict(zip(ACTS, act)),
            "intent": dict(zip(INTENTS, intent)),
        }

    def predict(self, wave16k: np.ndarray, context: str, questions: dict[str, dict],
                transcription: ASRTranscription | None = None) -> dict[str, Any]:
        if not questions:
            raise ValueError("At least one fixed semantic question is required")
        tasks = {name: fixed_semantic_task(question) for name, question in questions.items()}
        if any(task is None for task in tasks.values()):
            raise ValueError("Unsupported fixed semantic question")
        began = time.perf_counter()
        if transcription is None:
            asr_began = time.perf_counter(); transcription = self.transcribe_result(wave16k)
            asr_ms = (time.perf_counter() - asr_began) * 1000
        else:
            asr_ms = 0.0
        transcript = transcription.text
        semantic_began = time.perf_counter(); fixed = self._probabilities(transcript, context)
        answers = {}
        for name, question in questions.items():
            labels = fixed[tasks[name]]
            pairs = question_options(question)
            answers[name] = typed_answer(question, [labels[key] for key, _ in pairs])
        semantic_ms = (time.perf_counter() - semantic_began) * 1000
        return {
            "model": "qwen06-asr-semantic-natural-v8",
            "transcript": transcript,
            "answers": answers,
            "evidence": {
                "audio_duration_ms": len(wave16k) / 16,
                "transcript_source": MODEL_ID,
                "asr_revision": REVISION,
                "asr_adaptation": self.asr_adaptation,
                "detected_language": transcription.language,
                "transcript_reused": asr_ms == 0.0,
                "semantic_model": "modernbert-ja-130m-semantic-natural-v8",
                "fixed_tasks": tasks,
            },
            "latency_ms": {"asr": asr_ms, "semantic": semantic_ms,
                           "total": (time.perf_counter() - began) * 1000},
            "action_executable": False,
            "scope": "Selected fixed intent/speech-act classifier; probabilities are not calibrated correctness.",
        }