objects.py

import asyncio
import json
import quart
import os

from babel.languages import get_official_languages
from geoip2 import records

from typing import (
    Optional,
    Dict,
    Any,
)

from utils import (
    DISCORD_API_BASE_URL,
    LANGUAGES,
    LOGGER,
    requests_api
)

class Asset:
    def __init__(self, id: str, key: str):
        self.key: str = key
        self.url: str = f"https://cdn.discordapp.com/avatars/{id}/{key}.webp"

class User:
    def __init__(self, pool, data: Dict):
        self.id: str = data.get("id")
        self.name: str = data.get("global_name")
        self.avatar: Asset = Asset(self.id, data.get("avatar"))
        self.access_token: str = data.get("access_token")
        self.country: records.Country = data.get("country") 
        
        self.bot: Optional[Bot] = None
        self.guild: Optional[Guild] = None
        
        self._pool: UserPool = pool
        self._websocket: Optional[quart.Websocket] = None
    
    async def assign_bot(self, bot) -> None:
        if self.bot:
            if self.id in self.bot._users:
                del self.bot._users[self.id]
            self.bot = None
        
        if self.guild:
            await self.guild.remove_user(self)

        self.bot = bot
        self.bot._users[self.id] = self
        await self.send_to_bot({"op": "initUser"})
        await self.send_to_bot({"op": "initPlayer"})

    async def send_to_bot(self, payload: Dict) -> None:
        method = payload.get("op")
        if method == "heartbeat":
            return
        
        if method == "updateSelectedBot":
            bot = BotPool.get(payload.get("botId"))
            if not bot:
                return await self.send({"op": "botNotFound"})

            return await self.assign_bot(bot)

        elif method == "getMutualGuilds":
            try:
                resp = await requests_api(f'{DISCORD_API_BASE_URL}/users/@me/guilds', headers={'Authorization': f'Bearer {self.access_token}'})
                if resp is None:
                    LOGGER.warning(f"Failed to fetch guilds for user {self.name}({self.id}): Discord API returned non-200 or connection error (possibly expired token).")
                    quart.session.pop("discord_token", None)
                    return await self.send({"op": "errorMsg", "level": "error", "msg": "Su sesión de Discord ha expirado o es inválida. Por favor, cierre sesión e inicie sesión de nuevo."})

                guilds = {
                    guild["id"]: {
                        "avatar": f"https://cdn.discordapp.com/icons/{guild['id']}/{guild['icon']}.webp" if guild.get('icon') else None,
                        "banner": f"https://cdn.discordapp.com/banners/{guild['id']}/{guild['banner']}.webp?size=480&quality=lossless" if guild.get('banner') else None,
                        "name": guild['name']
                    }
                    for guild in resp if int(guild['permissions']) >= 1275593889
                }
                payload["guilds"] = guilds

            except Exception as e:
                LOGGER.error(f"Error fetching mutual guilds for {self.name}({self.id}): {e}")
                return await self.send({"op": "errorMsg", "level": "error", "msg": "Failed to retrieve guild information. Please try again later!"})
        
        payload["userId"] = self.id

        if self.guild:
            # FIX 3: Verificar que el bot esté conectado antes de reenviar
            if not self.guild.bot or not self.guild.bot.is_connected:
                return await self.send({"op": "errorMsg", "level": "error", "msg": "El bot está apagado o reiniciándose. Por favor, inténtelo de nuevo en unos momentos."})
            return await self.guild.send_to_bot(payload)
        
        elif self.bot:
            # FIX 3: Verificar que el bot esté conectado antes de reenviar
            if not self.bot.is_connected:
                return await self.send({"op": "errorMsg", "level": "error", "msg": "El bot está apagado o reiniciándose. Por favor, inténtelo de nuevo en unos momentos."})
            return await self.bot.send(payload)
    
    async def send(self, payload: Dict) -> None:
        if self._websocket:
            try:
                await self._websocket.send_json(payload)
            except Exception as e:
                LOGGER.warning("User (%s) send failed: %s", self.id, e)
            
    async def _listen(self) -> None:
        # Send immediate ping so HF proxy doesn't drop the idle connection
        await self.send({"op": "ping"})

        _no_bot_interval = 5.0
        _last_broadcast = 0.0
        import time as _time

        # Heartbeat task: keeps the WS alive against HF nginx proxy idle timeout
        async def _heartbeat():
            while True:
                await asyncio.sleep(8)
                try:
                    await self.send({"op": "ping"})
                except Exception:
                    break

        hb_task = asyncio.create_task(_heartbeat())
        try:
            while True:
                try:
                    if not self.bot:
                        _now = _time.monotonic()
                        if _now - _last_broadcast >= _no_bot_interval:
                            await BotPool.broadcast({"op": "initBot", "userId": self.id})
                            _last_broadcast = _now

                    data = await self._websocket.receive()
                    await self.send_to_bot(json.loads(data))
                except asyncio.CancelledError:
                    raise
                except Exception as e:
                    err = str(e).lower()
                    if any(x in err for x in ("connection", "disconnect", "closed", "1000", "1001", "1006", "reset", "wsproto")):
                        LOGGER.info("User %s(%s) WebSocket closed: %s", self.name, self.id, e)
                    else:
                        LOGGER.error("Unexpected error in User._listen for %s(%s): %s", self.name, self.id, e, exc_info=True)
                    break
        finally:
            hb_task.cancel()
            try:
                await hb_task
            except asyncio.CancelledError:
                pass

                
    async def connect(self, websocket: quart.Websocket) -> None:
        if self._websocket:
            await self.disconnect()
            
        self._websocket = websocket
                
        LOGGER.info(f"User {self.name}({self.id}) has been connected!")
        received = asyncio.create_task(self._listen())
        await asyncio.gather(received)

    async def disconnect(self) -> None:
        if self._websocket:
            if self.guild:
                try:
                    await self.guild.remove_user(self)
                except Exception as e:
                    LOGGER.warning("User (%s) disconnect: guild remove failed: %s", self.id, e)
            
            self.bot = None
            try:
                await self._websocket.close(1004)
            except Exception:
                pass
            self._websocket = None
            LOGGER.info(f"User {self.name}({self.id}) has been disconnected!")

    @property
    def is_connected(self) -> bool:
        return self._websocket
    
    @property
    def language_code(self) -> str:
        language = get_official_languages(self.country.iso_code if self.country else "US")
        return language[0] if language and language[0] in LANGUAGES else list(LANGUAGES.keys())[0]
        
    def __repr__(self) -> str:
        return f"ID={self.id} Name={self.name}, Guild={self.guild}"
    
