"""Kicks API — VAE kick drum synthesizer deployed on Hugging Face Spaces.

Uses BigVGAN vocoder on GPU (falls back to Griffin-LIM on CPU).
PCA parameters are pre-computed and loaded from pca_params.npz.
"""

import glob
import io
import os
import time
from collections import defaultdict

import numpy as np
import torch
import torch.nn as nn
import torchaudio
import soundfile as sf
from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse


# ──────────────────────────────────────────────
# Rate limiter
# ──────────────────────────────────────────────
class RateLimiter:
    """Simple in-memory sliding-window rate limiter."""

    def __init__(self, max_requests: int = 10, window_seconds: float = 60.0):
        self.max_requests = max_requests
        self.window = window_seconds
        self._clients: dict[str, list[float]] = defaultdict(list)

    def _prune(self, ip: str, now: float) -> None:
        cutoff = now - self.window
        self._clients[ip] = [t for t in self._clients[ip] if t > cutoff]
        if not self._clients[ip]:
            del self._clients[ip]

    def is_allowed(self, ip: str) -> bool:
        now = time.time()
        self._prune(ip, now)
        if len(self._clients[ip]) >= self.max_requests:
            return False
        self._clients[ip].append(now)
        return True


def _get_client_ip(request: Request) -> str:
    forwarded = request.headers.get("X-Forwarded-For")
    if forwarded:
        return forwarded.split(",")[0].strip()
    return request.client.host if request.client else "unknown"


# Per-endpoint rate limiters (stricter for expensive endpoints)
_generate_limiter = RateLimiter(max_requests=10, window_seconds=60)
_spectrogram_limiter = RateLimiter(max_requests=30, window_seconds=60)
_config_limiter = RateLimiter(max_requests=60, window_seconds=60)


async def rate_limit_generate(request: Request) -> None:
    ip = _get_client_ip(request)
    if not _generate_limiter.is_allowed(ip):
        raise HTTPException(status_code=429, detail="Too many requests. Wait before generating again.")


async def rate_limit_spectrogram(request: Request) -> None:
    ip = _get_client_ip(request)
    if not _spectrogram_limiter.is_allowed(ip):
        raise HTTPException(status_code=429, detail="Too many requests. Wait before requesting again.")

# ──────────────────────────────────────────────
# Constants (mirrors kicks/model.py)
# ──────────────────────────────────────────────
SAMPLE_RATE = 44100
AUDIO_LENGTH = 65536      # ~1.49 seconds
N_FFT = 1024
HOP_LENGTH = 256
WIN_SIZE = 1024
N_MELS = 128
FMIN = 0
FMAX = None

LOG_MEL_MIN = -11.5129    # ln(1e-5) — silence floor
LOG_MEL_MAX = 2.5         # headroom above observed max

N_PCS = 5
LATENT_DIM = 128

VOCODER_DIR = "models/vocoder"  # fine-tuned BigVGAN weights
BIGVGAN_MODEL = "nvidia/bigvgan_v2_44khz_128band_256x"


