import uuid
import asyncio
import logging
import os
from pathlib import Path
from contextlib import asynccontextmanager
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse

from config import HOST, PORT
from websocket.handler import manager

# Frontend dist directory
FRONTEND_DIR = Path(__file__).parent.parent / "frontend" / "dist"

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("Starting Telegram Web Client server...")
    yield
    logger.info("Shutting down server...")
    # Cleanup all connections
    for session_id in list(manager.active_connections.keys()):
        await manager.disconnect(session_id)


app = FastAPI(
    title="Telegram Web Client",
    description="WebSocket-based Telegram bot chat client",
    version="1.0.0",
    lifespan=lifespan
)

# CORS configuration
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Configure appropriately for production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/api")
async def api_root():
    return {"status": "ok", "message": "Telegram Web Client API"}


@app.get("/health")
async def health():
    return {"status": "healthy"}


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    """Main WebSocket endpoint for chat"""
    session_id = str(uuid.uuid4())

    try:
        await manager.connect(websocket, session_id)

        # Send session ID to client
        await manager.send_message(session_id, {
            "type": "connected",
            "data": {"session_id": session_id}
        })

        while True:
            data = await websocket.receive_json()
            await manager.handle_message(session_id, data)

    except WebSocketDisconnect:
        logger.info(f"Client disconnected: {session_id}")
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
    finally:
        await manager.disconnect(session_id)


@app.websocket("/ws/{session_id}")
async def websocket_with_session(websocket: WebSocket, session_id: str):
    """WebSocket endpoint with specific session ID (for reconnection)"""
    try:
        await manager.connect(websocket, session_id)

        await manager.send_message(session_id, {
            "type": "connected",
            "data": {"session_id": session_id}
        })

        while True:
            data = await websocket.receive_json()
            await manager.handle_message(session_id, data)

    except WebSocketDisconnect:
        logger.info(f"Client disconnected: {session_id}")
    except Exception as e:
        logger.error(f"WebSocket error: {e}")
    finally:
        await manager.disconnect(session_id)


# Mount static files if frontend is built
if FRONTEND_DIR.exists():
    app.mount("/assets", StaticFiles(directory=FRONTEND_DIR / "assets"), name="assets")

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

    @app.get("/{full_path:path}")
    async def serve_spa(full_path: str):
        # Serve static file if exists, otherwise index.html for SPA
        file_path = FRONTEND_DIR / full_path
        if file_path.exists() and file_path.is_file():
            return FileResponse(file_path)
        return FileResponse(FRONTEND_DIR / "index.html")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host=HOST, port=PORT)