# app.py - VERSÃO COMPLETA COM VOZ (BASE64) E VISÃO
import os
import base64
import io
import asyncio
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel 
from PIL import Image
from jade.core import JadeAgent 
from jade.scholar import ScholarAgent
from jade.heavy_mode import JadeHeavyAgent
from jade.webdev import WebDevAgent

print("Iniciando a J.A.D.E. com FastAPI...")
jade_agent = JadeAgent() 
scholar_agent = ScholarAgent()
# Instantiate Heavy Agent. It uses environment variables.
jade_heavy_agent = JadeHeavyAgent()
# Instantiate WebDev Agent
webdev_agent = WebDevAgent()

print("J.A.D.E. pronta para receber requisições.")

app = FastAPI(title="J.A.D.E. API")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
)

# Mount generated directory for static files (PDFs, Images)
os.makedirs("backend/generated", exist_ok=True)
app.mount("/generated", StaticFiles(directory="backend/generated"), name="generated")

# Dicionário global para armazenar sessões de usuários
# Structure: user_sessions[user_id] = { "jade": [...], "scholar": [...], "heavy": [...] }
user_sessions = {}

class UserRequest(BaseModel):
    user_input: str
    image_base64: str | None = None
    user_id: str | None = None
    agent_type: str = "jade" # "jade", "scholar", "heavy"
    web_search: bool = False  # Toggle para busca web na J.A.D.E.
    thinking_mode: bool = False  # Toggle para modo CoT/Thinking

@app.post("/chat")
async def handle_chat(request: UserRequest):
    try:
        user_id = request.user_id if request.user_id else "default_user"
        agent_type = request.agent_type.lower()
        
        if user_id not in user_sessions:
            print(f"Nova sessão criada para: {user_id}")
            user_sessions[user_id] = {
                "jade": [jade_agent.system_prompt],
                "scholar": [],
                "heavy": []
            }
        
        # Ensure sub-keys exist
        if "jade" not in user_sessions[user_id]: user_sessions[user_id]["jade"] = [jade_agent.system_prompt]
        if "scholar" not in user_sessions[user_id]: user_sessions[user_id]["scholar"] = []
        if "heavy" not in user_sessions[user_id]: user_sessions[user_id]["heavy"] = []

        vision_context = None
        if request.image_base64:
            try:
                header, encoded_data = request.image_base64.split(",", 1)
                image_bytes = base64.b64decode(encoded_data)
                pil_image = Image.open(io.BytesIO(image_bytes))
                # Jade handles vision processing
                vision_context = jade_agent.image_handler.process_pil_image(pil_image)
            except Exception as img_e:
                print(f"Erro ao processar imagem Base64: {img_e}")
                vision_context = "Houve um erro ao analisar a imagem."

        final_user_input = request.user_input if request.user_input else "Descreva a imagem em detalhes."

        bot_response_text = ""
        audio_path = None
        
        if agent_type == "scholar":
            current_history = user_sessions[user_id]["scholar"]
            bot_response_text, audio_path, updated_history = scholar_agent.respond(
                history=current_history,
                user_input=final_user_input,
                user_id=user_id,
                vision_context=vision_context
            )
            user_sessions[user_id]["scholar"] = updated_history
            
        elif agent_type == "heavy":
            current_history = user_sessions[user_id]["heavy"]
            # Heavy agent is async
            bot_response_text, audio_path, updated_history = await jade_heavy_agent.respond(
                history=current_history,
                user_input=final_user_input,
                user_id=user_id,
                vision_context=vision_context,
                web_search=request.web_search  # Passa web search para Heavy Mode
            )
            user_sessions[user_id]["heavy"] = updated_history
            
        else:
            # Default to J.A.D.E.
            current_history = user_sessions[user_id]["jade"]
            # Jade agent is synchronous, run directly
            bot_response_text, audio_path, updated_history = jade_agent.respond(
                history=current_history,
                user_input=final_user_input,
                user_id=user_id,
                vision_context=vision_context,
                web_search=request.web_search,
                thinking_mode=request.thinking_mode  # Passa o toggle de thinking
            )
            user_sessions[user_id]["jade"] = updated_history

        # Audio Logic
        audio_base64 = None
        if audio_path and os.path.exists(audio_path):
            print(f"Codificando arquivo de áudio: {audio_path}")
            with open(audio_path, "rb") as audio_file:
                audio_bytes = audio_file.read()
                audio_base64 = base64.b64encode(audio_bytes).decode('utf-8')
            
            # Remove only if temp file
            if "backend/generated" not in audio_path:
                os.remove(audio_path) 
        
        return {
            "success": True,
            "bot_response": bot_response_text,
            "audio_base64": audio_base64 
        }
    except Exception as e:
        print(f"Erro crítico no endpoint /chat: {e}")
        return {"success": False, "error": str(e)}

# WebDev Vibe Coder endpoint
class WebDevRequest(BaseModel):
    prompt: str
    existing_code: str | None = None
    mode: str = "html"  # "html" or "react"
    error_message: str | None = None  # For agentic error fixing
    model: str | None = None  # Model to use (e.g., "kimi-k2", "claude-sonnet")

@app.post("/webdev/generate")
async def handle_webdev(request: WebDevRequest):
    try:
        result = webdev_agent.generate(
            prompt=request.prompt,
            refine_code=request.existing_code,
            mode=request.mode,
            error_message=request.error_message,
            model=request.model
        )
        return result
    except Exception as e:
        print(f"Erro no endpoint /webdev: {e}")
        return {"success": False, "error": str(e)}

@app.get("/")
def root():
    return {"message": "Servidor J.A.D.E. com FastAPI está online."}

if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 7860))
    print(f"Iniciando o servidor Uvicorn em http://0.0.0.0:{port}")
    uvicorn.run(app, host="0.0.0.0", port=port)