TurboSu/server/game_manager.py

from __future__ import annotations

import io
import json
import shutil
import zipfile
import importlib.util
import sys
from dataclasses import dataclass, field, asdict
from pathlib import Path
from typing import Any

from utils.logger import get_logger

logger = get_logger(__name__)

GAMES_DIR = Path(__file__).resolve().parent.parent / "games"
BUILTIN_DIR = GAMES_DIR / "builtin"
USER_DIR = GAMES_DIR / "user"


@dataclass
class GamePlugin:
    id: str = ""
    name: str = ""
    parser_type: str = "forza"
    telemetry_format: str = ""
    description: str = ""
    author: str = ""
    version: str = "1.0.0"
    default_port: int = 20777
    icon: str = ""
    is_builtin: bool = True
    enabled: bool = True
    manifest_path: str = ""

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)

    @classmethod
    def from_manifest(cls, path: Path, is_builtin: bool = True) -> GamePlugin | None:
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            gp = cls(
                id=data.get("id", path.parent.name),
                name=data.get("name", path.parent.name),
                parser_type=data.get("parser_type", "forza"),
                telemetry_format=data.get("telemetry_format", ""),
                description=data.get("description", ""),
                author=data.get("author", ""),
                version=data.get("version", "1.0.0"),
                default_port=data.get("default_port", 20777),
                icon=data.get("icon", ""),
                is_builtin=is_builtin,
                enabled=True,
                manifest_path=str(path.parent),
            )
            return gp
        except Exception as e:
            logger.error("Failed to load game plugin manifest %s: %s", path, e)
            return None


class GamePluginManager:
    def __init__(self):
        BUILTIN_DIR.mkdir(parents=True, exist_ok=True)
        USER_DIR.mkdir(parents=True, exist_ok=True)
        self._plugins: dict[str, GamePlugin] = {}
        self._parser_cache: dict[str, Any] = {}
        self._scanned = False
        self._discover()

    def _discover(self):
        self._plugins.clear()
        for d in [BUILTIN_DIR, USER_DIR]:
            if not d.exists():
                continue
            is_builtin = (d == BUILTIN_DIR)
            for item in d.iterdir():
                if item.is_dir():
                    manifest = item / "manifest.json"
                    if manifest.exists():
                        gp = GamePlugin.from_manifest(manifest, is_builtin)
                        if gp:
                            self._plugins[gp.id] = gp
        self._scanned = True

    def _ensure_scanned(self):
        if not self._scanned:
            self._discover()

    def list_all(self) -> list[dict[str, Any]]:
        self._ensure_scanned()
        return [p.to_dict() for p in self._plugins.values()]

    def get(self, plugin_id: str) -> GamePlugin | None:
        self._ensure_scanned()
        return self._plugins.get(plugin_id)

    def get_parser(self, plugin_id: str) -> Any | None:
        self._discover()
        gp = self._plugins.get(plugin_id)
        if not gp:
            return None

        cache_key = gp.id
        if cache_key in self._parser_cache:
            return self._parser_cache[cache_key]

        parser_dir = Path(gp.manifest_path)
        parser_file = parser_dir / "parser.py"
        if not parser_file.exists():
            return None

        try:
            module_name = f"turbosu_game_{gp.id}"
            spec = importlib.util.spec_from_file_location(module_name, parser_file)
            if spec and spec.loader:
                module = importlib.util.module_from_spec(spec)
                sys.modules[module_name] = module
                spec.loader.exec_module(module)
                if hasattr(module, "get_parser"):
                    parser = module.get_parser()
                    self._parser_cache[cache_key] = parser
                    logger.info("Loaded parser for game: %s", gp.name)
                    return parser
        except Exception as e:
            logger.error("Failed to load parser for %s: %s", gp.id, e)

        return None

    def reload(self):
        self._discover()
        self._parser_cache.clear()

    def remove_plugin(self, plugin_id: str) -> bool:
        self._discover()
        gp = self._plugins.get(plugin_id)
        if gp and gp.is_builtin:
            logger.warning("Cannot remove builtin plugin: %s", plugin_id)
            return False
        dest_dir = USER_DIR / plugin_id
        if dest_dir.exists():
            shutil.rmtree(dest_dir)
            self._discover()
            logger.info("Plugin removed: %s", plugin_id)
            return True
        return False

    def install_plugin_zip(self, zip_data: bytes) -> bool:
        try:
            with zipfile.ZipFile(io.BytesIO(zip_data), 'r') as zf:
                names = zf.namelist()
                manifest_name = None
                for n in names:
                    if n.endswith('manifest.json'):
                        manifest_name = n
                        break
                if not manifest_name:
                    logger.error("No manifest.json found in plugin zip")
                    return False

                manifest = json.loads(zf.read(manifest_name).decode('utf-8'))
                plugin_id = manifest.get("id", "")
                if not plugin_id:
                    return False

                dest_dir = USER_DIR / plugin_id
                if dest_dir.exists():
                    shutil.rmtree(dest_dir)
                dest_dir.mkdir(parents=True)

                with open(dest_dir / "manifest.json", "w", encoding="utf-8") as f:
                    json.dump(manifest, f, indent=2, ensure_ascii=False)

                for name in names:
                    if name == manifest_name or name.endswith('/'):
                        continue
                    if name.endswith('parser.py'):
                        out_path = dest_dir / "parser.py"
                        out_path.parent.mkdir(parents=True, exist_ok=True)
                        out_path.write_bytes(zf.read(name))

                self._discover()
                logger.info("Plugin installed from zip: %s", plugin_id)
                return True
        except Exception as e:
            logger.error("Failed to install plugin zip: %s", e)
            return False

    def export_plugin_zip(self, plugin_id: str) -> bytes | None:
        gp = self._plugins.get(plugin_id)
        if not gp:
            return None
        parser_dir = Path(gp.manifest_path)
        if not parser_dir.exists():
            return None
        buf = io.BytesIO()
        with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
            for f in sorted(parser_dir.rglob('*')):
                if f.is_file():
                    zf.write(f, f.relative_to(parser_dir))
        logger.info("Plugin exported as zip: %s", plugin_id)
        return buf.getvalue()


game_plugin_manager = GamePluginManager()