from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel
from agent_workflow import workflow_app
import uuid
import asyncio
from fastapi.responses import FileResponse
from pydantic import BaseModel
from agent_workflow import workflow_app

app = FastAPI(title="Agentic AI Document Review System", version="1.0.0")

app.mount("/static", StaticFiles(directory="static"), name="static")

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

class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

manager = ConnectionManager()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    thread_id = None
    task = None
    config = None

    async def poll_status():
        try:
            while True:
                if config:
                    state = workflow_app.get_state(config)
                    if state.values:
                        next_node = state.next
                        values = state.values
                        current_status = values.get("status", "unknown")
                        is_waiting_for_human = "human_approval" in next_node
                        
                        await websocket.send_json({
                            "type": "status_update",
                            "status": current_status,
                            "revision_count": values.get("revision_count", 0),
                            "waiting_for_human": is_waiting_for_human,
                            "latest_feedback": values.get("review_feedback", ""),
                            "document": values.get("current_draft", "")
                        })
                await asyncio.sleep(1)
        except asyncio.CancelledError:
            pass
        except Exception as e:
            print(f"Error in poll_status: {e}")

    try:
        while True:
            data = await websocket.receive_json()
            action = data.get("action")

            if action == "start":
                raw_input = data.get("raw_input")
                thread_id = str(uuid.uuid4())
                config = {"configurable": {"thread_id": thread_id}}
                
                initial_state = {
                    "raw_input": raw_input,
                    "current_draft": "",
                    "review_feedback": "",
                    "status": "started",
                    "revision_count": 0
                }
                
                await websocket.send_json({
                    "type": "workflow_started",
                    "thread_id": thread_id
                })
                
                if task:
                    task.cancel()
                task = asyncio.create_task(poll_status())
                
                # Run the blocking workflow invoke in a thread to not block the WebSocket event loop
                await asyncio.to_thread(workflow_app.invoke, initial_state, config)

            elif action == "approve":
                if config:
                    state = workflow_app.get_state(config)
                    next_node = state.next
                    if "human_approval" in next_node:
                        # Resume workflow
                        await asyncio.to_thread(workflow_app.invoke, None, config)
                        await websocket.send_json({"type": "workflow_approved"})

    except WebSocketDisconnect:
        manager.disconnect(websocket)
        if task:
            task.cancel()

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="127.0.0.1", port=8000, reload=True)