from __future__ import annotations
from typing import Any, Literal, Optional, TypedDict
from fastapi import (
    Body,
    Cookie,
    FastAPI,
    Path,
    WebSocket,
    WebSocketDisconnect,
    Request,
    Response,
)
from fastapi.staticfiles import StaticFiles
import secrets
import os
import httpx
import json

from pydantic import BaseModel

from logger import logger

public_url = os.getenv("PUBLIC_URL", "/")
if not public_url.endswith("/"):
    public_url += "/"

app = FastAPI()

logger.info(f"Starting server with public URL: {public_url}")


class LobbyResponse(TypedDict):
    id: str
    name: str
    private: bool


class Session:
    _instances: list[Session] = []
    _save_file = "sessions.json"
    _loaded = False

    def __init__(self, id: str):
        logger.info(f"Instantiating new session {id}")
        self._instances.append(self)
        self.id = id
        self.short = id[:8]
        self.name = ""
        self.lobbies: list[Lobby] = []  # List of lobby IDs this session is in
        self.lobby_peers: dict[
            str, list[str]
        ] = {}  # lobby ID -> list of peer session IDs
        self.ws: WebSocket | None = None
        self.save()

    @classmethod
    def save(cls):
        data: list[dict[str, str | list[LobbyResponse]]] = [
            {
                "id": s.id,
                "name": s.name,
                "lobbies": [
                    {"id": lobby.id, "name": lobby.name, "private": lobby.private}
                    for lobby in s.lobbies
                ],
            }
            for s in cls._instances
        ]
        with open(cls._save_file, "w") as f:
            json.dump(data, f, indent=2)
        logger.info(f"Saved {len(data)} sessions to {cls._save_file}")

    @classmethod
    def load(cls):
        if not os.path.exists(cls._save_file):
            logger.info(f"No session save file found: {cls._save_file}")
            return
        with open(cls._save_file, "r") as f:
            data = json.load(f)
        for sdata in data:
            session = Session(sdata["id"])
            session.name = sdata["name"]
            for lobby in sdata.get("lobbies", []):
                session.lobbies.append(
                    Lobby(
                        name=lobby.get("name"),
                        id=lobby.get("id"),
                        private=lobby.get("private", False),
                    )
                )
            logger.info(
                f"Loaded session {session.getName()} with {len(session.lobbies)} lobbies"
            )
            for lobby in session.lobbies:
                lobbies[lobby.id] = Lobby(
                    name=lobby.name, id=lobby.id
                )  # Ensure lobby exists
        logger.info(f"Loaded {len(data)} sessions from {cls._save_file}")

    @classmethod
    def getSession(cls, id: str) -> Session | None:
        if not cls._loaded:
            cls.load()
            logger.info(f"Loaded {len(cls._instances)} sessions from disk...")
            cls._loaded = True
        for s in cls._instances:
            if s.id == id:
                return s
        return None

    @classmethod
    def isUniqueName(cls, name: str) -> bool:
        if not name:
            return False
        for s in cls._instances:
            if s.name.lower() == name.lower():
                return False
        return True

    def getName(self) -> str:
        return f"{self.short}:{self.name if self.name else unset_label}"

    def setName(self, name: str):
        self.name = name
        self.save()

    async def join(self, lobby: Lobby):
        if not self.ws:
            logger.error(
                f"{self.getName()} - No WebSocket connection. Lobby not available."
            )
            return

        if lobby.id in self.lobby_peers or self.id in lobby.sessions:
            logger.info(f"{self.getName()} - Already joined to {lobby.getName()}.")
            await self.ws.send_json(
                {
                    "type": "join_status",
                    "status": "Joined",
                    "message": f"Already joined to lobby {lobby.getName()}",
                }
            )
            return

        # Initialize the peer list for this lobby
        self.lobbies.append(lobby)
        self.lobby_peers[lobby.id] = []

        for peer_id in lobby.sessions:
            if peer_id == self.id:
                raise Exception(
                    "Should not happen: self in lobby.sessions while not in lobby."
                )

            peer_session = lobby.getSession(peer_id)
            if not peer_session or not peer_session.ws:
                logger.warning(
                    f"{self.getName()} - Live peer session {peer_id} not found in lobby {lobby.getName()}. Removing."
                )
                del lobby.sessions[peer_id]
                continue

            # Add the peer to session's RTC peer list
            self.lobby_peers[lobby.id].append(peer_id)

            # Add this user as an RTC peer to each existing peer
            peer_session.lobby_peers[lobby.id].append(self.id)

            logger.info(
                f"{self.getName()} -> {peer_session.getName()}:addPeer({self.getName(), lobby.getName()}, should_create_offer=False)"
            )
            await peer_session.ws.send_json(
                {
                    "type": "addPeer",
                    "data": {
                        "peer_id": self.id,
                        "peer_name": self.name,
                        "should_create_offer": False,
                    },
                }
            )

            # Add each other peer to the caller
            logger.info(
                f"{self.getName()} -> {self.getName()}:addPeer({peer_session.getName(), lobby.getName()}, should_create_offer=True)"
            )
            await self.ws.send_json(
                {
                    "type": "addPeer",
                    "data": {
                        "peer_id": peer_session.id,
                        "peer_name": peer_session.name,
                        "should_create_offer": True,
                    },
                }
            )

        # Add this user as an RTC peer
        await lobby.addSession(self)
        Session.save()

        await self.ws.send_json({"type": "join_status", "status": "Joined"})

    async def part(self, lobby: Lobby):
        if lobby.id not in self.lobby_peers or self.id not in lobby.sessions:
            logger.info(
                f"{self.getName()} - Attempt to part non-joined lobby {lobby.getName()}."
            )
            if self.ws:
                await self.ws.send_json(
                    {"type": "error", "error": "Attempt to part non-joined lobby"}
                )
            return

        logger.info(f"{self.getName()} <- part({lobby.getName()}) - Lobby part.")

        lobby_peers = self.lobby_peers[lobby.id]
        del self.lobby_peers[lobby.id]
        self.lobbies.remove(lobby)

        # Remove this peer from all other RTC peers, and remove each peer from this peer
        for peer_session_id in lobby_peers:
            peer_session = getSession(peer_session_id)
            if not peer_session:
                logger.warning(
                    f"{self.getName()} <- part({lobby.getName()}) - Peer session {peer_session_id} not found. Skipping."
                )
                continue
            if not peer_session.ws:
                logger.warning(
                    f"{self.getName()} <- part({lobby.getName()}) - No WebSocket connection for {peer_session.getName()}. Skipping."
                )
                continue
            logger.info(f"{peer_session.getName()} <- remove_peer({self.getName()})")
            await peer_session.ws.send_json(
                {
                    "type": "removePeer",
                    "data": {"peer_name": self.name, "peer_id": self.id},
                }
            )

            if not self.ws:
                logger.error(
                    f"{self.getName()} <- part({lobby.getName()}) - No WebSocket connection."
                )
                continue

            logger.info(f"{self.getName()} <- remove_peer({peer_session.getName()})")

            await self.ws.send_json(
                {
                    "type": "removePeer",
                    "data": {
                        "peer_name": peer_session.name,
                        "peer_id": peer_session.id,
                    },
                }
            )

        await lobby.removeSession(self)
        Session.save()


