"""AASIST3 Audio Deepfake Detection API.

Detects synthetic speech using the AASIST3 model architecture:
- Frontend: XLSR wav2vec 2.0 (HuggingFace Transformers)
- Backend: AASIST with KAN (Kolmogorov-Arnold Network) linear
  layers and Graph Attention Networks

Reference: https://github.com/AI4Bharat/AASIST3
"""

import base64
import io
import logging
import os
import platform
import time
from typing import Optional

import numpy as np
import soundfile as sf
import torch
import torch.nn.functional as F
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field

# Point transformers cache to pre-cached wav2vec2 config
# (must be set before importing model code)
os.environ["TRANSFORMERS_CACHE"] = "/app/w2v_cache"

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("aasist3_api")

# Import model class
try:
    from model import aasist3 as AASIST3Model
except ImportError as e:
    logger.error(f"Failed to import AASIST3 model: {e}")
    AASIST3Model = None

# Constants
MODEL_NAME = "aasist3"
MODEL_ID = "aasist3_kan_mlaad"
WEIGHTS_DIR = "/app/weights"


def _get_device():
    """Select optimal device: MPS (Apple) > CUDA (NVIDIA) > CPU."""
    override = os.environ.get("DEEPSAFE_DEVICE", "").strip().lower()
    if override == "cpu":
        return torch.device("cpu")
    if override == "cuda" and torch.cuda.is_available():
        return torch.device("cuda")
    if (
        override == "mps"
        and hasattr(torch.backends, "mps")
        and torch.backends.mps.is_available()
    ):
        return torch.device("mps")
    if override:
        pass  # Invalid override, fall through to auto-detect
    if (
        platform.system() == "Darwin"
        and hasattr(torch.backends, "mps")
        and torch.backends.mps.is_available()
    ):
        return torch.device("mps")
    if torch.cuda.is_available():
        return torch.device("cuda")
    return torch.device("cpu")


DEVICE = _get_device()

if DEVICE.type == "cuda":
    torch.backends.cudnn.benchmark = True
    torch.set_float32_matmul_precision("high")

if DEVICE.type == "cuda":
    logger.info(
        "Device: cuda (%s, %.1f GB VRAM)",
        torch.cuda.get_device_name(0),
        torch.cuda.get_device_properties(0).total_memory / 1024**3,
    )
else:
    logger.warning(
        "Device: %s (no CUDA available -- check nvidia-container-toolkit)",
        DEVICE,
    )

SAMPLE_RATE = 16000
TARGET_SAMPLES = 64600  # ~4.04 seconds at 16kHz

# Global model instance
model = None


class AudioInput(BaseModel):
    """Request schema for audio deepfake detection."""

    audio_data: str = Field(
        ..., description="Base64 encoded audio string (WAV/MP3/etc)"
    )
    threshold: Optional[float] = Field(
        0.5, ge=0.0, le=1.0, description="Classification threshold"
    )


app = FastAPI(
    title="AASIST3 Audio Deepfake Detection API",
    description=(
        "Service for detecting synthetic speech using the "
        "AASIST3 model (HuggingFace wav2vec 2.0 + AASIST "
        "with KAN layers)."
    ),
    version="1.0.0",
)


def load_model():
    """Load the AASIST3 model from pretrained weights.

    Returns:
        The loaded model, or None if loading fails.
    """
    global model
    if model is not None:
        return model

    logger.info(f"Loading AASIST3 model onto {DEVICE}...")

    if AASIST3Model is None:
        logger.error("AASIST3 model class not available.")
        return None

    weights_safetensors = os.path.join(WEIGHTS_DIR, "model.safetensors")
    weights_config = os.path.join(WEIGHTS_DIR, "config.json")

    if not os.path.exists(weights_safetensors):
        logger.error(f"Model weights not found at {weights_safetensors}")
        return None

    if not os.path.exists(weights_config):
        logger.error(f"Model config not found at {weights_config}")
        return None

    try:
        model = AASIST3Model.from_pretrained(WEIGHTS_DIR)
        model.to(DEVICE)
        # Set model to inference mode (disables dropout, batchnorm)
        model.train(False)

        logger.info("AASIST3 model loaded successfully.")
        return model
    except Exception as e:
        logger.exception(f"Failed to load AASIST3 model: {e}")
        model = None
        return None


@app.on_event("startup")
async def startup_event():
    """Load model on service startup."""
    load_model()


def _gpu_health_info() -> dict:
    """Return GPU metrics for the health endpoint."""
    if torch.cuda.is_available() and DEVICE.type == "cuda":
        return {
            "gpu_name": torch.cuda.get_device_name(0),
            "vram_used_mb": round(torch.cuda.memory_allocated(0) / 1024**2),
            "vram_total_mb": round(
                torch.cuda.get_device_properties(0).total_memory / 1024**2
            ),
        }
    return {}


@app.get("/health")
async def health():
    """Health check endpoint."""
    return {
        "status": "healthy" if model is not None else "degraded",
        "model": MODEL_NAME,
        "model_id": MODEL_ID,
        "device": str(DEVICE),
        "weights_found": os.path.exists(os.path.join(WEIGHTS_DIR, "model.safetensors")),
        **_gpu_health_info(),
    }


