from __future__ import annotations

import json
from datetime import datetime, timezone

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from ..chat.session import ChatSessionManager
from ..state import StateManager
from .constants import logger

router = APIRouter()


@router.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket) -> None:
    connected_clients: list[WebSocket] = websocket.app.state.connected_clients
    state_manager: StateManager = websocket.app.state.state_manager

    await websocket.accept()
    connected_clients.append(websocket)
    await state_manager.update_web_client_count(len(connected_clients))

    try:
        await websocket.send_text(
            json.dumps(
                {
                    "type": "state_update",
                    "data": state_manager.state.to_dict(),
                    "timestamp": datetime.now(tz=timezone.utc).isoformat(),
                }
            )
        )

        while True:
            _data = await websocket.receive_text()
    except WebSocketDisconnect:
        pass
    except OSError as e:
        logger.debug("WebSocket error: %s", e)
    finally:
        if websocket in connected_clients:
            connected_clients.remove(websocket)
        await state_manager.update_web_client_count(len(connected_clients))


@router.websocket("/ws/sessions/{session_id}/terminal")
async def session_terminal_websocket(websocket: WebSocket, session_id: str) -> None:
    session_manager: ChatSessionManager = websocket.app.state.session_manager
    session = session_manager.get_session(session_id)
    if not session:
        await websocket.close(code=4004, reason="Session not found")
        return

    await websocket.accept()

    query_params = dict(websocket.query_params)
    can_resize = query_params.get("resize", "false").lower() == "true"

    async def output_callback(data: bytes) -> None:
        try:
            await websocket.send_bytes(data)
        except WebSocketDisconnect:
            raise
        except OSError as e:
            logger.warning("Error sending bytes to terminal WebSocket: %s", e)

    try:
        history = session.attach_binary_callback(output_callback, can_resize=can_resize)

        await websocket.send_text(
            json.dumps(
                {
                    "type": "pty_size",
                    "rows": session.term_rows,
                    "cols": session.term_cols,
                }
            )
        )

        if history:
            await websocket.send_bytes(history)

        while True:
            message = await websocket.receive()
            msg_type = message.get("type", "")

            if msg_type == "websocket.disconnect":
                break

            if "bytes" in message and message["bytes"] is not None:
                await session.send_input(message["bytes"], source="ws")
            elif "text" in message and message["text"] is not None:
                text = message["text"]
                try:
                    ctrl_data = json.loads(text)
                    ctrl_type = ctrl_data.get("type", "")

                    if ctrl_type == "resize" and can_resize:
                        rows = ctrl_data.get("rows", 24)
                        cols = ctrl_data.get("cols", 80)
                        session.resize(rows, cols, requester=output_callback)
                except json.JSONDecodeError:
                    await session.send_input(text.encode("utf-8"), source="ws")

    except WebSocketDisconnect:
        logger.debug("Terminal WebSocket disconnected normally")
    except OSError as e:
        logger.warning("Terminal WebSocket error: %s", e, exc_info=True)
    finally:
        session.detach_binary_callback(output_callback)

