"""Transcription handler for local Apple Silicon backends."""

from __future__ import annotations

import shutil
import tempfile
import time
import wave
from contextlib import contextmanager
from dataclasses import dataclass, field
from importlib import import_module
from pathlib import Path
from threading import Event
from typing import Any
from urllib.parse import urljoin

import httpx
import structlog

from hive_compute.model_ids import TURBO_MODEL_ALIASES, canonical_model_identifier, split_model_identifier
from hive_compute.models import ComputeJob
from hive_compute.url_safety import validate_outbound_http_url

logger = structlog.get_logger()

# Max audio file size (500MB)
MAX_AUDIO_SIZE = 500 * 1024 * 1024


class StageCanceledError(RuntimeError):
    """Raised when a staged job is canceled before inference starts."""


@dataclass(frozen=True)
class WhisperModelSpec:
    """Resolved backend/model settings for a single transcription request."""

    identifier: str
    backend: str
    model: str
    quantization: str
    language: str
    repo: str | None = None


@dataclass
class PreparedWhisperJob:
    """Audio that has been downloaded and optionally decoded for inference."""

    job_id: str
    shard_id: str | None
    audio_url: str
    audio_path: Path
    model: WhisperModelSpec
    stage_timings: dict[str, float] = field(default_factory=dict)
    cleanup_paths: list[Path] = field(default_factory=list)

    def cleanup(self) -> None:
        for path in self.cleanup_paths:
            path.unlink(missing_ok=True)