def _load_audio_bytes(audio_bytes: bytes) -> tuple:
    """Load audio from raw bytes using soundfile with torchaudio fallback.

    Args:
        audio_bytes: Raw audio file bytes.

    Returns:
        Tuple of (audio_numpy_array, sample_rate).

    Raises:
        ValueError: If audio cannot be loaded by any backend.
    """
    # Try soundfile first (handles WAV, FLAC natively)
    sf_error = None
    try:
        audio, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32")
        if audio.ndim > 1:
            audio = audio.mean(axis=1)  # Convert to mono
        return audio, sr
    except Exception as sf_err:
        sf_error = sf_err
        logger.debug(f"soundfile failed, trying torchaudio: {sf_err}")

    # Fallback to torchaudio (handles MP3, compressed formats)
    try:
        import torchaudio

        buf = io.BytesIO(audio_bytes)
        waveform, sr = torchaudio.load(buf)
        if waveform.shape[0] > 1:
            waveform = waveform.mean(dim=0, keepdim=True)
        return waveform.squeeze(0).numpy(), sr
    except Exception as ta_err:
        raise ValueError(
            f"Failed to load audio with soundfile and torchaudio: "
            f"sf={sf_error}, ta={ta_err}"
        )


def preprocess_audio(audio_bytes: bytes) -> torch.Tensor:
    """Preprocess audio for AASIST3 inference.

    Loads audio, resamples to 16kHz mono, and zero-pads or
    truncates to TARGET_SAMPLES.

    Args:
        audio_bytes: Raw audio file bytes.

    Returns:
        Audio tensor of shape (1, TARGET_SAMPLES).

    Raises:
        ValueError: If audio preprocessing fails.
    """
    try:
        logger.info("Starting audio preprocessing...")
        audio, sr = _load_audio_bytes(audio_bytes)
        logger.info(f"Audio loaded. Length: {len(audio)} samples at {sr}Hz")

        # Resample to 16kHz if needed
        if sr != SAMPLE_RATE:
            import torchaudio

            resampler = torchaudio.transforms.Resample(
                orig_freq=sr, new_freq=SAMPLE_RATE
            )
            audio_tensor = torch.FloatTensor(audio).unsqueeze(0)
            audio_tensor = resampler(audio_tensor).squeeze(0)
            audio = audio_tensor.numpy()
            logger.info(
                f"Resampled from {sr}Hz to {SAMPLE_RATE}Hz. "
                f"New length: {len(audio)} samples"
            )

        # Zero-pad or truncate to TARGET_SAMPLES
        if len(audio) >= TARGET_SAMPLES:
            audio = audio[:TARGET_SAMPLES]
        else:
            pad_length = TARGET_SAMPLES - len(audio)
            audio = np.pad(audio, (0, pad_length), mode="constant")

        logger.info(f"Audio padded/trimmed to {TARGET_SAMPLES} samples")

        audio_tensor = torch.FloatTensor(audio).unsqueeze(0).to(DEVICE)
        return audio_tensor
    except Exception as e:
        logger.error(f"Error preprocessing audio: {e}")
        raise ValueError(f"Audio preprocessing failed: {str(e)}")


@app.post("/predict")
async def predict(input_data: AudioInput):
    """Run deepfake detection on base64-encoded audio.

    The model outputs 2 logits: [bonafide_score, spoof_score].
    Class 0 = bonafide (real), Class 1 = spoof (fake).
    The returned probability is the spoof/fake probability.
    """
    if model is None:
        if load_model() is None:
            raise HTTPException(status_code=503, detail="Model not loaded")

    try:
        start_time = time.time()
        logger.info(
            f"Prediction request. Data size: " f"{len(input_data.audio_data)} chars"
        )

        # Decode base64 audio
        audio_bytes = base64.b64decode(input_data.audio_data)

        # Preprocess
        audio_tensor = preprocess_audio(audio_bytes)

        # Inference
        logger.info("Starting model inference...")
        with torch.no_grad():
            output = model(audio_tensor)

            # output shape: [batch, 2]
            # Index 0 = bonafide logit, Index 1 = spoof logit
            probs = torch.softmax(output, dim=1)
            prob_fake = probs[0, 1].item()

        prediction = 1 if prob_fake >= input_data.threshold else 0
        verdict = "fake" if prediction == 1 else "real"
        inference_time = time.time() - start_time

        logger.info(
            f"Prediction: {verdict} (prob_fake={prob_fake:.4f}, "
            f"time={inference_time:.3f}s)"
        )

        return {
            "model": MODEL_NAME,
            "probability": float(prob_fake),
            "prediction": int(prediction),
            "class": verdict,
            "inference_time": float(inference_time),
        }

    except Exception as e:
        logger.exception(f"Error during prediction: {e}")
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    port = int(os.environ.get("MODEL_PORT", 8005))
    uvicorn.run(app, host="0.0.0.0", port=port)