class Guild:
    def __init__(self, bot, guild_id: str):
        self.bot: Bot = bot
        self.id: str = guild_id
        
        self._users: Dict[str, User] = {}
    
    async def add_user(self, user: User, init_player: bool = True) -> None:
        if not user.guild:
            user.guild = self
            self._users[user.id] = user

            if init_player:
                await self.bot.send({"op": "initPlayer", "userId": user.id, "guildId": self.id})
        
    async def remove_user(self, user: User) -> None:
        if user.id in self._users:
            await user.send({"op": "playerClose"})
            if len(self._users.keys()) <= 1:
                await user.send_to_bot({"op": "closeConnection", "guildId": self.id})
                
            user.guild = None
            del self._users[user.id]
    
    async def remove_all_user(self) -> None:
        for user_id, user in self._users.copy().items():
            await user.send({"op": "playerClose"})
            user.guild = None
            del self._users[user_id]
            
    async def broadcast(self, payload: Dict) -> None:
        skip_users = payload.get("skip_users", [])
        for user_id, user in self._users.copy().items():
            if user_id not in skip_users:
                await user.send(payload)
    
    async def send_to_bot(self, data: Dict) -> None:
        data["guildId"] = self.id
        await self.bot.send(data)
        
class Bot:
    def __init__(
        self, 
        pool,
        headers: Dict[str, str],
        websocket: quart.Websocket
    ):  
        self.id: str = headers.get("User-Id")
        self._websocket: quart.Websocket = websocket
        self._pool: BotPool = pool
        self._listen_task: Optional[asyncio.Task] = None

        self._guilds: Dict[str, Guild] = {}
        self._users: Dict[str, User] = {}
        self._pending_responses: Dict[str, asyncio.Future] = {}
    
    async def broadcast(self, payload: Dict):
        try:
            for guild in self._guilds.copy().values():
                await guild.broadcast(payload)
        except Exception as e:
            LOGGER.error("Something went wrong while broadcasting to the bot: %s", e, exc_info=True)
    
    async def send(self, payload: Dict) -> None:
        if self.is_connected:
            try:
                LOGGER.debug(f"Bot ({self.id}) sending message: {payload}")
                await self._websocket.send_json(payload)
            except Exception as e:
                LOGGER.warning("Bot (%s) send failed - marking disconnected: %s", self.id, e)
                self._websocket = None
    
    async def request(self, payload: Dict, timeout: float = 5.0) -> Optional[Dict]:
        """Envía un payload y espera la respuesta con el mismo op."""
        op = payload.get("op")
        if not op:
            return None
        loop = asyncio.get_event_loop()
        future = loop.create_future()
        self._pending_responses[op] = future
        try:
            await self.send(payload)
            return await asyncio.wait_for(future, timeout)
        except (asyncio.TimeoutError, Exception):
            return None
        finally:
            self._pending_responses.pop(op, None)
    
    async def _listen(self):
        ws = self._websocket  # Captura referencia local — inmune a reasignaciones concurrentes
        try:
            while True:
                try:
                    if self._websocket is not ws:
                        break  # Este task es obsoleto, el websocket fue reemplazado
                    data = await ws.receive()
                    data: Dict = json.loads(data)
                    LOGGER.debug(f"Bot ({self.id}) receiving message: {data}")
                    
                    method = data.get("op")
                    if not method:
                        continue

                    # Respuestas a peticiones request-response del dashboard (ej: getAdminUsers)
                    if method == "getAdminUsers":
                        future = self._pending_responses.pop("getAdminUsers", None)
                        if future and not future.done():
                            future.set_result(data)
                        continue
                    
                    guild: Optional[Guild] = None  # always initialize before conditional assignment
                    if (guild_id := data.get("guildId")):
                        guild = self.get_guild(guild_id)
                        if not guild:
                            guild = self.create_guild(guild_id)

                        if not guild.bot:
                            guild.bot = self
                    
                        if method == "updateGuild":
                            user: User = UserPool.get(user_id=data.get("user", {}).get("userId"))
                            if user:
                                await guild.add_user(user) if data.get("isJoined") else await guild.remove_user(user)
                        
                        elif method == "createPlayer":
                            for member_id in data.get("memberIds", []):
                                user = UserPool.get(user_id=member_id)
                                if user:
                                    await guild.add_user(user)
                            
                            continue
                        
                        elif method == "initPlayer":
                            user: User = UserPool.get(user_id=data.get("userId"))
                            if user:
                                await guild.add_user(user, init_player=False)

                        elif method == "playerClose":
                            guild = self._guilds.get(data.get("guildId"))
                            if guild:
                                await guild.remove_all_user()

                    if user_id := data.get("userId"):
                        user = UserPool.get(user_id=user_id)
                        if user:
                            await user.send(data)
                        
                    elif guild:
                        await guild.broadcast(data)

                except Exception as e:
                    LOGGER.error("Unexpected error in Bot._listen for bot %s: %s", self.id, e, exc_info=True)
                    break
        finally:
            # Limpieza: si este websocket sigue siendo el actual, marcar como desconectado
            if self._websocket is ws:
                try:
                    await self.disconnect()
                except Exception:
                    pass
            # Remover del pool para que los sockets muertos no rompan broadcasts
            if BotPool.get(self.id) is self:
                BotPool._bots.pop(self.id, None)

    async def disconnect(self) -> None:
        if self._websocket:
            try:
                await self._websocket.close(1004)
            except Exception:
                pass
            self._websocket = None
            
            for guild in self._guilds.values():
                await guild.remove_all_user()
            
            for user in self._users.values():
                user.bot = None
                user.guild = None
                await user.send({"op": "closeConnection"})

            self._guilds = {}
            self._users = {}

            LOGGER.info(f"Bot ({self.id}) has been disconnected!")

    def create_guild(self, guild_id: str) -> Guild:
        if guild_id in self._guilds:
            raise Exception("Guild already exists!")
        
        guild = Guild(self, guild_id)
        self._guilds[guild_id] = guild
        
        return guild
    
    def get_guild(self, guild_id: str) -> Optional[Guild]:
        return self._guilds.get(guild_id)
    
    @property
    def is_connected(self) -> bool:
        return self._websocket is not None
    
