from io import BytesIO
from contextlib import asynccontextmanager

import librosa
import numpy as np
import torch
from fastapi import FastAPI, File, HTTPException, UploadFile
from transformers import AutoFeatureExtractor, AutoModelForAudioClassification, AutoProcessor

MODEL_NAME = "garystafford/wav2vec2-deepfake-voice-detector"
TARGET_SAMPLE_RATE = 16000
THRESHOLD = 0.5
DEVICE = "cpu"
ALLOWED_EXTENSIONS = {".wav", ".mp3"}


@asynccontextmanager
async def lifespan(app: FastAPI):
    global processor, model
    try:
        processor = AutoProcessor.from_pretrained(MODEL_NAME)
    except TypeError:
        # Some audio checkpoints do not ship tokenizer files expected by AutoProcessor.
        processor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)
    model = AutoModelForAudioClassification.from_pretrained(MODEL_NAME)
    model.to(DEVICE)
    model.eval()
    yield


app = FastAPI(lifespan=lifespan)

processor = None
model = None


def _resolve_label_indices(id2label: dict[int, str]) -> tuple[int, int]:
    real_idx = None
    fake_idx = None

    for idx, label in id2label.items():
        normalized = label.strip().lower()
        if normalized == "real":
            real_idx = idx
        elif normalized == "fake":
            fake_idx = idx

    if real_idx is None or fake_idx is None:
        raise RuntimeError("Model labels must include both 'real' and 'fake'.")

    return real_idx, fake_idx


@app.get("/health")
def health() -> dict[str, str]:
    return {"status": "ok"}


@app.post("/predict")
async def predict(file: UploadFile = File(...)) -> dict:
    if processor is None or model is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet.")

    filename = file.filename or "uploaded_audio.wav"
    lowered = filename.lower()
    if not any(lowered.endswith(ext) for ext in ALLOWED_EXTENSIONS):
        raise HTTPException(status_code=400, detail="Only .wav and .mp3 files are supported.")

    try:
        audio_bytes = await file.read()
        if not audio_bytes:
            raise ValueError("Uploaded file is empty.")

        # librosa loads and converts to mono/16kHz in one step.
        audio, _ = librosa.load(BytesIO(audio_bytes), sr=TARGET_SAMPLE_RATE, mono=True)
        if audio.size == 0:
            raise ValueError("Audio content is empty.")

        audio = np.asarray(audio, dtype=np.float32)

        inputs = processor(
            audio,
            sampling_rate=TARGET_SAMPLE_RATE,
            return_tensors="pt",
            padding=True,
        )

        inputs = {key: value.to(DEVICE) for key, value in inputs.items()}

        with torch.no_grad():
            logits = model(**inputs).logits
            probabilities = torch.softmax(logits, dim=-1).cpu().numpy()[0]

        id2label = {int(k): v for k, v in model.config.id2label.items()}
        real_idx, fake_idx = _resolve_label_indices(id2label)

        real_score = float(probabilities[real_idx])
        fake_score = float(probabilities[fake_idx])

        is_fake = fake_score > THRESHOLD
        predicted_label = "fake" if is_fake else "real"
        predicted_index = fake_idx if is_fake else real_idx
        confidence = fake_score if is_fake else real_score

        return {
            "source": filename,
            "predicted_label": predicted_label,
            "predicted_index": int(predicted_index),
            "confidence": round(confidence, 6),
            "is_fake": is_fake,
            "fake_score": round(fake_score, 6),
            "real_score": round(real_score, 6),
            "threshold": THRESHOLD,
            "scores": {
                "real": round(real_score, 6),
                "fake": round(fake_score, 6),
            },
            "model_name": MODEL_NAME,
            "sample_rate": TARGET_SAMPLE_RATE,
            "device": DEVICE,
        }

    except HTTPException:
        raise
    except Exception as exc:
        raise HTTPException(status_code=400, detail=f"Failed to process audio: {exc}") from exc