app.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from fastapi.responses import Response
from typing import List, Optional
import aiohttp
import asyncio
import io
import logging

# Import all generators
from generators.legendary import PremiumGenerator
from generators.common import CommonGenerator
from generators.custom import CustomGenerator
from generators.profile import ProfileGenerator

app = FastAPI()
logging.basicConfig(level=logging.INFO)

# Pre-load the heavy VFX into the Space's 16GB of RAM on startup
PremiumGenerator.preload_vfx_cache()

# ==========================================
# 📦 PYDANTIC MODELS
# ==========================================
class CardRequest(BaseModel):
    image_url: str
    char_name: str
    print_num: int
    rarity: str
    element: str
    custom_frame_url: Optional[str] = None

class ProfileCardData(BaseModel):
    name: str
    print_num: int
    rarity: str
    element: str
    image_url: str
    custom_frame_url: Optional[str] = None

class ProfileRequest(BaseModel):
    username: str
    team_data: List[Optional[ProfileCardData]]

# ==========================================
# 🌐 ASYNC DOWNLOAD HELPERS
# ==========================================
async def download_image(url: str) -> bytes:
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as resp:
            if resp.status != 200:
                raise HTTPException(status_code=400, detail=f"Failed to download asset: {url}")
            return await resp.read()

async def fetch_profile_card(card: ProfileCardData, session: aiohttp.ClientSession) -> Optional[dict]:
    """Fetches the character art and frame concurrently for a single profile slot."""
    if not card:
        return None
    
    img_bytes = None
    frame_bytes = None
    
    try:
        # We fire both downloads at the exact same time to cut network waiting in half
        tasks = [session.get(card.image_url)]
        if card.custom_frame_url:
            tasks.append(session.get(card.custom_frame_url))
            
        responses = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Parse Character Art
        if not isinstance(responses[0], Exception) and responses[0].status == 200:
            img_bytes = await responses[0].read()
            
        # Parse Custom Frame (If requested)
        if len(responses) > 1 and not isinstance(responses[1], Exception) and responses[1].status == 200:
            frame_bytes = await responses[1].read()
            
    except Exception as e:
        logging.error(f"Failed to fetch profile images for {card.name}: {e}")
        return None
        
    return {
        "bytes": img_bytes,
        "frame_bytes": frame_bytes,
        "name": card.name,
        "print_num": card.print_num,
        "rarity": card.rarity,
        "element": card.element,
        "custom_frame_url": card.custom_frame_url
    }

# ==========================================
# 🚀 API ENDPOINTS
# ==========================================
@app.post("/generate")
async def generate_card(req: CardRequest):
    try:
        image_bytes = await download_image(req.image_url)
        frame_bytes = await download_image(req.custom_frame_url) if req.custom_frame_url else None

        if req.custom_frame_url and frame_bytes:
            generator = CustomGenerator()
            template = generator.build_template(image_bytes, frame_bytes, req.char_name)
        elif req.rarity in ["SSR", "UR"]:
            generator = PremiumGenerator()
            template = generator.build_template(image_bytes, req.element)
        else:
            generator = CommonGenerator()
            template = generator.build_template(image_bytes, req.element)

        final_buffer = generator.apply_text(template, req.char_name, req.print_num, req.element, req.rarity)
        template.close()
        
        return Response(content=final_buffer.getvalue(), media_type="image/webp")

    except Exception as e:
        logging.error(f"Generation failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.post("/generate_profile")
async def generate_profile(req: ProfileRequest):
    try:
        async with aiohttp.ClientSession() as session:
            # Launch all 8 slot downloads simultaneously
            tasks = [fetch_profile_card(card, session) if card else asyncio.sleep(0) for card in req.team_data]
            processed_team_data = await asyncio.gather(*tasks)
            
            # Clean out the asyncio.sleep(0) results and replace with None
            processed_team_data = [res if isinstance(res, dict) else None for res in processed_team_data]
        
        generator = ProfileGenerator()
        final_buffer = generator.build_profile(
            username=req.username, 
            stats={}, 
            showcase_bytes=None, 
            team_data=processed_team_data, 
            avatar_bytes=None
        )
        
        return Response(content=final_buffer.getvalue(), media_type="image/webp")
        
    except Exception as e:
        logging.error(f"Profile Generation failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/")
def health_check():
    return {"status": "Forge is online and ready."}