import os
import sys
import uvicorn
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
from fastapi.responses import HTMLResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates
from fastapi.middleware.cors import CORSMiddleware
from dotenv import load_dotenv

from src.websocket.websocket_manager import ConnectionManager
from src.users.router import router as user_router
from src.files.router import router as file_router
from src.bot.router import router as bot_router
from src.etl.router import router as etl_router
from src.login.router import router as login_router
from src.conversation.router import router as conversation_router

from src.chat.index import chat_manager
from src.chat.callback import StreamingLLMCallbackHandler
from src.chat.schemas import ChatResponse

load_dotenv()

tags_metadata = [
    {
        "name": "users",
        "description": "Operations with users.",
    },
    {
        "name": "files",
        "description": "Operations with file handling.",
    },
    {
        "name": "bot",
        "description": "Operations with bot.",
    },
    {
        "name": "etl",
        "description": "Operations with ETL process.",
    },
     {
        "name": "conversations",
        "description": "Operations with conversations process.",
    },

]


app = FastAPI(openapi_tags=tags_metadata)


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

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


app.include_router(user_router)
app.include_router(file_router)
app.include_router(bot_router)
app.include_router(etl_router)
app.include_router(login_router)
app.include_router(conversation_router)


@app.get("/",response_class=HTMLResponse)
async def read_index(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})


@app.get("/knowledge-base",response_class=HTMLResponse, tags=["ui"])
async def read_index(request: Request):
    return templates.TemplateResponse("knowledge-base.html", {"request": request})


@app.get("/websocket-tester",response_class=HTMLResponse,  tags=["tester"], summary="UI for testing websocket")
async def read_index(request: Request):
    return templates.TemplateResponse("websocket-tester.html", {"request": request})

@app.get("/upload-tester",response_class=HTMLResponse,  tags=["tester"], summary="UI for testing file upload")
async def read_index(request: Request):
    return templates.TemplateResponse("upload-tester.html", {"request": request})


manager = ConnectionManager()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)

    while True:
        try:
            data = await websocket.receive_json()
            q = data["message"]

            stream_handler = StreamingLLMCallbackHandler(websocket)
            await websocket.send_text(".-botResponseStartStreaming")
            response = await chat_manager.chat_response({
                "question": q,
                "chat_session": data["chat_session"],
                "bot_id": data["bot_id"],
                "stream_handler": stream_handler,
            })

            await websocket.send_text(".-botResponseEndStreaming")
        except WebSocketDisconnect:
            manager.disconnect(websocket)
            # await websocket.send_text(f"Connection lost")

if __name__ == "__main__":
    uvicorn.run(app, host='0.0.0.0', port=int(os.environ.get('PORT', 8000)), reload=True)