from __future__ import annotations

import asyncio
import hmac
import json
import math
import os
import tempfile
from pathlib import Path
from typing import Any

from fastapi import APIRouter, HTTPException, Request

router = APIRouter(prefix="/v1/diarization", tags=["diarization"])

_pipeline: Any = None
_pipeline_lock = asyncio.Lock()
_inference_lock = asyncio.Lock()


def _max_input_bytes() -> int:
    try:
        value = int(os.getenv("DIARIZATION_MAX_INPUT_BYTES", "1073741824"))
    except ValueError as exc:
        raise HTTPException(status_code=500, detail="Invalid diarization size configuration") from exc
    if value < 1_048_576:
        raise HTTPException(status_code=500, detail="Diarization size limit is too small")
    return value


def _max_response_spans() -> int:
    try:
        value = int(os.getenv("DIARIZATION_MAX_RESPONSE_SPANS", "10000"))
    except ValueError as exc:
        raise HTTPException(status_code=500, detail="Invalid diarization response limit") from exc
    if value < 1:
        raise HTTPException(status_code=500, detail="Diarization response limit is invalid")
    return value


def _max_speakers() -> int:
    try:
        value = int(os.getenv("DIARIZATION_MAX_SPEAKERS", "32"))
    except ValueError as exc:
        raise HTTPException(status_code=500, detail="Invalid diarization speaker limit") from exc
    if not 1 <= value <= 256:
        raise HTTPException(status_code=500, detail="Diarization speaker limit is invalid")
    return value


def _max_duration_seconds() -> float:
    try:
        value = float(os.getenv("DIARIZATION_MAX_DURATION_SECONDS", "14400"))
    except ValueError as exc:
        raise HTTPException(status_code=500, detail="Invalid diarization duration limit") from exc
    if not 1 <= value <= 86400:
        raise HTTPException(status_code=500, detail="Diarization duration limit is invalid")
    return value


def _inference_timeout_seconds() -> float:
    try:
        value = float(os.getenv("DIARIZATION_INFERENCE_TIMEOUT_SECONDS", "900"))
    except ValueError as exc:
        raise HTTPException(status_code=500, detail="Invalid diarization inference timeout") from exc
    if value < 1:
        raise HTTPException(status_code=500, detail="Diarization inference timeout is invalid")
    return value


def _configured_api_key() -> str:
    return os.getenv("DIARIZATION_API_KEY", "").strip()


def _authorize(request: Request) -> None:
    expected = _configured_api_key()
    if not expected:
        raise HTTPException(status_code=503, detail="Diarization API authentication is not configured")
    authorization = request.headers.get("authorization", "")
    scheme, _, presented = authorization.partition(" ")
    if scheme.lower() != "bearer" or not presented or not hmac.compare_digest(
        presented.strip(), expected
    ):
        raise HTTPException(status_code=401, detail="Invalid diarization API credentials")


async def _get_pipeline() -> Any:
    global _pipeline
    if _pipeline is not None:
        return _pipeline

    async with _pipeline_lock:
        if _pipeline is not None:
            return _pipeline

        token = os.getenv("HF_TOKEN", "").strip()
        model_id = os.getenv("DIARIZATION_MODEL_ID", "pyannote/speaker-diarization-3.1").strip()
        if not token:
            raise HTTPException(status_code=503, detail="Diarization model credentials are not configured")

        try:
            from pyannote.audio import Pipeline
        except ImportError as exc:
            raise HTTPException(status_code=503, detail="Diarization backend is not installed") from exc

        try:
            pipeline = Pipeline.from_pretrained(model_id, token=token)
            try:
                import torch
                pipeline.to(torch.device("cpu"))
            except Exception:
                pass
            _pipeline = pipeline
            return pipeline
        except Exception as exc:
            raise HTTPException(status_code=503, detail="Diarization model could not be loaded") from exc