class WhisperHandler:
    """Transcribe audio using local backends such as mlx-whisper, parakeet-mlx, and whisper.cpp."""

    def __init__(self, config: dict):
        self.model = "large-v3"
        self.quantization = "4bit"
        self.language = "en"
        self.backend = "mlx-whisper"
        self.repo = ""
        self.cache_dir = Path("~/.ambient/audio-cache").expanduser()
        self.supported_models: list[str] = []
        self.prewarm_model_on_startup = True
        self._parakeet_model: Any | None = None
        self._parakeet_model_id: str | None = None
        self.update_config(config)

    def update_config(self, config: dict) -> bool:
        """Apply handler config changes and report whether anything changed."""
        next_model = config.get("model", "large-v3")
        next_quantization = config.get("quantization", "4bit")
        next_language = config.get("language", "en")
        next_backend = config.get("backend", "mlx-whisper")
        raw_repo = config.get("repo", "")
        next_repo = "" if raw_repo is None else str(raw_repo).strip()
        next_cache_dir = Path(config.get("cache_dir", "~/.ambient/audio-cache")).expanduser()
        next_supported_models = list(config.get("supported_models") or [])
        next_prewarm = bool(config.get("prewarm_model_on_startup", True))
        next_cache_dir.mkdir(parents=True, exist_ok=True)

        changed = (
            self.model != next_model
            or self.quantization != next_quantization
            or self.language != next_language
            or self.backend != next_backend
            or self.repo != next_repo
            or self.cache_dir != next_cache_dir
            or self.supported_models != next_supported_models
            or self.prewarm_model_on_startup != next_prewarm
        )
        current_parakeet_repo = self._resolve_parakeet_repo(self.model) if self.backend == "parakeet-mlx" else None
        next_parakeet_repo = self._resolve_parakeet_repo(next_model) if next_backend == "parakeet-mlx" else None
        if self.backend != next_backend or current_parakeet_repo != next_parakeet_repo:
            self._parakeet_model = None
            self._parakeet_model_id = None
        self.model = next_model
        self.quantization = next_quantization
        self.language = next_language
        self.backend = next_backend
        self.repo = next_repo
        self.cache_dir = next_cache_dir
        self.supported_models = next_supported_models
        self.prewarm_model_on_startup = next_prewarm
        return changed

    def run(self, job: ComputeJob) -> dict[str, Any]:
        """Download audio, trim to shard boundaries if applicable, transcribe."""
        prepared = self.prepare_job(job)
        try:
            return self.run_prepared(prepared)
        finally:
            prepared.cleanup()

    def prepare_job(self, job: ComputeJob, *, cancel_event: Event | None = None) -> PreparedWhisperJob:
        """Stage audio download and decode work ahead of inference."""
        audio_url = job.input_url or job.input_data.get("audio_url")
        if not audio_url:
            raise ValueError("No audio_url in job input")

        model_spec = self._resolve_model_spec_from_input(job.input_data or {})
        work_unit_id = self._work_unit_id(job.id, job.shard_id)

        started_at = time.perf_counter()
        audio_path = self._download_audio(audio_url, work_unit_id, cancel_event=cancel_event)
        download_elapsed = time.perf_counter() - started_at
        cleanup_paths = [audio_path]
        local_audio_path = audio_path
        decode_elapsed = 0.0

        if job.start_seconds is not None and job.end_seconds is not None:
            if cancel_event and cancel_event.is_set():
                audio_path.unlink(missing_ok=True)
                raise StageCanceledError("Staging canceled before audio decode")

            started_at = time.perf_counter()
            trimmed_path = self._trim_audio(audio_path, job.start_seconds, job.end_seconds, work_unit_id)
            decode_elapsed = time.perf_counter() - started_at
            local_audio_path = trimmed_path
            cleanup_paths.append(trimmed_path)

        if cancel_event and cancel_event.is_set():
            for path in cleanup_paths:
                path.unlink(missing_ok=True)
            raise StageCanceledError("Staging canceled before inference")

        return PreparedWhisperJob(
            job_id=job.id,
            shard_id=job.shard_id,
            audio_url=audio_url,
            audio_path=local_audio_path,
            model=model_spec,
            stage_timings={
                "download_seconds": round(download_elapsed, 3),
                "decode_seconds": round(decode_elapsed, 3),
            },
            cleanup_paths=cleanup_paths,
        )

    def run_prepared(self, prepared: PreparedWhisperJob) -> dict[str, Any]:
        """Run inference against already-staged audio."""
        original_settings = (self.model, self.quantization, self.language, self.repo)
        try:
            self.model = prepared.model.model
            self.quantization = prepared.model.quantization
            self.language = prepared.model.language
            self.repo = prepared.model.repo or ""
            started_at = time.perf_counter()
            result = self._transcribe(prepared.audio_path)
            inference_elapsed = time.perf_counter() - started_at
            result["pipeline_timings"] = {
                **prepared.stage_timings,
                "inference_seconds": round(inference_elapsed, 3),
                "total_seconds": round(
                    prepared.stage_timings.get("download_seconds", 0.0)
                    + prepared.stage_timings.get("decode_seconds", 0.0)
                    + inference_elapsed,
                    3,
                ),
            }
            return result
        finally:
            self.model, self.quantization, self.language, self.repo = original_settings

    def model_capabilities(self, models: list[str] | None = None) -> list[dict[str, Any]]:
        """Describe whether configured models are warmed, installed, or downloadable."""
        capabilities: list[dict[str, Any]] = []
        for identifier in models or self.configured_model_identifiers():
            requested_model, requested_quantization = split_model_identifier(
                identifier,
                default_quantization=self.quantization,
            )
            model_spec = self._resolve_model_spec(requested_model, requested_quantization)
            warmed = self._is_model_warmed(model_spec)
            installed = warmed or self._is_model_installed(model_spec)
            downloadable = self._is_model_downloadable(model_spec)
            if warmed:
                state = "warmed"
            elif installed:
                state = "installed"
            elif downloadable:
                state = "downloadable"
            else:
                state = "unavailable"
            capabilities.append(
                {
                    "model": model_spec.identifier,
                    "backend": model_spec.backend,
                    "repo": model_spec.repo,
                    "state": state,
                    "installed": installed,
                    "warmed": warmed,
                    "downloadable": downloadable,
                    "primary": model_spec.identifier == self.current_model_identifier(),
                }
            )
        return capabilities

    def pull_model(self, model: str | None = None, quantization: str | None = None) -> dict[str, Any]:
        """Download a model through the backend's normal load path."""
        return self.prewarm(model=model, quantization=quantization, reason="pull")

    def prewarm(
        self,
        model: str | None = None,
        quantization: str | None = None,
        *,
        reason: str = "startup",
    ) -> dict[str, Any]:
        """Intentionally load the configured model so the next real shard is warm."""
        model_spec = self._resolve_model_spec(model, quantization)
        if self._is_model_warmed(model_spec):
            return {
                "status": "already_warm",
                "model": model_spec.identifier,
                "backend": model_spec.backend,
                "repo": model_spec.repo,
                "reason": reason,
            }

        original_settings = (self.model, self.quantization, self.language, self.repo)
        try:
            self.model = model_spec.model
            self.quantization = model_spec.quantization
            self.language = model_spec.language
            self.repo = model_spec.repo or ""
            with self._temporary_silent_wav() as silent_wav:
                started_at = time.perf_counter()
                self._transcribe(silent_wav)
                elapsed = time.perf_counter() - started_at
            return {
                "status": "warmed",
                "model": model_spec.identifier,
                "backend": model_spec.backend,
                "repo": model_spec.repo,
                "elapsed_seconds": round(elapsed, 3),
                "reason": reason,
            }
        except (FileNotFoundError, ImportError, ModuleNotFoundError) as exc:
            return {
                "status": "skipped",
                "model": model_spec.identifier,
                "backend": model_spec.backend,
                "repo": model_spec.repo,
                "error": str(exc),
                "reason": reason,
            }
        except Exception as exc:
            return {
                "status": "failed",
                "model": model_spec.identifier,
                "backend": model_spec.backend,
                "repo": model_spec.repo,
                "error": str(exc),
                "reason": reason,
            }
        finally:
            self.model, self.quantization, self.language, self.repo = original_settings

    def configured_model_identifiers(self) -> list[str]:
        configured = [*self.supported_models, self.current_model_identifier()]
        models: list[str] = []
        seen: set[str] = set()
        for model in configured:
            identifier = canonical_model_identifier(model, default_quantization="none")
            if identifier in seen:
                continue
            seen.add(identifier)
            models.append(identifier)
        return models

    def current_model_identifier(self) -> str:
        return canonical_model_identifier(self.model, self.quantization)

    def _resolve_model_spec_from_input(self, input_config: dict[str, Any]) -> WhisperModelSpec:
        requested_model, requested_quantization = split_model_identifier(
            input_config.get("model", self.model),
            input_config.get("quantization"),
            default_quantization=self.quantization,
        )
        requested_language = input_config.get("language", self.language)
        return self._resolve_model_spec(
            requested_model,
            requested_quantization,
            language=requested_language,
        )

    def _resolve_model_spec(
        self,
        model: str | None = None,
        quantization: str | None = None,
        *,
        language: str | None = None,
    ) -> WhisperModelSpec:
        requested_model, requested_quantization = split_model_identifier(
            model or self.model,
            quantization,
            default_quantization=self.quantization,
        )
        requested_language = str(language or self.language).strip() or self.language
        repo = None
        if self.backend == "mlx-whisper":
            repo = self._resolve_mlx_repo(requested_model, requested_quantization, repo_override=self.repo or None)
        elif self.backend == "parakeet-mlx":
            repo = self._resolve_parakeet_repo(requested_model)
        elif self.backend == "whisper-cpp":
            repo = (self.repo or requested_model or self.model).strip()
        return WhisperModelSpec(
            identifier=canonical_model_identifier(requested_model, requested_quantization),
            backend=self.backend,
            model=requested_model,
            quantization=requested_quantization,
            language=requested_language,
            repo=repo or None,
        )

    @staticmethod
    def _work_unit_id(job_id: str, shard_id: str | None) -> str:
        if shard_id:
            return f"{job_id}-{shard_id}"
        return job_id

    @staticmethod
    def _backend_module_available(module_name: str) -> bool:
        try:
            import_module(module_name)
        except Exception:
            return False
        return True

    def _is_model_downloadable(self, model_spec: WhisperModelSpec) -> bool:
        if self.backend in {"mlx-whisper", "parakeet-mlx"}:
            return bool(model_spec.repo)
        if self.backend == "whisper-cpp":
            return Path(str(model_spec.repo or model_spec.model)).expanduser().exists()
        return False

    def _is_model_installed(self, model_spec: WhisperModelSpec) -> bool:
        if self.backend == "whisper-cpp":
            return Path(str(model_spec.repo or model_spec.model)).expanduser().exists()
        if not model_spec.repo:
            return False
        try:
            hub = import_module("huggingface_hub")
        except Exception:
            return False

        snapshot_download = getattr(hub, "snapshot_download", None)
        if snapshot_download is None:
            return False

        try:
            snapshot_download(model_spec.repo, local_files_only=True)
        except Exception:
            return False
        return True

    def _is_model_warmed(self, model_spec: WhisperModelSpec) -> bool:
        if self.backend == "parakeet-mlx":
            return self._parakeet_model is not None and self._parakeet_model_id == model_spec.repo
        if self.backend == "mlx-whisper":
            try:
                transcribe_module = import_module("mlx_whisper.transcribe")
            except Exception:
                return False
            model_holder = getattr(transcribe_module, "ModelHolder", None)
            return bool(model_holder and getattr(model_holder, "model_path", None) == model_spec.repo)
        return False

    @staticmethod
    @contextmanager
    def _temporary_silent_wav():
        tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
        tmp.close()
        try:
            with wave.open(tmp.name, "wb") as wav_file:
                wav_file.setnchannels(1)
                wav_file.setsampwidth(2)
                wav_file.setframerate(16000)
                wav_file.writeframes(b"\x00\x00" * 1600)
            yield Path(tmp.name)
        finally:
            Path(tmp.name).unlink(missing_ok=True)

    def _download_audio(self, url: str, job_id: str, *, cancel_event: Event | None = None) -> Path:
        """Download audio file to local cache."""
        audio_path = self.cache_dir / f"{job_id}.audio"
        logger.info("Downloading audio", url=url[:80], dest=str(audio_path))

        current_url = url
        for _ in range(4):
            validate_outbound_http_url(current_url, purpose="worker audio_url")
            with httpx.stream("GET", current_url, follow_redirects=False, timeout=300) as resp:
                if resp.status_code in {301, 302, 303, 307, 308}:
                    location = resp.headers.get("location")
                    if not location:
                        raise ValueError("Redirected audio URL did not include a location header")
                    current_url = urljoin(current_url, location)
                    continue

                resp.raise_for_status()
                total = 0
                with open(audio_path, "wb") as f:
                    for chunk in resp.iter_bytes(chunk_size=8192):
                        if cancel_event and cancel_event.is_set():
                            audio_path.unlink(missing_ok=True)
                            raise StageCanceledError("Staging canceled during download")
                        total += len(chunk)
                        if total > MAX_AUDIO_SIZE:
                            audio_path.unlink(missing_ok=True)
                            raise ValueError(f"Audio file too large (>{MAX_AUDIO_SIZE // 1024 // 1024}MB)")
                        f.write(chunk)
                break
        else:
            raise ValueError("Audio URL redirected too many times")

        logger.info("Audio downloaded", size_mb=round(total / 1024 / 1024, 1))
        return audio_path

    def _trim_audio(self, audio_path: Path, start: float, end: float, job_id: str) -> Path:
        """Trim audio to [start, end] seconds using ffmpeg."""
        import subprocess

        if not shutil.which("ffmpeg"):
            raise RuntimeError(
                "ffmpeg is required for shard processing but was not found on PATH. "
                "Install ffmpeg (e.g. 'brew install ffmpeg') and retry."
            )

        shard_label = f"{start:.3f}-{end:.3f}".replace(".", "_")
        trimmed = self.cache_dir / f"{job_id}.{shard_label}.trimmed.wav"
        duration = end - start

        logger.info("Trimming audio", start=start, end=end, duration=duration)
        try:
            subprocess.run(
                [
                    "ffmpeg", "-y",
                    "-ss", str(start),
                    "-t", str(duration),
                    "-i", str(audio_path),
                    "-ar", "16000",
                    "-ac", "1",
                    str(trimmed),
                ],
                check=True,
                capture_output=True,
                text=True,
            )
        except subprocess.CalledProcessError as exc:
            trimmed.unlink(missing_ok=True)
            stderr = (exc.stderr or "")[:500]
            raise RuntimeError(f"ffmpeg trim failed (exit {exc.returncode}): {stderr}") from exc

        return trimmed

    def _transcribe(self, audio_path: Path) -> dict[str, Any]:
        """Run Whisper on the audio file."""
        if self.backend == "mlx-whisper":
            return self._transcribe_mlx(audio_path)
        elif self.backend == "parakeet-mlx":
            return self._transcribe_parakeet(audio_path)
        elif self.backend == "whisper-cpp":
            return self._transcribe_cpp(audio_path)
        else:
            raise ValueError(f"Unknown Whisper backend: {self.backend}")

    @staticmethod
    def _resolve_mlx_repo(model: str, quantization: str = "none", repo_override: str | None = None) -> str:
        """Resolve model name and quantization to a HuggingFace MLX repo."""
        if repo_override:
            return repo_override
        turbo_repos = {
            "none": "mlx-community/whisper-large-v3-turbo",
            "4bit": "mlx-community/whisper-large-v3-turbo-q4",
            "8bit": "LibraxisAI/whisper-large-v3-turbo-mlx-q8",
        }
        if model in TURBO_MODEL_ALIASES:
            return turbo_repos.get(quantization, turbo_repos["none"])

        repo = f"mlx-community/whisper-{model}-mlx"
        if quantization == "none":
            return repo
        return f"{repo}-{quantization}"

    @staticmethod
    def _resolve_parakeet_repo(model: str) -> str:
        """Resolve model aliases to the default parakeet-mlx repository."""
        overrides = {
            "large-v3": "mlx-community/parakeet-tdt-0.6b-v3",
            "large-v3-turbo": "mlx-community/parakeet-tdt-0.6b-v3",
            "turbo": "mlx-community/parakeet-tdt-0.6b-v3",
            "parakeet-v3": "mlx-community/parakeet-tdt-0.6b-v3",
            "parakeet-tdt-0.6b-v3": "mlx-community/parakeet-tdt-0.6b-v3",
        }
        resolved = overrides.get(model, model)
        if "/" in resolved:
            return resolved
        if resolved.startswith("parakeet-"):
            return f"mlx-community/{resolved}"
        return resolved

    @staticmethod
    def _extract_alignment_end(result: Any) -> float | None:
        sentences = getattr(result, "sentences", None)
        if not sentences:
            return None

        end_times: list[float] = []
        for sentence in sentences:
            if isinstance(sentence, dict):
                end = sentence.get("end")
            else:
                end = getattr(sentence, "end", None)
            if isinstance(end, int | float):
                end_times.append(float(end))
        return max(end_times) if end_times else None

    def _transcribe_mlx(self, audio_path: Path) -> dict[str, Any]:
        """Transcribe using mlx-whisper (Apple MLX framework)."""
        import traceback

        try:
            import mlx_whisper
            transcribe_module = import_module("mlx_whisper.transcribe")
        except Exception as exc:
            raise RuntimeError(f"mlx-whisper import failed: {exc}\n{traceback.format_exc()}") from exc

        repo = self._resolve_mlx_repo(self.model, self.quantization, repo_override=self.repo)
        cache_state = "warm"
        model_holder = getattr(transcribe_module, "ModelHolder", None)
        if model_holder is None or getattr(model_holder, "model_path", None) != repo:
            cache_state = "cold"

        logger.info(
            "Transcribing with mlx-whisper",
            model=self.model,
            quantization=self.quantization,
            repo=repo,
            cache_state=cache_state,
        )
        started_at = time.perf_counter()
        try:
            result = mlx_whisper.transcribe(
                str(audio_path),
                path_or_hf_repo=repo,
                language=self.language,
            )
        except Exception as exc:
            raise RuntimeError(f"mlx-whisper transcription failed: {exc}\n{traceback.format_exc()}") from exc
        elapsed = time.perf_counter() - started_at
        log_data = {
            "model": self.model,
            "quantization": self.quantization,
            "repo": repo,
            "cache_state": cache_state,
            "elapsed_seconds": round(elapsed, 3),
        }
        if repo == "mlx-community/whisper-large-v3-mlx":
            log_data["estimated_vram_gb"] = 3.0
        elif repo == "mlx-community/whisper-large-v3-mlx-8bit":
            log_data["estimated_vram_gb"] = 1.5
        elif repo == "mlx-community/whisper-large-v3-mlx-4bit":
            log_data["estimated_vram_gb"] = 0.8
        logger.info("mlx-whisper transcription finished", **log_data)

        transcript = result.get("text", "").strip()
        duration = result.get("duration")
        model_id = canonical_model_identifier(self.model, self.quantization)

        return {
            "transcript": transcript,
            "model": f"whisper-{model_id}",
            "backend": "mlx-whisper",
            "language": self.language,
            "duration_seconds": duration,
            "word_count": len(transcript.split()) if transcript else 0,
        }

    def _transcribe_parakeet(self, audio_path: Path) -> dict[str, Any]:
        """Transcribe using parakeet-mlx with an in-process warm model cache."""
        parakeet_module = import_module("parakeet_mlx")
        repo = self._resolve_parakeet_repo(self.model)
        cache_state = "warm"
        if self._parakeet_model is None or self._parakeet_model_id != repo:
            cache_state = "cold"
            self._parakeet_model = parakeet_module.from_pretrained(repo)
            self._parakeet_model_id = repo

        logger.info("Transcribing with parakeet-mlx", model=self.model, repo=repo, cache_state=cache_state)
        started_at = time.perf_counter()
        result = self._parakeet_model.transcribe(str(audio_path))
        elapsed = time.perf_counter() - started_at
        logger.info(
            "parakeet-mlx transcription finished",
            model=self.model,
            repo=repo,
            cache_state=cache_state,
            elapsed_seconds=round(elapsed, 3),
        )

        transcript = getattr(result, "text", "").strip()
        duration = self._extract_alignment_end(result)
        return {
            "transcript": transcript,
            "model": repo.rsplit("/", 1)[-1],
            "backend": "parakeet-mlx",
            "language": self.language,
            "duration_seconds": duration,
            "word_count": len(transcript.split()) if transcript else 0,
        }

    def _transcribe_cpp(self, audio_path: Path) -> dict[str, Any]:
        """Transcribe using whisper.cpp CLI."""
        import subprocess

        logger.info("Transcribing with whisper.cpp", model=self.model)

        with tempfile.NamedTemporaryFile(suffix=".txt", delete=False) as tmp:
            output_path = Path(tmp.name)

        try:
            subprocess.run(
                [
                    "whisper-cpp",
                    "--model", self.model,
                    "--language", self.language,
                    "--output-txt",
                    "--output-file", str(output_path.with_suffix("")),
                    str(audio_path),
                ],
                check=True,
                capture_output=True,
                text=True,
            )
            transcript = output_path.read_text().strip()
        finally:
            output_path.unlink(missing_ok=True)
        model_id = canonical_model_identifier(self.model, self.quantization)

        return {
            "transcript": transcript,
            "model": f"whisper-{model_id}",
            "backend": "whisper-cpp",
            "language": self.language,
            "word_count": len(transcript.split()) if transcript else 0,
        }