class BotPool:
    _bots: Dict[str, Bot] = {}
    
    @classmethod
    async def create(cls, bot_id: str, websocket: quart.Websocket) -> None:
        bot: Optional[Bot] = None
        try:
            header = websocket.headers
            
            bot = cls.get(bot_id)
            if bot:
                # Cancelar el task de escucha anterior ANTES de desconectar/reasignar
                if bot._listen_task and not bot._listen_task.done():
                    bot._listen_task.cancel()
                    try:
                        await bot._listen_task
                    except (asyncio.CancelledError, Exception):
                        pass

                if bot.is_connected:
                    await bot.disconnect()
                bot._websocket = websocket
                # El finally de _listen pudo haber removido al bot del pool; re-registrar
                cls._bots[bot_id] = bot
            else:  
                bot = Bot(cls, header, websocket)
                cls._bots[bot_id] = bot
            
            LOGGER.info(f"Bot ({bot.id}) has been connected!")

            bot._listen_task = asyncio.create_task(bot._listen())
            await bot._listen_task
        except asyncio.CancelledError:
            pass
        except Exception:
            if bot:
                await bot.disconnect()
            raise
    
    @classmethod
    def get(cls, bot_id: str) -> Optional[Bot]:
        return cls._bots.get(bot_id)
    
    @classmethod
    async def broadcast(cls, data: Dict) -> None:
        for bot_id, bot in list(cls._bots.items()):
            if not bot.is_connected:
                cls._bots.pop(bot_id, None)
                continue
            try:
                await bot.send(data)
            except Exception as e:
                LOGGER.error("BotPool.broadcast: error sending to bot %s: %s", bot_id, e)
                try:
                    await bot.disconnect()
                except Exception:
                    pass
                cls._bots.pop(bot_id, None)
            
