from __future__ import annotations

import base64
import io
import json
import os
from typing import Optional

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from safetensors.torch import load_file
from transformers import Wav2Vec2Model, Wav2Vec2Config, pipeline

app = FastAPI(title="Speech Emotion Recognition")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── v1: Custom Wav2Vec2 + SERHead (7 classes) ──────────────────────────
MODEL_CLASSES_V1 = ["angry", "disgust", "fear", "happy", "neutral", "pleasant_surprise", "sad"]

# ── v2: HuggingFace pipeline (8 classes) ────────────────────────────────
V2_LABELS = ["ANG", "CAL", "DIS", "FEA", "HAP", "NEU", "SAD", "SUR"]

_model_v1 = None
_model_v2 = None


class SERHead(nn.Module):
    def __init__(self):
        super().__init__()
        self.projector = nn.Linear(768, 256)
        self.classifier = nn.Linear(256, 7)
        self.layer_weights = nn.Parameter(torch.ones(13) / 13)

    def forward(self, hidden_states):
        stacked = torch.stack(list(hidden_states), dim=0)
        w = F.softmax(self.layer_weights, dim=0)
        weighted = (stacked * w.view(-1, 1, 1, 1)).sum(dim=0)
        pooled = weighted.mean(dim=1)
        return self.classifier(F.relu(self.projector(pooled)))


def load_model_v1():
    global _model_v1
    device = "cuda" if torch.cuda.is_available() else "cpu"
    config = Wav2Vec2Config.from_pretrained("facebook/wav2vec2-base")
    backbone = Wav2Vec2Model(config).to(device).eval()
    head = SERHead().to(device).eval()

    model_path = "model.safetensors"
    if not os.path.exists(model_path):
        print("[WARN] model.safetensors not found — v1 unavailable")
        _model_v1 = "unavailable"
        return

    state = load_file(model_path)
    backbone_prefix = "wav2vec2."
    backbone_state = {k[len(backbone_prefix):]: v for k, v in state.items() if k.startswith(backbone_prefix)}
    backbone.load_state_dict(backbone_state, strict=False)
    head.load_state_dict(state, strict=False)
    _model_v1 = {"backbone": backbone, "head": head, "device": device}
    print("[INFO] v1 (Wav2Vec2) loaded")


def load_model_v2():
    global _model_v2
    try:
        _model_v2 = pipeline(
            "audio-classification",
            model="prithivMLmods/Speech-Emotion-Classification",
        )
        print("[INFO] v2 (prithivMLmods) loaded")
    except Exception as e:
        print(f"[WARN] v2 failed to load: {e}")
        _model_v2 = "unavailable"


@app.on_event("startup")
async def startup():
    load_model_v1()
    load_model_v2()


def _decode_audio(wav_bytes: bytes):
    import soundfile as sf
    import librosa
    buf = io.BytesIO(wav_bytes)
    waveform, sr = sf.read(buf)
    if waveform.ndim > 1:
        waveform = waveform.mean(axis=1)
    if sr != 16000:
        waveform = librosa.resample(y=waveform, orig_sr=sr, target_sr=16000)
        sr = 16000
    return waveform, sr


def _predict_v1(waveform: np.ndarray) -> dict:
    if _model_v1 is None or _model_v1 == "unavailable":
        return {"emotion": "neutral", "confidence": 0.0, "probabilities": {}}

    backbone = _model_v1["backbone"]
    head = _model_v1["head"]
    device = _model_v1["device"]

    wav_t = torch.from_numpy(waveform).float().unsqueeze(0).to(device)
    with torch.no_grad():
        outputs = backbone(wav_t, output_hidden_states=True)
        logits = head(outputs.hidden_states)
        probs = F.softmax(logits, dim=-1).squeeze(0)

    probs_np = probs.cpu().numpy()
    pred_idx = int(probs_np.argmax())
    emotion = MODEL_CLASSES_V1[pred_idx]
    prob_map = {c: round(float(probs_np[i]), 4) for i, c in enumerate(MODEL_CLASSES_V1)}
    return {"emotion": emotion, "confidence": round(float(probs_np[pred_idx]), 4), "probabilities": prob_map}


def _predict_v2(waveform: np.ndarray, sr: int) -> dict:
    if _model_v2 is None or _model_v2 == "unavailable":
        return {"emotion": "neutral", "confidence": 0.0, "probabilities": {}}

    result = _model_v2(waveform, top_k=8)
    probs = {r["label"]: r["score"] for r in result}
    top = result[0]
    return {
        "emotion": top["label"],
        "confidence": round(float(top["score"]), 4),
        "probabilities": probs,
    }


@app.get("/")
@app.get("/health")
async def health():
    return {"status": "ok", "v1_loaded": _model_v1 is not None and _model_v1 != "unavailable", "v2_loaded": _model_v2 is not None and _model_v2 != "unavailable"}


@app.post("/predict_b64")
async def predict_b64(request: Request):
    try:
        body = await request.body()
        content_type = request.headers.get("content-type", "")

        if "application/json" in content_type or body.startswith(b"{"):
            payload = json.loads(body)
            b64_str = payload.get("audio") or payload.get("image", "")
            model_ver = payload.get("model", "v1")
        else:
            import urllib.parse
            parsed = urllib.parse.parse_qs(body.decode())
            raw = parsed.get("data", [None])[0]
            if raw is None:
                raise HTTPException(status_code=400, detail="Missing 'data' field")
            payload = json.loads(raw)
            b64_str = payload.get("audio") or payload.get("image", "") or raw
            model_ver = payload.get("model", "v1")

        if not b64_str:
            raise HTTPException(status_code=400, detail="No audio data found")

        wav_bytes = base64.b64decode(b64_str)
        waveform, sr = _decode_audio(wav_bytes)

        if model_ver == "v2":
            result = _predict_v2(waveform, sr)
        else:
            result = _predict_v1(waveform)

        result["model"] = model_ver
        return result
    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))