from __future__ import annotations

import argparse
import tempfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

import uvicorn
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from faster_whisper import WhisperModel

import sys
sys.path.insert(0, "/app")
from asr.config import (
    COMPUTE_TYPE,
    CUSTOM_MODEL_DIR,
    DEFAULT_MODEL,
    DEVICE,
    MODEL_CACHE,
    PYANNOTE_HF_TOKEN,
    ensure_runtime_dirs,
    resolve_custom_model_path,
)

import os
# 컨테이너는 983 유저로 실행되는데 HOME(/app)이 root 소유라, pyannote.audio가 쓰는
# numba/matplotlib/torch가 각자 ~/.cache, ~/.config 아래에 쓰려다 실패한다.
os.environ["HOME"] = "/tmp"  # 컨테이너 기본 HOME=/app은 983 유저가 쓰기 불가 (setdefault로는 덮어쓰기 안 됨)
os.environ.setdefault("NUMBA_CACHE_DIR", "/tmp/numba_cache")

app = FastAPI(title="ASR Faster-Whisper Worker")
_MODEL_CACHE: Dict[str, WhisperModel] = {}
_DIARIZATION_PIPELINE: Any = None


@app.on_event("startup")
def startup() -> None:
    ensure_runtime_dirs()


@app.get("/health")
def health() -> Dict[str, Any]:
    return {
        "status": "ok",
        "device": DEVICE,
        "compute_type": COMPUTE_TYPE,
        "diarization_available": bool(PYANNOTE_HF_TOKEN),
    }


@app.post("/transcribe")
async def transcribe(
    file: UploadFile = File(...),
    model: str = Form(DEFAULT_MODEL),
    custom_model_path: Optional[str] = Form(None),
    language: Optional[str] = Form(None),
    task: str = Form("transcribe"),
    beam_size: int = Form(5),
    temperature: float = Form(0.0),
    word_timestamps: bool = Form(False),
    diarize: bool = Form(False),
    num_speakers: Optional[int] = Form(None),
    min_speakers: Optional[int] = Form(None),
    max_speakers: Optional[int] = Form(None),
    no_repeat_ngram_size: int = Form(0),
    repetition_penalty: float = Form(1.0),
    compression_ratio_threshold: float = Form(2.4),
    log_prob_threshold: float = Form(-1.0),
    no_speech_threshold: float = Form(0.6),
    condition_on_previous_text: bool = Form(True),
) -> Dict[str, Any]:
    ensure_runtime_dirs()
    suffix = Path(file.filename or "upload.bin").suffix or ".bin"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        tmp.write(await file.read())
        tmp_path = Path(tmp.name)

    try:
        resolved_model = resolve_custom_model_path(custom_model_path) or model
        result = _transcribe(
            audio_path=tmp_path,
            model_name=resolved_model,
            language=language or None,
            task=task,
            beam_size=beam_size,
            temperature=temperature,
            word_timestamps=word_timestamps,
            no_repeat_ngram_size=no_repeat_ngram_size,
            repetition_penalty=repetition_penalty,
            compression_ratio_threshold=compression_ratio_threshold,
            log_prob_threshold=log_prob_threshold,
            no_speech_threshold=no_speech_threshold,
            condition_on_previous_text=condition_on_previous_text,
        )
        if diarize:
            result = _apply_diarization(
                audio_path=tmp_path,
                result=result,
                num_speakers=num_speakers,
                min_speakers=min_speakers,
                max_speakers=max_speakers,
            )
        return result
    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=500, detail=f"Worker failure: {exc}") from exc
    finally:
        tmp_path.unlink(missing_ok=True)


def _load_model(model_name: str) -> WhisperModel:
    if model_name not in _MODEL_CACHE:
        _MODEL_CACHE[model_name] = WhisperModel(
            model_name,
            device=DEVICE,
            compute_type=COMPUTE_TYPE,
            download_root=str(MODEL_CACHE),
        )
    return _MODEL_CACHE[model_name]


