"""
Hugging Face Inference API handler for MultimodalFusionModel.

Handles POST requests with:
  - "text": patient speech transcript (required)
  - "audio": base64-encoded WAV audio, 16kHz mono float32 (optional)

Returns:
  { "repetition": 0.xx, "confusion": 0.xx, "distress": 0.xx, "topic_switch": 0.xx }
"""
import os, io, base64, struct, json
import numpy as np
import torch

CLASSES = ["repetition", "confusion", "distress", "topic_switch"]

_model = None
_device = None


def load_model():
    """Load MultimodalFusionModel from the repo's best_model.pt."""
    global _model, _device
    if _model is not None:
        return _model

    # Import model class from repo (model.py is in the same directory)
    import sys
    sys.path.insert(0, os.path.dirname(__file__))
    from model import MultimodalFusionModel

    _device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = MultimodalFusionModel(
        audio_emb_dim=768,
        text_emb_dim=768,
        fusion_hidden=256,
        lstm_layers=2,
        num_heads=4,
        num_classes=4,
        freeze_encoders=True,
    )

    ckpt_path = os.path.join(os.path.dirname(__file__), "best_model.pt")
    state = torch.load(ckpt_path, map_location=_device, weights_only=False)

    # Handle different checkpoint formats
    if isinstance(state, dict) and "model_state_dict" in state:
        state = state["model_state_dict"]
    elif isinstance(state, dict) and "state_dict" in state:
        state = state["state_dict"]

    # Strip "module." prefix (DataParallel / DistributedDataParallel)
    cleaned = {}
    for k, v in state.items():
        cleaned[k.replace("module.", "", 1) if k.startswith("module.") else k] = v

    # Load only matching keys (skip encoder weights if not in checkpoint)
    model_state = model.state_dict()
    matched = {k: v for k, v in cleaned.items() if k in model_state and v.shape == model_state[k].shape}
    model_state.update(matched)
    model.load_state_dict(model_state)

    model.eval().to(_device)
    _model = model
    return _model


def _decode_audio_b64(b64_str):
    """Decode base64 WAV (16kHz mono) to float32 numpy array."""
    raw = base64.b64decode(b64_str)
    # Parse WAV header: skip 44 bytes for standard PCM WAV
    if raw[:4] == b"RIFF":
        # Find "data" chunk
        pos = 12
        while pos < len(raw) - 8:
            chunk_id = raw[pos:pos+4]
            chunk_size = struct.unpack("<I", raw[pos+4:pos+8])[0]
            if chunk_id == b"data":
                audio_data = raw[pos+8:pos+8+chunk_size]
                break
            pos += 8 + chunk_size
        else:
            audio_data = raw[44:]  # fallback
    else:
        audio_data = raw  # raw PCM

    # Try to parse as float32 PCM (model expects float waveforms)
    samples = np.frombuffer(audio_data, dtype=np.float32)
    if samples.shape[0] == 0:
        # Try int16 → float32
        samples = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0
    return samples


def predict(data):
    """
    Run inference.

    Parameters:
        data: dict with keys:
            - "text" (str, required): transcribed speech
            - "audio" (str, optional): base64-encoded 16kHz mono float32 WAV
    Returns:
        dict: { class_name: probability float }
    """
    model = load_model()
    text = data.get("text", "")
    audio_b64 = data.get("audio", None)

    # Build waveform
    if audio_b64:
        waveform = _decode_audio_b64(audio_b64)
        waveform = torch.tensor(waveform, dtype=torch.float32).unsqueeze(0).to(_device)
    else:
        # No audio — use silence (model still works via text encoder)
        waveform = torch.zeros(1, 16000, dtype=torch.float32).to(_device)

    # Tokenize text
    enc = model.tokenizer(
        text, truncation=True, padding=True,
        max_length=128, return_tensors="pt"
    )
    input_ids = enc["input_ids"].to(_device)
    attention_mask = enc["attention_mask"].to(_device)

    # Compute embeddings → forward → sigmoid
    with torch.no_grad():
        audio_emb, text_emb = model.get_embeddings(waveform, input_ids, attention_mask)
        # Add turn dimension: (1, 1, emb_dim)
        audio_emb = audio_emb.unsqueeze(1)
        text_emb = text_emb.unsqueeze(1)
        logits = model.forward(audio_embeds=audio_emb, text_embeds=text_emb)
        probs = torch.sigmoid(logits).squeeze(0).cpu().numpy()

    return {CLASSES[i]: round(float(probs[i]), 4) for i in range(len(CLASSES))}