async def _save_body(request: Request, destination: Path) -> int:
    limit = _max_input_bytes()
    content_length = request.headers.get("content-length")
    if content_length:
        try:
            declared = int(content_length)
        except ValueError as exc:
            raise HTTPException(status_code=400, detail="Invalid Content-Length") from exc
        if declared < 0:
            raise HTTPException(status_code=400, detail="Invalid Content-Length")
        if declared > limit:
            raise HTTPException(status_code=413, detail="Diarization input exceeds the configured size limit")

    total = 0
    try:
        with destination.open("wb") as handle:
            async for chunk in request.stream():
                if not chunk:
                    continue
                total += len(chunk)
                if total > limit:
                    raise HTTPException(status_code=413, detail="Diarization input exceeds the configured size limit")
                handle.write(chunk)
    except HTTPException:
        raise
    except OSError as exc:
        raise HTTPException(status_code=500, detail="Unable to store diarization input") from exc
    if total == 0:
        raise HTTPException(status_code=400, detail="Diarization input is empty")
    return total


async def _probe_duration(source: Path) -> float:
    binary = os.getenv("FFPROBE_BINARY", "ffprobe")
    command = [binary, "-v", "error", "-show_entries", "format=duration", "-of", "json", str(source)]
    try:
        process = await asyncio.create_subprocess_exec(
            *command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.DEVNULL
        )
    except OSError as exc:
        raise HTTPException(status_code=503, detail="Diarization media probe is unavailable") from exc
    try:
        stdout, _ = await asyncio.wait_for(process.communicate(), timeout=30)
    except asyncio.TimeoutError as exc:
        process.kill()
        await process.communicate()
        raise HTTPException(status_code=422, detail="Diarization media probe timed out") from exc
    if process.returncode != 0:
        raise HTTPException(status_code=422, detail="Unable to inspect diarization media")
    try:
        payload = json.loads(stdout)
        duration = float(payload["format"]["duration"])
    except (json.JSONDecodeError, KeyError, TypeError, ValueError) as exc:
        raise HTTPException(status_code=422, detail="Diarization media has no valid duration") from exc
    if not math.isfinite(duration) or duration <= 0:
        raise HTTPException(status_code=422, detail="Diarization media has invalid duration")
    if duration > _max_duration_seconds():
        raise HTTPException(status_code=413, detail="Diarization media exceeds the configured duration limit")
    return duration


async def _extract_audio(source: Path, destination: Path) -> None:
    command = [
        os.getenv("FFMPEG_BINARY", "ffmpeg"), "-hide_banner", "-loglevel", "error", "-nostdin",
        "-i", str(source), "-vn", "-ac", "1", "-ar", "16000", "-c:a", "pcm_s16le", "-f", "wav",
        str(destination),
    ]
    try:
        process = await asyncio.create_subprocess_exec(
            *command, stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.PIPE
        )
    except OSError as exc:
        raise HTTPException(status_code=422, detail="Unable to decode diarization media") from exc
    try:
        _, _ = await asyncio.wait_for(process.communicate(), timeout=600)
    except asyncio.TimeoutError as exc:
        process.kill()
        await process.communicate()
        raise HTTPException(status_code=422, detail="Diarization media decoding timed out") from exc
    if process.returncode != 0 or not destination.is_file() or destination.stat().st_size == 0:
        raise HTTPException(status_code=422, detail="Unable to decode diarization media")


def _normalize_speaker(value: Any) -> str:
    if value is None:
        raise HTTPException(status_code=422, detail="Diarization returned an invalid speaker label")
    speaker = str(value).strip()
    if not speaker or len(speaker) > 128:
        raise HTTPException(status_code=422, detail="Diarization returned an invalid speaker label")
    if any(ord(char) < 32 or ord(char) == 127 for char in speaker):
        raise HTTPException(status_code=422, detail="Diarization returned an invalid speaker label")
    return speaker


def _normalize_spans(
    raw_spans: list[dict[str, Any]],
    duration: float,
    max_spans: int,
    max_speakers: int,
) -> list[dict[str, Any]]:
    if not math.isfinite(duration) or duration <= 0:
        raise HTTPException(status_code=422, detail="Diarization media duration is invalid")
    normalized: list[dict[str, Any]] = []
    speakers: set[str] = set()
    for raw in raw_spans:
        if len(normalized) >= max_spans:
            raise HTTPException(status_code=422, detail="Diarization produced too many spans")
        try:
            start = float(raw["start"])
            end = float(raw["end"])
        except (KeyError, TypeError, ValueError) as exc:
            raise HTTPException(status_code=422, detail="Diarization returned an invalid span") from exc
        if not math.isfinite(start) or not math.isfinite(end):
            raise HTTPException(status_code=422, detail="Diarization returned non-finite timestamps")
        if start < 0 or end <= start:
            raise HTTPException(status_code=422, detail="Diarization returned an invalid timestamp range")
        if end > duration + 0.1:
            raise HTTPException(status_code=422, detail="Diarization span exceeds media duration")
        end = min(end, duration)
        speaker = _normalize_speaker(raw.get("speaker"))
        speakers.add(speaker)
        if len(speakers) > max_speakers:
            raise HTTPException(status_code=422, detail="Diarization returned too many speakers")
        normalized.append({"speaker": speaker, "start": start, "end": end})
    normalized.sort(key=lambda item: (item["start"], item["end"], item["speaker"]))
    return normalized