def _transcribe(
    audio_path: Path,
    model_name: str,
    language: Optional[str],
    task: str,
    beam_size: int,
    temperature: float,
    word_timestamps: bool,
    no_repeat_ngram_size: int,
    repetition_penalty: float,
    compression_ratio_threshold: float,
    log_prob_threshold: float,
    no_speech_threshold: float,
    condition_on_previous_text: bool,
) -> Dict[str, Any]:
    model = _load_model(model_name)

    temperatures: Union[float, List[float]] = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0] if temperature == 0.0 else temperature

    segments_iter, info = model.transcribe(
        str(audio_path),
        language=language,
        task=task,
        beam_size=beam_size,
        temperature=temperatures,
        word_timestamps=word_timestamps,
        no_repeat_ngram_size=no_repeat_ngram_size,
        repetition_penalty=repetition_penalty,
        compression_ratio_threshold=compression_ratio_threshold,
        log_prob_threshold=log_prob_threshold,
        no_speech_threshold=no_speech_threshold,
        condition_on_previous_text=condition_on_previous_text,
    )

    segment_list = []
    full_text = []
    for seg in segments_iter:
        seg_dict: Dict[str, Any] = {
            "id": seg.id,
            "start": round(float(seg.start), 3),
            "end": round(float(seg.end), 3),
            "text": seg.text,
            "avg_logprob": round(float(seg.avg_logprob), 4) if seg.avg_logprob is not None else None,
            "compression_ratio": round(float(seg.compression_ratio), 4) if seg.compression_ratio is not None else None,
            "no_speech_prob": round(float(seg.no_speech_prob), 4) if seg.no_speech_prob is not None else None,
        }
        if word_timestamps and getattr(seg, "words", None):
            seg_dict["words"] = [
                {
                    "start": round(float(w.start), 3),
                    "end": round(float(w.end), 3),
                    "word": w.word,
                    "probability": round(float(w.probability), 4),
                }
                for w in seg.words
            ]
        segment_list.append(seg_dict)
        full_text.append(seg.text.strip())

    duration = None
    try:
        duration = round(float(info.duration), 3)
    except Exception:
        pass

    return {
        "backend": "faster-whisper",
        "model": model_name,
        "language": getattr(info, "language", language),
        "language_probability": round(float(getattr(info, "language_probability", 0) or 0), 4),
        "duration": duration,
        "text": " ".join(x for x in full_text if x),
        "segments": segment_list,
        "diarized": False,
    }


def _get_diarization_pipeline():
    global _DIARIZATION_PIPELINE
    if _DIARIZATION_PIPELINE is None:
        try:
            from pyannote.audio import Pipeline
            import torch
        except ImportError as exc:
            raise RuntimeError("pyannote.audio가 설치되지 않았습니다.") from exc
        if not PYANNOTE_HF_TOKEN:
            raise RuntimeError("화자 분리를 사용하려면 PYANNOTE_HF_TOKEN 환경변수를 설정하세요.")
        _DIARIZATION_PIPELINE = Pipeline.from_pretrained(
            "pyannote/speaker-diarization-3.1",
            use_auth_token=PYANNOTE_HF_TOKEN,
        )
        if DEVICE == "cuda":
            import torch
            _DIARIZATION_PIPELINE = _DIARIZATION_PIPELINE.to(torch.device("cuda"))
    return _DIARIZATION_PIPELINE


def _apply_diarization(
    audio_path: Path,
    result: Dict[str, Any],
    num_speakers: Optional[int],
    min_speakers: Optional[int],
    max_speakers: Optional[int],
) -> Dict[str, Any]:
    pipeline = _get_diarization_pipeline()
    kwargs: Dict[str, Any] = {}
    if num_speakers is not None:
        kwargs["num_speakers"] = num_speakers
    else:
        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)
    segments = result.get("segments", [])
    for seg in segments:
        seg_start = float(seg.get("start", 0))
        seg_end = float(seg.get("end", 0))
        speaker_times: Dict[str, float] = {}
        for turn, _, speaker in diarization.itertracks(yield_label=True):
            overlap_start = max(seg_start, turn.start)
            overlap_end = min(seg_end, turn.end)
            if overlap_end > overlap_start:
                speaker_times[speaker] = speaker_times.get(speaker, 0.0) + (overlap_end - overlap_start)
        seg["speaker"] = max(speaker_times, key=speaker_times.get) if speaker_times else "UNKNOWN"

    lines = []
    current_speaker: Optional[str] = None
    current_texts: list = []
    for seg in segments:
        speaker = seg.get("speaker", "UNKNOWN")
        text = seg.get("text", "").strip()
        if not text:
            continue
        if speaker != current_speaker:
            if current_texts and current_speaker is not None:
                lines.append(f"[{current_speaker}]: {' '.join(current_texts)}")
            current_speaker = speaker
            current_texts = [text]
        else:
            current_texts.append(text)
    if current_texts and current_speaker is not None:
        lines.append(f"[{current_speaker}]: {' '.join(current_texts)}")

    result["text"] = "\n".join(lines)
    result["diarized"] = True
    return result


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8001)
    args = parser.parse_args()
    ensure_runtime_dirs()
    uvicorn.run(app, host=args.host, port=args.port)