import os
import tempfile

import librosa
import numpy as np
import torch
from fastapi import FastAPI, File, Header, HTTPException, UploadFile

MODEL_ID = os.environ.get("MODEL_ID", "awsaf49/sonics-spectttra-alpha-120s")
VOICE_MODEL_ID = os.environ.get("VOICE_MODEL_ID", "MattyB95/AST-ASVspoof5-Synthetic-Voice-Detection")
API_KEY = os.environ.get("DETECT_API_KEY", "")
MAX_BYTES = 25 * 1024 * 1024
MODEL_SR = 16000

torch.set_num_threads(2)

app = FastAPI(title="SONICS Detect API")
model = None
voice_model = None
voice_extractor = None


@app.on_event("startup")
def load_model():
    global model, voice_model, voice_extractor
    from sonics import HFAudioClassifier

    m = HFAudioClassifier.from_pretrained(MODEL_ID)
    m.eval()
    model = m

    try:
        from transformers import AutoFeatureExtractor, AutoModelForAudioClassification

        voice_extractor = AutoFeatureExtractor.from_pretrained(VOICE_MODEL_ID)
        vm = AutoModelForAudioClassification.from_pretrained(VOICE_MODEL_ID)
        vm.eval()
        voice_model = vm
    except Exception:
        voice_model = None
        voice_extractor = None


@app.get("/")
def health():
    return {
        "ok": True,
        "model": MODEL_ID,
        "loaded": model is not None,
        "voice_model": VOICE_MODEL_ID,
        "voice_loaded": voice_model is not None,
    }


# Krumhansl-Schmuckler key profiles (major/minor).
KS_MAJOR = np.array([6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88])
KS_MINOR = np.array([6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17])
PITCHES = ["C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B"]


def estimate_key(y: np.ndarray, sr: int) -> dict:
    chroma = librosa.feature.chroma_cqt(y=y, sr=sr).mean(axis=1)
    if chroma.sum() <= 0:
        return {"key": None, "mode": None, "confidence": None}
    scores = []
    for shift in range(12):
        rolled = np.roll(chroma, -shift)
        for mode, profile in (("major", KS_MAJOR), ("minor", KS_MINOR)):
            r = float(np.corrcoef(rolled, profile)[0, 1])
            scores.append((r, PITCHES[shift], mode))
    scores.sort(reverse=True)
    best, second = scores[0], scores[1]
    confidence = max(0.0, min(1.0, (best[0] - second[0]) * 5 + 0.5))
    return {"key": best[1], "mode": best[2], "confidence": round(confidence, 2)}


def load_upload(audio: UploadFile, sr: int | None = None):
    data = audio.file.read()
    if not data:
        raise HTTPException(status_code=400, detail="Empty file")
    if len(data) > MAX_BYTES:
        raise HTTPException(status_code=413, detail="File too large (max 25MB)")
    suffix = os.path.splitext(audio.filename or "")[1] or ".mp3"
    try:
        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
            tmp.write(data)
            tmp_path = tmp.name
        try:
            y, sr_out = librosa.load(tmp_path, sr=sr, mono=True)
        finally:
            os.unlink(tmp_path)
    except HTTPException:
        raise
    except Exception:
        raise HTTPException(status_code=400, detail="Could not decode audio")
    return y, sr_out


@app.post("/analyze")
def analyze(audio: UploadFile = File(...), x_detect_key: str = Header(default="")):
    """Music utilities: BPM + musical key + basic facts. No AI model involved."""
    if API_KEY and x_detect_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid detect key")

    y, sr = load_upload(audio, sr=None)
    if y.size < sr * 3:
        raise HTTPException(status_code=400, detail="Audio too short (min 3s)")

    tempo_bpm = None
    tempo_alt = None
    try:
        tempo, beats = librosa.beat.beat_track(y=y, sr=sr, hop_length=512)
        t = float(np.atleast_1d(tempo)[0]) if tempo is not None else 0.0
        if t > 0:
            tempo_bpm = round(t, 1)
            # Common octave error alternative (half/double time).
            tempo_alt = round(t * 2, 1) if t < 90 else round(t / 2, 1)
    except Exception:
        pass

    key = estimate_key(y, sr)

    return {
        "bpm": tempo_bpm,
        "bpm_alternative": tempo_alt,
        "key": key["key"],
        "mode": key["mode"],
        "key_confidence": key["confidence"],
        "duration_s": round(len(y) / sr, 1),
        "sample_rate": int(sr),
    }


