import os
import uvicorn
import json
from fastapi import FastAPI, UploadFile, File, Form
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import List, Optional

try:
    from logic import GaiaApp
except ImportError:
    class GaiaApp:
        def process_input(self, msg, hist, files):
            return "Error: logic.py not found", hist, None

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

gaia_engine = GaiaApp()

class Message(BaseModel):
    role: str
    content: str

class ChatRequest(BaseModel):
    message: str
    history: List[Message]

def format_history(messages: List[Message]) -> List[tuple]:
    """Convert message history from React format to agent format."""
    formatted = []
    for msg in messages:
        if msg.role == 'user':
            formatted.append((msg.content, None))
        elif msg.role == 'model' and formatted and formatted[-1][1] is None:
            formatted[-1] = (formatted[-1][0], msg.content)
    return formatted

@app.post("/api/chat")
async def chat_endpoint(
    message: str = Form(...),
    history: str = Form("[]"),
    files: List[UploadFile] = File(default=[])
):
    """Process chat message with optional file uploads and return agent response."""
    try:
        messages = json.loads(history)
        formatted_history = format_history([Message(**msg) for msg in messages])
        
        file_paths = []
        if files:
            os.makedirs("uploads", exist_ok=True)
            for file in files:
                file_path = f"uploads/{file.filename}"
                with open(file_path, "wb") as f:
                    content = await file.read()
                    f.write(content)
                file_paths.append(file_path)
        
        generator = gaia_engine.process_input(message, formatted_history, file_paths or None)
        
        final_state = None
        for step in generator:
            final_state = step
            
        if final_state and final_state[1]:
            # final_state[1] is the history as list of dicts with role/content
            last_message = final_state[1][-1]
            if isinstance(last_message, dict):
                bot_response = last_message.get("content", "No response")
            else:
                # Fallback for tuple format
                bot_response = last_message[1] if len(last_message) > 1 else "No response"
            return {"reply": bot_response}
        
        return {"reply": "No response from agent."}
    except Exception as e:
        print(f"Chat endpoint error: {str(e)}")
        import traceback
        traceback.print_exc()
        return {"reply": f"Error: {str(e)}"}

ALLOWED_EXTENSIONS = {".css", ".js", ".png", ".html"}

@app.get("/")
async def serve_index():
    return FileResponse("index.html")

@app.get("/{filename}")
async def serve_static(filename: str):
    ext = os.path.splitext(filename)[1]
    if ext in ALLOWED_EXTENSIONS and os.path.exists(filename):
        return FileResponse(filename)
    return FileResponse("index.html", status_code=404)

if __name__ == "__main__":
    print("\n" + "="*60)
    print("🚀 GAIA Agent Starting...")
    print("="*60)
    print(f"📍 Open in browser: http://localhost:7860")
    print(f"📍 Or use: http://127.0.0.1:7860")
    print("="*60 + "\n")
    uvicorn.run(
        "app:app",
        host="0.0.0.0",
        port=7860,
        reload=True
    )