class Lobby:
    def __init__(self, name: str, id: str | None = None, private: bool = False):
        self.id = secrets.token_hex(16) if id is None else id
        self.short = self.id[:8]
        self.name = name
        self.sessions: dict[str, Session] = {}  # All lobby members
        self.private = private

    def getName(self) -> str:
        return f"{self.short}:{self.name}"

    async def update_state(self, requesting_session: Session | None = None):
        users: list[dict[str, str | bool]] = [
            {"name": s.name, "live": True if s.ws else False, "session_id": s.id}
            for s in self.sessions.values()
            if s.name
        ]
        if requesting_session:
            logger.info(
                f"{requesting_session.getName()} -> lobby_state({self.getName()})"
            )
            if requesting_session.ws:
                await requesting_session.ws.send_json(
                    {"type": "lobby_state", "data": {"participants": users}}
                )
            else:
                logger.warning(
                    f"{requesting_session.getName()} - No WebSocket connection."
                )
        else:
            for s in self.sessions.values():
                logger.info(f"{s.getName()} -> lobby_state({self.getName()})")
                if s.ws:
                    await s.ws.send_json(
                        {"type": "lobby_state", "data": {"participants": users}}
                    )

    def getSession(self, id: str) -> Session | None:
        return self.sessions.get(id, None)

    async def addSession(self, session: Session) -> None:
        if session.id in self.sessions:
            logger.warning(f"{session.getName()} - Already in lobby {self.getName()}.")
            return None
        self.sessions[session.id] = session
        await self.update_state()

    async def removeSession(self, session: Session) -> None:
        if session.id not in self.sessions:
            logger.warning(f"{session.getName()} - Not in lobby {self.getName()}.")
            return None
        del self.sessions[session.id]
        await self.update_state()


