import os
from tempfile import NamedTemporaryFile
from typing import Dict, List

import numpy as np
import torch
import torch.nn as nn
import torchvision.transforms as T
from decord import VideoReader, cpu
from fastapi import FastAPI, File, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from torchvision.models.video import R2Plus1D_18_Weights, r2plus1d_18

app = FastAPI(title="Elephant Stress API")

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

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
CKPT_PATH = "elephant_stress_r2plus1d18.pt"
CLIP_LEN = 16
NUM_CLASSES = 2

EVAL_FRAME_TRANSFORM = T.Compose(
    [
        T.ToPILImage(),
        T.Resize((128, 171)),
        T.CenterCrop((112, 112)),
        T.ToTensor(),
    ]
)

KIN_MEAN = torch.tensor([0.43216, 0.394666, 0.37645]).view(3, 1, 1)
KIN_STD = torch.tensor([0.22803, 0.22145, 0.216989]).view(3, 1, 1)


def _normalize_state_dict(loaded_obj: object) -> Dict[str, torch.Tensor]:
    if isinstance(loaded_obj, dict):
        for key in ("state_dict", "model_state_dict", "model"):
            nested = loaded_obj.get(key)
            if isinstance(nested, dict):
                loaded_obj = nested
                break
    if not isinstance(loaded_obj, dict):
        raise ValueError("Checkpoint does not contain a valid state dict.")

    normalized: Dict[str, torch.Tensor] = {}
    for key, value in loaded_obj.items():
        clean_key = key[7:] if key.startswith("module.") else key
        normalized[clean_key] = value
    return normalized


def _load_model() -> torch.nn.Module:
    if not os.path.exists(CKPT_PATH):
        raise FileNotFoundError(
            f"Missing checkpoint: {CKPT_PATH}. Add it to repo root with git-lfs."
        )

    model = r2plus1d_18(weights=R2Plus1D_18_Weights.DEFAULT)
    model.fc = nn.Linear(model.fc.in_features, NUM_CLASSES)

    loaded = torch.load(CKPT_PATH, map_location=DEVICE)
    state_dict = _normalize_state_dict(loaded)
    model.load_state_dict(state_dict, strict=True)
    model.to(DEVICE)
    model.eval()
    return model


MODEL = _load_model()


def sample_indices(num_frames: int, clip_len: int = CLIP_LEN) -> np.ndarray:
    if num_frames <= clip_len:
        return np.linspace(0, max(num_frames - 1, 0), clip_len).astype(int)
    start = np.random.randint(0, num_frames - clip_len + 1)
    return np.arange(start, start + clip_len)


def get_clip_tensor(video_reader: VideoReader, inds: np.ndarray) -> torch.Tensor:
    frames = video_reader.get_batch(inds).asnumpy()  # (T, H, W, C)
    transformed: List[torch.Tensor] = []
    for i in range(frames.shape[0]):
        frame = EVAL_FRAME_TRANSFORM(frames[i])
        frame = (frame - KIN_MEAN) / KIN_STD
        transformed.append(frame)

    # (C, T, H, W)
    return torch.stack(transformed, dim=1)


@torch.no_grad()
def predict_path(video_path: str, n_clips: int = 10) -> Dict[str, object]:
    vr = VideoReader(video_path, ctx=cpu(0))
    total_frames = len(vr)
    if total_frames == 0:
        raise ValueError("Video has no readable frames.")

    logits_sum = torch.zeros((1, NUM_CLASSES), device=DEVICE)
    for _ in range(n_clips):
        inds = sample_indices(total_frames, CLIP_LEN)
        clip = get_clip_tensor(vr, inds).unsqueeze(0).to(DEVICE)
        logits_sum += MODEL(clip)

    logits = logits_sum / n_clips
    probs = torch.softmax(logits, dim=1)[0].detach().cpu().numpy().tolist()
    pred = int(np.argmax(probs))

    return {
        "prediction": "Stress (1)" if pred == 1 else "Normal (0)",
        "probs": {"normal": probs[0], "stress": probs[1]},
        "n_clips": n_clips,
    }


@app.get("/")
def health() -> Dict[str, str]:
    return {"status": "ok"}


@app.post("/predict")
async def predict_api(file: UploadFile = File(...)) -> Dict[str, object]:
    suffix = os.path.splitext(file.filename or "")[1] or ".mp4"

    with NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        tmp.write(await file.read())
        tmp_path = tmp.name

    try:
        return predict_path(tmp_path)
    except Exception as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    finally:
        try:
            os.remove(tmp_path)
        except OSError:
            pass