@app.post("/detect-voice")
def detect_voice(audio: UploadFile = File(...), x_detect_key: str = Header(default="")):
    """AI voice / speech deepfake detection (AST fine-tuned on ASVspoof 5)."""
    if API_KEY and x_detect_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid detect key")
    if voice_model is None or voice_extractor is None:
        raise HTTPException(status_code=503, detail="Voice model unavailable")

    y, sr = load_upload(audio, sr=16000)
    if y.size < sr * 2:
        raise HTTPException(status_code=400, detail="Audio too short (min 2s)")

    # Score up to three 10s windows (start / middle / end) and average.
    win = sr * 10
    starts = [0]
    if len(y) > win * 2:
        starts.append((len(y) - win) // 2)
    if len(y) > win:
        starts.append(max(0, len(y) - win))
    probs = []
    id2label = voice_model.config.id2label
    for s in dict.fromkeys(starts):
        chunk = y[s : s + win]
        inputs = voice_extractor(chunk, sampling_rate=sr, return_tensors="pt")
        with torch.no_grad():
            logits = voice_model(**inputs).logits
        p = torch.softmax(logits, dim=-1)[0]
        spoof_idx = next(
            (i for i, lbl in id2label.items() if "spoof" in lbl.lower() or "fake" in lbl.lower()),
            1,
        )
        probs.append(float(p[int(spoof_idx)]))

    ai_prob = float(np.mean(probs))
    return {
        "ai_prob": round(ai_prob, 4),
        "windows": len(probs),
        "duration_s": round(len(y) / sr, 1),
        "model": VOICE_MODEL_ID,
    }


def compute_signals(y: np.ndarray, sr: int) -> dict:
    """Audio-forensic descriptors computed at the file's native sample rate."""
    out: dict = {"native_sample_rate": int(sr)}
    try:
        S = np.abs(librosa.stft(y, n_fft=2048, hop_length=512))
        freqs = librosa.fft_frequencies(sr=sr, n_fft=2048)
        mag = S.mean(axis=1)
        power = mag**2
        total = float(power.sum()) or 1e-12

        # Spectral cutoff: frequency below which 99% of energy lives.
        cum = np.cumsum(power)
        idx = int(np.searchsorted(cum, 0.99 * cum[-1]))
        out["spectral_cutoff_hz"] = int(freqs[min(idx, len(freqs) - 1)])

        # Share of energy above 10 kHz (only meaningful when sr allows it).
        if sr >= 32000:
            out["hf_energy_ratio"] = round(float(power[freqs >= 10000].sum() / total), 5)
        else:
            out["hf_energy_ratio"] = None

        # Dynamic range: dB spread between loud and quiet frames.
        rms = librosa.feature.rms(y=y, hop_length=512)[0]
        rms = rms[rms > 1e-6]
        if rms.size:
            out["dynamic_range_db"] = round(
                float(20 * np.log10(np.percentile(rms, 95) / max(np.percentile(rms, 10), 1e-9))), 1
            )
        else:
            out["dynamic_range_db"] = None

        # Spectral flatness: noisiness/synthetic-ness of the average spectrum.
        out["spectral_flatness"] = round(float(librosa.feature.spectral_flatness(y=y).mean()), 4)

        # Tempo and beat regularity (coefficient of variation of inter-beat intervals).
        try:
            tempo, beats = librosa.beat.beat_track(y=y, sr=sr, hop_length=512)
            tempo_val = float(np.atleast_1d(tempo)[0]) if tempo is not None else 0.0
            out["tempo_bpm"] = round(tempo_val, 1) if tempo_val > 0 else None
            if beats is not None and len(beats) > 8:
                ibis = np.diff(librosa.frames_to_time(beats, sr=sr, hop_length=512))
                out["tempo_cv"] = round(float(np.std(ibis) / max(np.mean(ibis), 1e-9)), 4)
            else:
                out["tempo_cv"] = None
        except Exception:
            out["tempo_bpm"] = None
            out["tempo_cv"] = None
    except Exception:
        pass
    return out


@app.post("/detect")
def detect(audio: UploadFile = File(...), x_detect_key: str = Header(default="")):
    if API_KEY and x_detect_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid detect key")
    if model is None:
        raise HTTPException(status_code=503, detail="Model still loading, retry shortly")

    data = audio.file.read()
    if not data:
        raise HTTPException(status_code=400, detail="Empty file")
    if len(data) > MAX_BYTES:
        raise HTTPException(status_code=413, detail="File too large (max 25MB)")

    suffix = os.path.splitext(audio.filename or "")[1] or ".mp3"
    try:
        with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
            tmp.write(data)
            tmp_path = tmp.name
        try:
            # Native rate for forensic signals (high-frequency artifacts live here).
            y_native, sr_native = librosa.load(tmp_path, sr=None, mono=True)
        finally:
            os.unlink(tmp_path)
    except Exception:
        raise HTTPException(status_code=400, detail="Could not decode audio")

    if y_native.size < sr_native * 3:
        raise HTTPException(status_code=400, detail="Audio too short (min 3s)")

    signals = compute_signals(y_native, sr_native)

    # Model expects 16 kHz; score the middle max_time window (official demo logic).
    y = (
        librosa.resample(y_native, orig_sr=sr_native, target_sr=MODEL_SR)
        if sr_native != MODEL_SR
        else y_native
    )
    max_time = model.config.audio.max_time
    chunk_samples = int(max_time * MODEL_SR)
    total_chunks = len(y) // chunk_samples
    middle_idx = total_chunks // 2
    start = middle_idx * chunk_samples
    chunk = y[start : start + chunk_samples]
    if len(chunk) < chunk_samples:
        chunk = np.pad(chunk, (0, chunk_samples - len(chunk)))

    with torch.no_grad():
        t = torch.from_numpy(chunk).float().unsqueeze(0)
        ai_prob = float(torch.sigmoid(model(t)).cpu().numpy().reshape(-1)[0])

    return {
        "ai_prob": round(ai_prob, 4),
        "duration_s": round(len(y_native) / sr_native, 1),
        "model": MODEL_ID,
        "signals": signals,
    }