import io
import traceback
import numpy as np
from typing import List
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse
from PIL import Image

from data_loader import load_music_engine, load_image_processor
from music_engine.llm_bridge import LLMAcousticBridge

app = FastAPI(title="EmoM API", version="1.0.0")

ml_context = {
    "image_processor": None,
    "music_matcher": None,
    "llm_bridge": None
}

@app.on_event("startup")
async def startup_event():
    print("Loading ML models...")
    ml_context["image_processor"] = load_image_processor()
    ml_context["music_matcher"] = load_music_engine()
    ml_context["llm_bridge"] = LLMAcousticBridge()
    print("Initialization complete.")

@app.post("/analyze")
async def analyze_event_endpoint(files: List[UploadFile] = File(...)):
    try:
        images = []
        for file in files:
            image_bytes = await file.read()
            img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
            images.append(img)
            
        print(f"Processing batch: {len(images)} images.")
        
        img_processor = ml_context["image_processor"]
        matcher = ml_context["music_matcher"]
        llm = ml_context["llm_bridge"]

        all_v, all_a = [], []
        all_objects = []

        for img in images:
            embedding = img_processor.extract_embedding(img)
            v, a = matcher.predict_va(embedding)
            all_v.append(v)
            all_a.append(a)
            
            caption = img_processor.describe_scene(img)
            all_objects.append(caption)

        target_v = float(np.mean(all_v))
        target_a = float(np.mean(all_a))
        unique_semantics = list(set(all_objects))

        llm_profile = llm.get_acoustic_profile(target_v, target_a, unique_semantics)

        playlist_df = matcher.find_nearest_tracks(target_v, target_a, llm_profile=llm_profile, top_k=15)
        tracks_list = playlist_df.to_dict(orient="records")

        return JSONResponse(content={
            "status": "success",
            "images_processed": len(images),
            "target_v": target_v,
            "target_a": target_a,
            "llm_profile": llm_profile,
            "semantics": unique_semantics,
            "tracks": tracks_list
        })
        
    except Exception as e:
        print(traceback.format_exc()) 
        raise HTTPException(status_code=500, detail=str(e))