def getName(session: Session | None) -> str | None:
    if session and session.name:
        return session.name
    return None


lobbies: dict[str, Lobby] = {}


def getSession(session_id: str) -> Session | None:
    return Session.getSession(session_id)


def getLobby(lobby_id: str) -> Lobby:
    lobby = lobbies.get(lobby_id, None)
    if not lobby:
        logger.error(f"Lobby not found: {lobby_id}")
        raise Exception(f"Lobby not found: {lobby_id}")
    return lobby


def getLobbyByName(lobby_name: str) -> Lobby | None:
    for lobby in lobbies.values():
        if lobby.name == lobby_name:
            return lobby
    return None


# API endpoints
@app.get(f"{public_url}api/health")
def health():
    logger.info("Health check endpoint called.")
    return {
        "status": "ok",
    }


# A session (cookie) is bound to a single user (name).
# A user can be in multiple lobbies, but a session is unique to a single user.
# A user can change their name, but the session ID remains the same and the name
# updates for all lobbies.
@app.get(f"{public_url}api/session")
async def session(
    request: Request, response: Response, session_id: str | None = Cookie(default=None)
) -> dict[str, str | list[LobbyResponse]]:
    if session_id is None:
        session_id = secrets.token_hex(16)
        response.set_cookie(key="session_id", value=session_id)
    # Validate that session_id is a hex string of length 32
    elif len(session_id) != 32 or not all(c in "0123456789abcdef" for c in session_id):
        return {"error": "Invalid session_id"}

    print(f"[{session_id[:8]}]: Browser hand-shake achieved.")

    session = getSession(session_id)
    if not session:
        session = Session(session_id)
        logger.info(f"{session.getName()}: New session created.")
    else:
        logger.info(f"{session.getName()}: Existing session resumed.")
        # Part all lobbies for this session that have no active websocket
        for lobby_id in list(session.lobby_peers.keys()):
            lobby = None
            try:
                lobby = getLobby(lobby_id)
            except Exception as e:
                logger.error(
                    f"{session.getName()} - Error getting lobby {lobby_id}: {e}"
                )
                continue
            await session.part(lobby)

    return {
        "id": session_id,
        "name": session.name if session.name else "",
        "lobbies": [
            {"id": lobby.id, "name": lobby.name, "private": lobby.private}
            for lobby in session.lobbies
        ],
    }


@app.get(public_url + "api/lobby")
async def get_lobbies(request: Request, response: Response):
    return {
        "lobbies": [
            {"id": lobby.id, "name": lobby.name}
            for lobby in lobbies.values()
            if not lobby.private
        ]
    }


class LobbyCreateData(BaseModel):
    name: str
    private: Optional[bool] = False


class LobbyCreateRequest(BaseModel):
    type: Literal["lobby_create"]
    data: LobbyCreateData


