import os
import uuid
import asyncio
from typing import List
from fastapi import FastAPI, UploadFile, File, HTTPException
from dotenv import load_dotenv

from binary_classifier.classify import GeminiSportsClassifier
from embeddings_ import process_inputs
from s3_file_uploader import upload_to_s3

load_dotenv()

app = FastAPI(title="Sports Classification API")
classifier = GeminiSportsClassifier()

UPLOAD_DIR = "temp_uploads"
os.makedirs(UPLOAD_DIR, exist_ok=True)

@app.post("/classify")
async def classify_media(files: List[UploadFile] = File(...)):
    if not files:
        raise HTTPException(status_code=400, detail="No files uploaded")
    file_paths = []
    path_to_original = {}
    try:
        for file in files:
            ext = os.path.splitext(file.filename)[1]
            temp_name = f"{uuid.uuid4()}{ext}"
            temp_path = os.path.join(UPLOAD_DIR, temp_name)
            with open(temp_path, "wb") as f:
                content = await file.read()
                f.write(content)
            file_paths.append(temp_path)
            path_to_original[temp_path] = file.filename
        first_level_results = await classifier.classify_batch(file_paths)
        allowed = {"not_sports", "error"}
        filtered_items = []

        for item in first_level_results:
            if item.get("answer") not in allowed:
                path = item["file"]
                file_name = path_to_original.get(path)
                ext = os.path.splitext(path)[1].lower()

                media_type = "video" if ext in [".mp4", ".avi", ".mov"] else "image"

                filtered_items.append({
                    "file_path": path,
                    "file_name": file_name,
                    "media_type": media_type,
                    "answer": item.get("answer")
                })

        if not filtered_items:
            return []

        embed_paths = [item["file_path"] for item in filtered_items]
        embedding_result = process_inputs(embed_paths)
        image_embeddings = embedding_result.get("images", [])
        video_embeddings = embedding_result.get("videos", [])
        img_idx = 0
        vid_idx = 0

        for item in filtered_items:
            if item["media_type"] == "image":
                item["embedding"] = image_embeddings[img_idx].tolist() if len(image_embeddings) > 0 else None
                img_idx += 1
            else:
                item["embedding"] = video_embeddings[vid_idx].tolist() if len(video_embeddings) > 0 else None
                vid_idx += 1

        async def upload_single(item):
            url = await upload_to_s3(item["file_path"], item["file_name"])
            item["s3_url"] = url
            return item

        uploaded_results = await asyncio.gather(*[upload_single(item) for item in filtered_items])

        return uploaded_results

    finally:
        for path in file_paths:
            if os.path.exists(path):
                os.remove(path)