import os
import cv2
import torch
import tempfile
import httpx
import numpy as np
import yt_dlp
from pydantic import BaseModel
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from transformers import VideoMAEForVideoClassification

app = FastAPI(
    title="Video Activity Recognition API",
    description="Classifies actions in a video. Supports file uploads and generic video links."
)

# --- MODEL SETUP ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_path = "model_final"

print("Loading model...")
eval_model = VideoMAEForVideoClassification.from_pretrained(model_path)

if torch.cuda.is_available():
    eval_model = eval_model.half()
eval_model = eval_model.to(device).eval()

MEAN = np.array(getattr(eval_model.config, "mean", [0.485, 0.456, 0.406]), dtype=np.float32)
STD = np.array(getattr(eval_model.config, "std", [0.229, 0.224, 0.225]), dtype=np.float32)
RESIZE_TO = tuple(getattr(eval_model.config, "resize_to", [224, 224]))

DIRECT_VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".flv", ".m4v")

# --- PYDANTIC MODELS ---
class URLRequest(BaseModel):
    url: str

# --- HELPER FUNCTIONS ---

def extract_and_preprocess_frames(video_path: str, num_frames: int = 16) -> torch.Tensor:
    cap = cv2.VideoCapture(video_path)
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    if frame_count <= 0:
        cap.release()
        raise ValueError("Could not read video or video has no frames.")

    indices = set(np.linspace(0, frame_count - 1, num_frames, dtype=int))
    frames_dict = {}

    for idx in range(frame_count):
        ret, frame = cap.read()
        if not ret:
            break
        if idx in indices:
            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
            frame_resized = cv2.resize(frame_rgb, RESIZE_TO)
            frame_scaled = frame_resized.astype(np.float32) / 255.0
            frame_normalized = (frame_scaled - MEAN) / STD
            frame_transposed = np.transpose(frame_normalized, (2, 0, 1))
            frames_dict[idx] = frame_transposed

    cap.release()

    processed_frames = [frames_dict[i] for i in sorted(frames_dict.keys())]

    while len(processed_frames) < num_frames:
        processed_frames.append(processed_frames[-1])

    processed_frames = processed_frames[:num_frames]
    return torch.tensor(np.array(processed_frames))


def _is_direct_video_url(url: str) -> bool:
    """Returns True if the URL points directly to a video file by extension."""
    clean = url.split("?")[0].split("#")[0].lower()
    return clean.endswith(DIRECT_VIDEO_EXTENSIONS)


def _download_direct(url: str) -> str:
    """Downloads a direct video URL via HTTP streaming. Returns temp file path."""
    tmp_fd, tmp_path = tempfile.mkstemp(suffix=".mp4")
    os.close(tmp_fd)
    try:
        with httpx.Client(follow_redirects=True, timeout=60) as client:
            with client.stream("GET", url) as response:
                response.raise_for_status()
                with open(tmp_path, "wb") as f:
                    for chunk in response.iter_bytes(chunk_size=8192):
                        f.write(chunk)
        if os.path.getsize(tmp_path) == 0:
            raise ValueError("Downloaded file is empty.")
        return tmp_path
    except Exception as e:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)
        raise ValueError(f"Direct download failed: {e}")


def _download_with_ytdlp(url: str) -> str:
    """Downloads a platform video URL (YouTube, TikTok, etc.) via yt-dlp. Returns temp file path."""
    tmp_fd, tmp_path = tempfile.mkstemp(suffix=".mp4")
    os.close(tmp_fd)

    # yt-dlp appends its own extension — use a base template and resolve afterward
    base_path = tmp_path.replace(".mp4", "")

    ydl_opts = {
        "format": "bestvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best",
        "outtmpl": base_path + ".%(ext)s",
        "quiet": True,
        "no_warnings": True,
        "merge_output_format": "mp4",
        "socket_timeout": 30,
        "retries": 3,
        "fragment_retries": 3,
    }

    ext = "mp4"
    try:
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            info = ydl.extract_info(url, download=True)
            ext = info.get("ext", "mp4")

        # Resolve actual output path (yt-dlp appends extension)
        expected = base_path + f".{ext}"
        if not os.path.exists(expected):
            expected = base_path + ".mp4"

        if not os.path.exists(expected) or os.path.getsize(expected) == 0:
            raise ValueError("yt-dlp produced an empty or missing file.")

        # Clean up the original empty placeholder if different
        if os.path.exists(tmp_path) and tmp_path != expected:
            os.remove(tmp_path)

        return expected

    except yt_dlp.utils.DownloadError as e:
        for p in [tmp_path, base_path + f".{ext}", base_path + ".mp4"]:
            if os.path.exists(p):
                os.remove(p)
        raise ValueError(f"yt-dlp download failed: {e}")
    except Exception as e:
        for p in [tmp_path, base_path + f".{ext}", base_path + ".mp4"]:
            if os.path.exists(p):
                os.remove(p)
        raise ValueError(f"Unexpected download error: {e}")


def download_video_from_url(url: str) -> str:
    """
    Smart router: uses httpx for direct video file URLs (.mp4, .avi, etc.)
    and yt-dlp for platform URLs (YouTube, TikTok, Twitter/X, etc.).
    Returns the path to a downloaded temp file. Caller must delete it.
    """
    if _is_direct_video_url(url):
        return _download_direct(url)
    else:
        return _download_with_ytdlp(url)


def run_inference(video_path: str) -> dict:
    """Handles the core inference logic for any valid local video path."""
    try:
        frames_tensor = extract_and_preprocess_frames(video_path, num_frames=eval_model.config.num_frames)
    except Exception as e:
        raise HTTPException(status_code=422, detail=f"Frame extraction failed: {e}")

    try:
        video = frames_tensor.unsqueeze(0).to(device)
        if torch.cuda.is_available():
            video = video.half()

        with torch.no_grad():
            logits = eval_model(pixel_values=video).logits
            probs = torch.softmax(logits, dim=1)[0]

        predicted_id = logits.argmax(1).cpu().item()

        def get_label(class_id):
            return eval_model.config.id2label.get(
                class_id,
                eval_model.config.id2label.get(str(class_id), f"Class_{class_id}")
            )

        predicted_label = get_label(predicted_id)

        all_scores = {
            get_label(i): round(prob.item(), 4)
            for i, prob in enumerate(probs.cpu())
        }

        del video, logits
        torch.cuda.empty_cache()

        return {
            "predicted_activity": predicted_label,
            "confidence": all_scores.get(predicted_label, 0.0),
            "all_scores": all_scores
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Inference error: {e}")


# --- API ENDPOINTS ---

@app.post("/predict/file")
async def predict_from_file(file: UploadFile = File(...)):
    """Endpoint for direct video file uploads."""
    temp_video_path = None
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_file:
            temp_video_path = tmp_file.name
            content = await file.read()
            tmp_file.write(content)

        result = run_inference(temp_video_path)
        return JSONResponse(result)

    finally:
        if temp_video_path and os.path.exists(temp_video_path):
            os.remove(temp_video_path)


@app.post("/predict/url")
async def predict_from_url(request: URLRequest):
    """Endpoint for video URLs — direct files (.mp4, .avi, etc.) or platform links (YouTube, TikTok, etc.)."""
    temp_video_path = None
    try:
        temp_video_path = download_video_from_url(request.url)
        result = run_inference(temp_video_path)
        return JSONResponse(result)

    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))

    finally:
        if temp_video_path and os.path.exists(temp_video_path):
            os.remove(temp_video_path)


@app.get("/health")
def health_check():
    return {"status": "healthy", "device": str(device)}