@app.post(public_url + "api/lobby/{session_id}")
async def lobby_create(
    request: Request,
    response: Response,
    session_id: str = Path(...),
    create_request: LobbyCreateRequest = Body(...),
) -> dict[str, str | dict[str, str | bool | int]]:
    if create_request.type != "lobby_create":
        return {"error": "Invalid request type"}

    data = create_request.data
    session = getSession(session_id)
    if not session:
        return {"error": f"Session not found ({session_id})"}
    logger.info(
        f"{session.getName()} lobby_create: {data.name} (private={data.private})"
    )

    lobby = getLobbyByName(data.name)
    if not lobby:
        lobby = Lobby(
            data.name,
            private=data.private if data.private is not None else False,
        )
        lobbies[lobby.id] = lobby
        logger.info(f"{session.getName()} <- lobby_create({lobby.short}:{lobby.name})")

    return {
        "type": "lobby_created",
        "data": {
            "id": lobby.id,
            "name": lobby.name,
            "private": lobby.private,
        },
    }


all_label = "[   all  ]"
info_label = "[  info  ]"
todo_label = "[  todo  ]"
unset_label = "[  ----  ]"


# Register websocket endpoint directly on app with full public_url path
@app.websocket(f"{public_url}" + "ws/lobby/{lobby_id}/{session_id}")
async def lobby_join(
    websocket: WebSocket,
    lobby_id: str | None = Path(...),
    session_id: str | None = Path(...),
):
    await websocket.accept()
    if lobby_id is None:
        await websocket.send_json(
            {"type": "error", "error": "Invalid or missing lobby"}
        )
        await websocket.close()
        return
    if session_id is None:
        await websocket.send_json(
            {"type": "error", "error": "Invalid or missing session"}
        )
        await websocket.close()
        return
    session = getSession(session_id)
    if not session:
        # logger.error(f"Invalid session ID {session_id}")
        await websocket.send_json(
            {"type": "error", "error": f"Invalid session ID {session_id}"}
        )
        await websocket.close()
        return

    lobby = None
    try:
        lobby = getLobby(lobby_id)
    except Exception as e:
        await websocket.send_json({"type": "error", "error": str(e)})
        await websocket.close()
        return

    logger.info(f"{session.getName()} <- lobby_joined({lobby.getName()})")

    session.ws = websocket
    if session.id in lobby.sessions:
        logger.info(
            f"{session.getName()} - Stale session in lobby {lobby.getName()}. Re-joining."
        )
        await session.part(lobby)
        await lobby.removeSession(session)

    for peer_id in lobby.sessions:
        peer_session = lobby.getSession(peer_id)
        if not peer_session or not peer_session.ws:
            logger.warning(
                f"{session.getName()} - Live peer session {peer_id} not found in lobby {lobby.getName()}. Removing."
            )
            del lobby.sessions[peer_id]
            continue
        logger.info(f"{session.getName()} -> user_joined({peer_session.getName()})")
        await peer_session.ws.send_json(
            {
                "type": "user_joined",
                "data": {
                    "session_id": session.id,
                    "name": session.name,
                },
            }
        )

    try:
        while True:
            message = await websocket.receive_json()
            type = message.get("type", None)
            data: dict[str, Any] | None = message.get("data", None)
            if not type:
                logger.error(f"{session.getName()} - Invalid request: {message}")
                await websocket.send_json({"type": "error", "error": "Invalid request"})
                continue
            # logger.info(f"{session.getName()} <- RAW Rx: {data}")
            match type:
                case "set_name":
                    if not data:
                        logger.error(f"{session.getName()} - set_name missing data")
                        await websocket.send_json(
                            {"type": "error", "error": "set_name missing data"}
                        )
                        continue
                    name = data.get("name")
                    logger.info(f"{session.getName()} <- set_name({name})")
                    if not name:
                        logger.error(f"{session.getName()} - Name required")
                        await websocket.send_json(
                            {"type": "error", "error": "Name required"}
                        )
                        continue
                    # Check for duplicate name
                    if not Session.isUniqueName(name):
                        logger.warning(f"{session.getName()} - Name already taken")
                        await websocket.send_json(
                            {"type": "error", "error": "Name already taken"}
                        )
                        continue

                    session.setName(name)
                    logger.info(f"{session.getName()}: -> update('name', {name})")
                    await websocket.send_json({"type": "update", "name": name})
                    # For any clients in any lobby with this session, update their user lists
                    await lobby.update_state()

                case "list_users":
                    await lobby.update_state(session)

                case "join":
                    logger.info(f"{session.getName()} <- join({lobby.getName()})")
                    await session.join(lobby=lobby)

                case "part":
                    logger.info(f"{session.getName()} <- part {lobby.getName()}")
                    await session.part(lobby=lobby)

                case "relayICECandidate":
                    logger.info(f"{session.getName()} <- relayICECandidate")
                    if not data:
                        logger.error(
                            f"{session.getName()} - relayICECandidate missing data"
                        )
                        await websocket.send_json(
                            {"type": "error", "error": "relayICECandidate missing data"}
                        )
                        continue

                    if (
                        lobby.id not in session.lobby_peers
                        or session.id not in lobby.sessions
                    ):
                        logger.error(
                            f"{session.short}:{session.name} <- relayICECandidate - Not an RTC peer ({session.id})"
                        )
                        await websocket.send_json(
                            {"type": "error", "error": "Not joined to lobby"}
                        )
                        continue
                    session_peers = session.lobby_peers[lobby.id]
                    peer_id = data.get("peer_id")
                    if peer_id not in session_peers:
                        logger.error(
                            f"{session.getName()} <- relayICECandidate - Not an RTC peer({peer_id}) in {session_peers}"
                        )
                        await websocket.send_json(
                            {
                                "type": "error",
                                "error": f"Target peer {peer_id} not found",
                            }
                        )
                        continue

                    candidate = data.get("candidate")

                    message: dict[str, Any] = {
                        "type": "iceCandidate",
                        "data": {
                            "peer_id": session.id,
                            "peer_name": session.name,
                            "candidate": candidate,
                        },
                    }

                    peer_session = lobby.getSession(peer_id)
                    if not peer_session or not peer_session.ws:
                        logger.warning(
                            f"{session.getName()} - Live peer session {peer_id} not found in lobby {lobby.getName()}."
                        )
                        break
                    logger.info(
                        f"{session.getName()} -> iceCandidate({peer_session.getName()})"
                    )
                    await peer_session.ws.send_json(message)

                case "relaySessionDescription":
                    logger.info(f"{session.getName()} <- relaySessionDescription")
                    if not data:
                        logger.error(
                            f"{session.getName()} - relaySessionDescription missing data"
                        )
                        await websocket.send_json(
                            {
                                "type": "error",
                                "error": "relaySessionDescription missing data",
                            }
                        )
                        continue

                    if (
                        lobby.id not in session.lobby_peers
                        or session.id not in lobby.sessions
                    ):
                        logger.error(
                            f"{session.short}:{session.name} <- relaySessionDescription - Not an RTC peer ({session.id})"
                        )
                        await websocket.send_json(
                            {"type": "error", "error": "Not joined to lobby"}
                        )
                        continue

                    lobby_peers = session.lobby_peers[lobby.id]
                    peer_id = data.get("peer_id")
                    if peer_id not in lobby_peers:
                        logger.error(
                            f"{session.getName()} <- relaySessionDescription - Not an RTC peer({peer_id}) in {lobby_peers}"
                        )
                        await websocket.send_json(
                            {
                                "type": "error",
                                "error": f"Target peer {peer_id} not found",
                            }
                        )
                        continue

                    peer_id = data.get("peer_id", None)
                    if not peer_id:
                        logger.error(
                            f"{session.getName()} - relaySessionDescription missing peer_id"
                        )
                        await websocket.send_json(
                            {
                                "type": "error",
                                "error": "relaySessionDescription missing peer_id",
                            }
                        )
                        continue
                    peer_session = lobby.getSession(peer_id)
                    if not peer_session or not peer_session.ws:
                        logger.warning(
                            f"{session.getName()} - Live peer session {peer_id} not found in lobby {lobby.getName()}."
                        )
                        break

                    session_description = data.get("session_description")
                    message = {
                        "type": "sessionDescription",
                        "data": {
                            "peer_id": session.id,
                            "peer_name": session.name,
                            "session_description": session_description,
                        },
                    }

                    logger.info(
                        f"{session.getName()} -> sessionDescription({peer_session.getName()})"
                    )
                    await peer_session.ws.send_json(message)

                case _:
                    await websocket.send_json(
                        {
                            "type": "error",
                            "error": f"Unknown request type: {type}",
                        }
                    )

    except WebSocketDisconnect:
        logger.info(f"{session.getName()} <- WebSocket disconnected for user.")
        # Cleanup: remove session from lobby and sessions dict
        session.ws = None
        if session.id in lobby.sessions:
            await session.part(lobby)

        await lobby.update_state()

        # Clean up empty lobbies
        if not lobby.sessions:
            if lobby.id in lobbies:
                del lobbies[lobby.id]
                logger.info(f"Cleaned up empty lobby {lobby.getName()}")


