import base64
import shutil
from pathlib import Path
from typing import Dict, List, Optional
from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from agent import run_agent


class ChatRequest(BaseModel):
    user_id: str
    user_message: str
    api_key: str
    provider: str = "openrouter"
    model: str = "perplexity/llama-3-sonar-large-32k-online"
    current_chat_history: Optional[List[Dict[str, str]]] = None
    user_files: Optional[Dict[str, str]] = None


class ChatResponse(BaseModel):
    agent_response: str
    updated_chat_history: List[Dict[str, str]]
    updated_files: Dict[str, str]


@asynccontextmanager
async def lifespan(app: FastAPI):
    yield


app = FastAPI(title="React Agent API", lifespan=lifespan)


@app.post("/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    if not request.api_key:
        raise HTTPException(status_code=400, detail="api_key is required")

    workspace = Path(f"/tmp/{request.user_id}")
    workspace.mkdir(parents=True, exist_ok=True)

    try:
        if request.user_files:
            for filename, content_b64 in request.user_files.items():
                filepath = workspace / filename
                filepath.parent.mkdir(parents=True, exist_ok=True)
                filepath.write_bytes(base64.b64decode(content_b64))

        result = await run_agent(
            api_key=request.api_key,
            model=request.model,
            provider=request.provider,
            workspace=workspace,
            user_message=request.user_message,
            chat_history=request.current_chat_history or [],
        )

        updated_files = {}
        if workspace.exists():
            for filepath in sorted(workspace.rglob("*")):
                if filepath.is_file():
                    rel_path = str(filepath.relative_to(workspace))
                    updated_files[rel_path] = base64.b64encode(
                        filepath.read_bytes()
                    ).decode()

        return ChatResponse(
            agent_response=result["agent_response"],
            updated_chat_history=result["updated_chat_history"],
            updated_files=updated_files,
        )

    finally:
        if workspace.exists():
            shutil.rmtree(workspace, ignore_errors=True)