import asyncio
import uuid
from typing import Annotated, Optional

from fastapi import (
    Cookie,
    Depends,
    Request,
    Response,
    WebSocket,
    WebSocketDisconnect,
    WebSocketException,
    status,
)
from fastapi.responses import FileResponse
from fastapi.staticfiles import StaticFiles

from kennel.app import app, logger
from kennel.config import config
from kennel.engine.systems.network import (
    EntityBornEvent,
    Event,
    EventType,
    InitialStateEvent,
    Publishable,
    SetControllableEvent,
)
from kennel.engine.systems.system import SystemType
from kennel.kennel import (
    create_session_controllable_entities,
    entity_manager,
    kennel,
    system_manager,
)

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


loop = asyncio.get_event_loop()


@app.on_event("startup")
async def startup_event():
    logger.info("Starting Kennel...")
    loop.create_task(kennel.run())


@app.on_event("shutdown")
async def shutdown_event():
    logger.info("Stopping Kennel...")
    kennel.stop()
    loop.stop()
    logger.info("Kennel stopped")


@app.get("/")
async def index(request: Request):
    return FileResponse("static/dist/index.html", media_type="text/html")


@app.get("/assign")
async def assign(
    response: Response,
    session: Annotated[Optional[str], Cookie()] = None,
):
    if session is None:
        session = str(uuid.uuid4().hex)
        response.set_cookie(key="session", value=session, max_age=config.COOKIE_MAX_AGE)
    return {"session": session}


async def get_cookie_or_token(
    websocket: WebSocket,
    session: Annotated[str | None, Cookie()] = None,
):
    if session is None:
        raise WebSocketException(code=status.WS_1008_POLICY_VIOLATION)
    return session


class WebSocketClient(Publishable):
    def __init__(self, websocket: WebSocket):
        self.websocket = websocket

    async def publish(self, event: Event):
        await self.websocket.send_json([event.event_type, event.data])


@app.websocket("/ws")
async def websocket_endpoint(
    websocket: WebSocket,
    session: Annotated[str, Depends(get_cookie_or_token)],
):
    await websocket.accept()
    client = WebSocketClient(websocket)
    logger.info(f"Websocket connection established for session {session}")

    initial_state = InitialStateEvent(
        config.WORLD_WIDTH, config.WORLD_HEIGHT, kennel.entity_manager.to_dict()
    )
    await client.publish(initial_state)

    session_entities = create_session_controllable_entities(session)
    logger.info(f"Creating {len(session_entities)} entities for session {session}")
    try:
        network_system = system_manager.get_system(SystemType.NETWORK)
        if network_system is None:
            raise "Network system not found"

        for entity in session_entities:
            logger.info(f"Adding entity {entity.id} for session {session}")
            entity_manager.add_entity(entity)
            network_system.add_event(EntityBornEvent(entity))

            set_controllable_event = SetControllableEvent(entity.id, session)
            await client.publish(set_controllable_event)

        network_system.add_client(session, WebSocketClient(websocket))
        while True:
            message = await websocket.receive_json()
            if not isinstance(message, list):
                message = [message]
            events = [Event.from_dict(event) for event in message]
            if not all([event is not None for event in events]):
                logger.info(f"Invalid events in: {message}"[:100])
                continue
            for event in events:
                network_system.client_event(session, event)
    except WebSocketDisconnect:
        logger.info(f"Websocket connection closed by client: {session}")
    except Exception as e:
        logger.error("Exception occurred", exc_info=e)
    finally:
        logger.info("Websocket connection closed")
        for entity in session_entities:
            logger.info(f"Removing entity {entity.id} for session {session}")
            entity_manager.remove_entity(entity.id)
            network_system.add_event(Event(EventType.ENTITY_DEATH, {"id": entity.id}))

        network_system.remove_client(session)
        await websocket.close()


@app.get("/healthcheck")
async def healthcheck():
    return Response("healthy")