from fastapi import FastAPI, UploadFile, File, Form
import google.generativeai as genai
import os
from PIL import Image
import io

app = FastAPI()

# 1. More aggressive System Instruction
system_prompt = (
    "You are a strict emotion classifier. "
    "Rules: \n"
    "1. Only output ONE word. \n"
    "2. No explanations. \n"
    "3. No punctuation. \n"
    "Example Output: HAPPY"
)

genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))

model = genai.GenerativeModel(
    model_name='gemma-4-26b-a4b-it',
    system_instruction=system_prompt
)

@app.post("/analyze")
async def analyze(text: str = Form(None), image: UploadFile = File(None)):
    try:
        user_content = []
        if text:
            # We add a clear command at the end of the text
            user_content.append(f"Input: {text}\n\nEmotion (one word only):")
        
        if image:
            img_bytes = await image.read()
            img = Image.open(io.BytesIO(img_bytes))
            user_content.append(img)
            user_content.append("Emotion in image (one word only):")

        if not user_content:
            return {"error": "No input"}

        response = model.generate_content(user_content)
        
        # We take only the LAST word just in case the AI still talks too much
        raw_text = response.text.strip()
        one_word = raw_text.split()[-1].replace(".", "").upper()
        
        return {"emotion": one_word}
    
    except Exception as e:
        return {"error": str(e)}