class UserPool:
    _users: Dict[str, User] = {}
    
    @classmethod
    def add(cls, data: Dict) -> User:
        user = User(cls, data)
        cls._users[user.id] = user
        return user
    
    @classmethod
    def get(cls, *, user_id: str = None, token: str = None) -> Optional[User]:
        if user_id:
            return cls._users.get(user_id)
        
        if token:
            for user in cls._users.values():
                if user.access_token == token:
                    return user

class Settings:
    def __init__(self, settings_file: str = "settings.json"):
        self.settings_file = settings_file
        self.settings = self.load()

        self.host: str = (os.getenv("HOST") or self.get_setting("host") or "").strip()
        self.port: int = int((os.getenv("PORT") or self.get_setting("port") or 5000))
        self.password: str = (os.getenv("PASSWORD") or self.get_setting("password") or "").strip()
        self.client_id: str = (os.getenv("CLIENT_ID") or self.get_setting("client_id") or "").strip()
        self.client_secret_id: str = (os.getenv("CLIENT_SECRET_ID") or self.get_setting("client_secret_id") or "").strip()
        self.secret_key: str = (os.getenv("SECRET_KEY") or self.get_setting("secret_key") or "").strip()
        self.redirect_url: str = (os.getenv("REDIRECT_URL") or self.get_setting("redirect_url") or "").strip()

        self.logging: Dict[str, Any] = self.get_setting("logging")

    def get_setting(self, key: str, default: Optional[str] = None) -> Optional[str]:
        return self.settings.get(key, default)

    def load(self) -> Dict:
        try:
            with open(self.settings_file, "r") as file:
                return json.load(file)
        except FileNotFoundError as e:
            LOGGER.error(f"Unable to load the settings file.", exc_info=e)
            return {}