# ──────────────────────────────────────────────
# VAE model (identical architecture to kicks/model.py)
# ──────────────────────────────────────────────
class VAE(nn.Module):
    """2D Conv VAE for log-mel spectrograms (1, 128, 256)."""

    def __init__(self, latent_dim: int = 128):
        super().__init__()
        self.latent_dim = latent_dim

        self.encoder = nn.Sequential(
            nn.Conv2d(1, 32, 3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU(),
            nn.Conv2d(32, 64, 3, stride=2, padding=1), nn.BatchNorm2d(64), nn.ReLU(),
            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.BatchNorm2d(128), nn.ReLU(),
            nn.Conv2d(128, 256, 3, stride=2, padding=1), nn.BatchNorm2d(256), nn.ReLU(),
        )
        self.fc_mu = nn.Linear(256 * 8 * 16, latent_dim)
        self.fc_logvar = nn.Linear(256 * 8 * 16, latent_dim)
        self.fc_decode = nn.Linear(latent_dim, 256 * 8 * 16)
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(256, 128, 3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(128), nn.ReLU(),
            nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(64), nn.ReLU(),
            nn.ConvTranspose2d(64, 32, 3, stride=2, padding=1, output_padding=1),
            nn.BatchNorm2d(32), nn.ReLU(),
            nn.ConvTranspose2d(32, 1, 3, stride=2, padding=1, output_padding=1),
            nn.Sigmoid(),
        )

    def decode(self, z: torch.Tensor) -> torch.Tensor:
        h = self.fc_decode(z)
        h = h.view(-1, 256, 8, 16)
        return self.decoder(h)


# ──────────────────────────────────────────────
# BigVGAN vocoder
# ──────────────────────────────────────────────
def load_bigvgan(device: torch.device):
    """Load BigVGAN from HF Hub, optionally with fine-tuned weights."""
    import bigvgan as _bigvgan

    # Patch for huggingface_hub >= 1.0 compatibility
    original = _bigvgan.BigVGAN._from_pretrained.__func__

    @classmethod
    def _patched(cls, **kwargs):
        kwargs.setdefault("proxies", None)
        kwargs.setdefault("resume_download", False)
        return original(cls, **kwargs)

    _bigvgan.BigVGAN._from_pretrained = _patched

    model = _bigvgan.BigVGAN.from_pretrained(BIGVGAN_MODEL, use_cuda_kernel=False)

    # Load fine-tuned weights if available
    pth_files = sorted(glob.glob(os.path.join(VOCODER_DIR, "*.pth")))
    if pth_files:
        pth_path = pth_files[0]
        checkpoint = torch.load(pth_path, map_location=device, weights_only=True)
        model.load_state_dict(checkpoint["generator"])
        print(f"Loaded fine-tuned BigVGAN from {pth_path}")
    else:
        print("Using pretrained BigVGAN from HuggingFace Hub")

    model.remove_weight_norm()
    return model.eval().to(device)


# ──────────────────────────────────────────────
# Griffin-LIM fallback vocoder
# ──────────────────────────────────────────────
class GriffinLimVocoder:
    """Mel spectrogram → audio via pseudo-inverse + Griffin-LIM."""

    def __init__(self, n_iter: int = 64):
        n_stft = N_FFT // 2 + 1
        fb = torchaudio.functional.melscale_fbanks(
            n_freqs=n_stft,
            f_min=FMIN,
            f_max=FMAX if FMAX is not None else SAMPLE_RATE / 2.0,
            n_mels=N_MELS,
            sample_rate=SAMPLE_RATE,
        )
        self.fb_pinv = torch.linalg.pinv(fb.T)
        self.griffin_lim = torchaudio.transforms.GriffinLim(
            n_fft=N_FFT, win_length=WIN_SIZE,
            hop_length=HOP_LENGTH, n_iter=n_iter,
        )

    def __call__(self, log_mel: torch.Tensor) -> torch.Tensor:
        mel_linear = torch.exp(log_mel).cpu()
        linear_spec = torch.clamp(self.fb_pinv @ mel_linear, min=0.0)
        linear_spec = linear_spec + 1e-4
        return self.griffin_lim(linear_spec)


# ──────────────────────────────────────────────
# Normalization helpers
# ──────────────────────────────────────────────
def denormalize(spec: torch.Tensor) -> torch.Tensor:
    """Convert [0, 1] normalized spectrogram back to log-mel scale."""
    return spec * (LOG_MEL_MAX - LOG_MEL_MIN) + LOG_MEL_MIN


# ──────────────────────────────────────────────
# Application state
# ──────────────────────────────────────────────
class AppState:
    device: torch.device
    model: VAE
    vocoder: object
    vocoder_type: str
    pca_components: np.ndarray
    pca_mean: np.ndarray
    pc_names: list[str]
    pc_mins: list[float]
    pc_maxs: list[float]


state = AppState()


# ──────────────────────────────────────────────
# FastAPI app
# ──────────────────────────────────────────────
app = FastAPI(title="Kicks API")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.on_event("startup")
async def startup():
    # Auto-detect device
    if torch.cuda.is_available():
        state.device = torch.device("cuda")
        print("Using CUDA GPU")
    else:
        state.device = torch.device("cpu")
        print("CUDA not available — falling back to CPU")

    # Load PCA parameters
    pca_path = os.environ.get("PCA_PARAMS_PATH", "pca_params.npz")
    pca_data = np.load(pca_path, allow_pickle=False)
    state.pca_components = pca_data["components"]
    state.pca_mean = pca_data["mean"]
    state.pc_names = list(pca_data["pc_names"])
    state.pc_mins = list(pca_data["pc_mins"])
    state.pc_maxs = list(pca_data["pc_maxs"])
    print(f"PCA loaded: {state.pc_names}")

    # Load VAE checkpoint
    ckpt_path = os.environ.get("VAE_CHECKPOINT_PATH", "vae_best.pth")
    checkpoint = torch.load(ckpt_path, map_location=state.device, weights_only=True)
    latent_dim = checkpoint.get("latent_dim", LATENT_DIM)
    state.model = VAE(latent_dim=latent_dim)
    state.model.load_state_dict(checkpoint["model"])
    state.model.to(state.device)
    state.model.eval()
    print(f"VAE loaded (latent_dim={latent_dim})")

    # Load vocoder — try BigVGAN with CUDA, fall back to Griffin-LIM
    if state.device.type == "cuda":
        try:
            state.vocoder = load_bigvgan(state.device)
            state.vocoder_type = "bigvgan"
            print("BigVGAN vocoder ready")
        except Exception as e:
            print(f"BigVGAN loading failed ({e}), falling back to Griffin-LIM")
            state.vocoder = GriffinLimVocoder()
            state.vocoder_type = "griffinlim"
    else:
        state.vocoder = GriffinLimVocoder()
        state.vocoder_type = "griffinlim"

    print("API ready")


def _parse_pc_values(request: Request) -> list[float]:
    values = []
    for i in range(N_PCS):
        raw = float(request.query_params.get(f"pc{i + 1}", "0.5"))
        val = state.pc_mins[i] + raw * (state.pc_maxs[i] - state.pc_mins[i])
        values.append(max(state.pc_mins[i], min(state.pc_maxs[i], val)))
    return values


def _pc_to_latent(pc_values: list[float]) -> torch.Tensor:
    z_np = (np.array(pc_values) @ state.pca_components) + state.pca_mean
    return torch.tensor(z_np, dtype=torch.float32).to(state.device).unsqueeze(0)


def _spec_to_audio(spec: torch.Tensor) -> np.ndarray:
    """Convert normalized spectrogram to waveform using the loaded vocoder."""
    log_mel = denormalize(spec)
    log_mel = log_mel.squeeze(1)  # (1, 128, 256)

    if state.vocoder_type == "bigvgan":
        with torch.no_grad():
            waveform = state.vocoder(log_mel.to(state.device))  # (1, 1, T)
        waveform = waveform.squeeze(1).cpu()  # (1, T)
        waveform = torchaudio.functional.highpass_biquad(waveform, SAMPLE_RATE, cutoff_freq=25.0)
        waveform = torchaudio.functional.lowpass_biquad(waveform, SAMPLE_RATE, cutoff_freq=20000.0)
        waveform = waveform / (waveform.abs().max() + 1e-8)
    else:
        waveform = state.vocoder(log_mel)  # (1, T)
        waveform = waveform / (waveform.abs().max() + 1e-8)

    return waveform.squeeze(0).numpy()


# ──────────────────────────────────────────────
# Endpoints
# ──────────────────────────────────────────────

@app.get("/config")
async def get_config():
    return {
        "sliders": [
            {
                "id": i + 1,
                "name": state.pc_names[i],
                "min": 0.0, "max": 1.0,
                "default": 0.5, "step": 0.01,
            }
            for i in range(N_PCS)
        ],
        "vocoder": state.vocoder_type,
    }


@app.get("/generate")
async def generate(request: Request):
    await rate_limit_generate(request)
    pc_values = _parse_pc_values(request)
    z = _pc_to_latent(pc_values)

    with torch.no_grad():
        spec = state.model.decode(z)
        audio = _spec_to_audio(spec)

    buf = io.BytesIO()
    sf.write(buf, audio, SAMPLE_RATE, format="WAV")
    buf.seek(0)

    return StreamingResponse(buf, media_type="audio/wav")


@app.get("/spectrogram")
async def spectrogram_data(request: Request):
    await rate_limit_spectrogram(request)
    pc_values = _parse_pc_values(request)
    z = _pc_to_latent(pc_values)

    with torch.no_grad():
        spec = state.model.decode(z)

    return {"data": spec.squeeze().cpu().numpy().tolist()}


@app.get("/health")
async def health():
    return {"status": "ok", "vocoder": state.vocoder_type, "device": str(state.device)}