# Serve static files or proxy to frontend development server
PRODUCTION = os.getenv("PRODUCTION", "false").lower() == "true"
client_build_path = os.path.join(os.path.dirname(__file__), "/client/build")

if PRODUCTION:
    logger.info(f"Serving static files from: {client_build_path} at {public_url}")
    app.mount(
        public_url, StaticFiles(directory=client_build_path, html=True), name="static"
    )


else:
    logger.info(f"Proxying static files to http://static-frontend:3000 at {public_url}")

    import ssl

    @app.api_route(
        f"{public_url}{{path:path}}",
        methods=["GET", "POST", "PUT", "DELETE", "OPTIONS", "HEAD", "PATCH"],
    )
    async def proxy_static(request: Request, path: str):
        # Do not proxy API or websocket paths
        if path.startswith("api/") or path.startswith("ws/"):
            return Response(status_code=404)
        url = f"{request.url.scheme}://static-frontend:3000/{public_url.strip('/')}/{path}"
        if not path:
            url = f"{request.url.scheme}://static-frontend:3000/{public_url.strip('/')}"
        headers = dict(request.headers)
        try:
            # Accept self-signed certs in dev
            async with httpx.AsyncClient(verify=False) as client:
                proxy_req = client.build_request(
                    request.method, url, headers=headers, content=await request.body()
                )
                proxy_resp = await client.send(proxy_req, stream=True)
                content = await proxy_resp.aread()

                # Remove problematic headers for browser decoding
                filtered_headers = {
                    k: v
                    for k, v in proxy_resp.headers.items()
                    if k.lower()
                    not in ["content-encoding", "transfer-encoding", "content-length"]
                }
                return Response(
                    content=content,
                    status_code=proxy_resp.status_code,
                    headers=filtered_headers,
                )
        except Exception as e:
            logger.error(f"Proxy error for {url}: {e}")
            return Response("Proxy error", status_code=502)

    # WebSocket proxy for /ws (for React DevTools, etc.)
    import websockets
    import asyncio
    from starlette.websockets import WebSocket as StarletteWebSocket

    @app.websocket("/ws")
    async def websocket_proxy(websocket: StarletteWebSocket):
        logger.info("REACT: WebSocket proxy connection established.")
        # Get scheme from websocket.url (should be 'ws' or 'wss')
        scheme = websocket.url.scheme if hasattr(websocket, "url") else "ws"
        target_url = f"{scheme}://static-frontend:3000/ws"
        await websocket.accept()
        try:
            # Accept self-signed certs in dev for WSS
            ssl_ctx = ssl.create_default_context()
            ssl_ctx.check_hostname = False
            ssl_ctx.verify_mode = ssl.CERT_NONE
            async with websockets.connect(target_url, ssl=ssl_ctx) as target_ws:

                async def client_to_server():
                    while True:
                        msg = await websocket.receive_text()
                        await target_ws.send(msg)

                async def server_to_client():
                    while True:
                        msg = await target_ws.recv()
                        if isinstance(msg, str):
                            await websocket.send_text(msg)
                        else:
                            await websocket.send_bytes(msg)

                try:
                    await asyncio.gather(client_to_server(), server_to_client())
                except (WebSocketDisconnect, websockets.ConnectionClosed):
                    logger.info("REACT: WebSocket proxy connection closed.")
        except Exception as e:
            logger.error(f"REACT: WebSocket proxy error: {e}")
            await websocket.close()