def _run_pipeline(
    pipeline: Any,
    audio_path: Path,
    min_speakers: int | None,
    max_speakers: int | None,
    duration: float,
    max_spans: int,
    max_speakers_allowed: int,
) -> list[dict[str, Any]]:
    kwargs: dict[str, int] = {}
    if min_speakers is not None:
        kwargs["min_speakers"] = min_speakers
    if max_speakers is not None:
        kwargs["max_speakers"] = max_speakers

    diarization = pipeline(str(audio_path), **kwargs)
    raw_spans: list[dict[str, Any]] = []
    for turn, _, speaker in diarization.itertracks(yield_label=True):
        try:
            start = float(turn.start)
            end = float(turn.end)
        except (AttributeError, TypeError, ValueError) as exc:
            raise HTTPException(status_code=422, detail="Diarization returned an invalid span") from exc
        if len(raw_spans) >= max_spans:
            raise HTTPException(status_code=422, detail="Diarization produced too many spans")
        raw_spans.append({"speaker": speaker, "start": start, "end": end})
    return _normalize_spans(raw_spans, duration, max_spans, max_speakers_allowed)


async def _run_inference(
    pipeline: Any,
    audio_path: Path,
    min_speakers: int | None,
    max_speakers: int | None,
    duration: float,
    max_spans: int,
    max_speakers_allowed: int,
) -> list[dict[str, Any]]:
    timeout = _inference_timeout_seconds()
    async with _inference_lock:
        try:
            return await asyncio.wait_for(
                asyncio.to_thread(
                    _run_pipeline, pipeline, audio_path, min_speakers, max_speakers,
                    duration, max_spans, max_speakers_allowed,
                ),
                timeout=timeout,
            )
        except asyncio.TimeoutError as exc:
            raise HTTPException(status_code=504, detail="Speaker diarization timed out") from exc
        except HTTPException:
            raise
        except Exception as exc:
            raise HTTPException(status_code=422, detail="Speaker diarization failed") from exc


@router.post("/diarize")
async def diarize(request: Request) -> dict[str, Any]:
    _authorize(request)
    filename = request.headers.get("x-filename", "input.bin")[:255]
    try:
        min_speakers = int(request.headers["x-min-speakers"]) if request.headers.get("x-min-speakers") else None
        max_speakers = int(request.headers["x-max-speakers"]) if request.headers.get("x-max-speakers") else None
    except ValueError as exc:
        raise HTTPException(status_code=400, detail="Speaker bounds must be integers") from exc

    if min_speakers is not None and not 1 <= min_speakers <= 32:
        raise HTTPException(status_code=400, detail="X-Min-Speakers must be between 1 and 32")
    if max_speakers is not None and not 1 <= max_speakers <= 32:
        raise HTTPException(status_code=400, detail="X-Max-Speakers must be between 1 and 32")
    if min_speakers is not None and max_speakers is not None and min_speakers > max_speakers:
        raise HTTPException(status_code=400, detail="X-Min-Speakers cannot exceed X-Max-Speakers")

    max_spans = _max_response_spans()
    max_speakers_allowed = _max_speakers()

    with tempfile.TemporaryDirectory(prefix="diarization-") as directory:
        root = Path(directory)
        source = root / "input"
        audio = root / "audio.wav"
        await _save_body(request, source)
        duration = await _probe_duration(source)
        await _extract_audio(source, audio)
        pipeline = await _get_pipeline()
        spans = await _run_inference(
            pipeline, audio, min_speakers, max_speakers, duration, max_spans, max_speakers_allowed
        )

    return {
        "spans": spans,
        "speakers": list(dict.fromkeys(item["speaker"] for item in spans)),
        "filename": filename,
    }