FabledSteward/fabledscryer/core/plugin_manager.py

# fabledscryer/core/plugin_manager.py
from __future__ import annotations
import hashlib
import importlib
import logging
import os
import sys
import tempfile
import zipfile
from pathlib import Path
from typing import TYPE_CHECKING

import yaml
from packaging.version import Version

if TYPE_CHECKING:
    from quart import Quart

logger = logging.getLogger(__name__)

# Track which plugins were successfully loaded so hot-reload knows
# whether to register a blueprint (new plugin) or skip it (already registered).
_LOADED_PLUGINS: set[str] = set()


def load_plugins(app: "Quart") -> None:
    """Import all enabled plugins and register their blueprints and scheduled tasks.

    For each plugin listed as enabled in config.yaml under 'plugins':
    1. Locate the plugin directory under PLUGIN_DIR
    2. Load and validate plugin.yaml (name must match dir, min_app_version check)
    3. Merge plugin.yaml config defaults with user overrides in app.config["PLUGINS"]
    4. Import the plugin Python package
    5. Validate required exports: setup() and get_scheduled_tasks()
    6. Call setup(app)
    7. Register blueprint at /plugins/<name>/ if get_blueprint() exists
    8. Append get_scheduled_tasks() results to app._task_registry
    """
    import fabledscryer

    plugin_dir = Path(app.config["PLUGIN_DIR"])
    plugins_cfg: dict = app.config["PLUGINS"]

    # Ensure plugin_dir is on sys.path so plugins are importable by name
    plugin_dir_str = str(plugin_dir.resolve())
    if plugin_dir_str not in sys.path:
        sys.path.insert(0, plugin_dir_str)

    for name, cfg in list(plugins_cfg.items()):
        if not cfg.get("enabled", False):
            continue

        plugin_path = plugin_dir / name
        if not plugin_path.exists():
            logger.error("Plugin %r: directory %s not found, skipping", name, plugin_path)
            continue

        # Load and validate plugin.yaml
        yaml_path = plugin_path / "plugin.yaml"
        if not yaml_path.exists():
            logger.error("Plugin %r: plugin.yaml not found, skipping", name)
            continue

        try:
            with yaml_path.open() as f:
                meta = yaml.safe_load(f) or {}
        except Exception:
            logger.exception("Plugin %r: failed to read plugin.yaml, skipping", name)
            continue

        if meta.get("name") != name:
            logger.error(
                "Plugin %r: plugin.yaml name=%r does not match directory name, skipping",
                name, meta.get("name"),
            )
            continue

        min_ver = meta.get("min_app_version")
        if min_ver:
            try:
                if Version(fabledscryer.__version__) < Version(min_ver):
                    logger.error(
                        "Plugin %r: requires fabledscryer>=%s but running %s, skipping",
                        name, min_ver, fabledscryer.__version__,
                    )
                    continue
            except Exception:
                logger.exception("Plugin %r: invalid min_app_version %r, skipping", name, min_ver)
                continue

        # Merge plugin.yaml config defaults with user overrides (non-"enabled" keys)
        defaults = meta.get("config", {})
        user_overrides = {k: v for k, v in cfg.items() if k != "enabled"}
        plugins_cfg[name] = {"enabled": True, **defaults, **user_overrides}

        # Import the plugin package
        try:
            module = importlib.import_module(name)
        except ImportError:
            logger.exception("Plugin %r: import failed, skipping", name)
            continue

        # Validate required exports
        if not hasattr(module, "setup"):
            logger.error("Plugin %r: missing required export 'setup', skipping", name)
            continue
        if not hasattr(module, "get_scheduled_tasks"):
            logger.error(
                "Plugin %r: missing required export 'get_scheduled_tasks', skipping", name
            )
            continue

        # Call setup
        try:
            module.setup(app)
        except Exception:
            logger.exception("Plugin %r: setup() raised, skipping", name)
            continue

        # Register blueprint (optional)
        if hasattr(module, "get_blueprint"):
            try:
                bp = module.get_blueprint()
                app.register_blueprint(bp, url_prefix=f"/plugins/{name}")
                logger.info("Plugin %r: blueprint registered at /plugins/%s/", name, name)
            except Exception:
                logger.exception("Plugin %r: get_blueprint() raised", name)

        # Register scheduled tasks
        try:
            tasks = module.get_scheduled_tasks()
            app._task_registry.extend(tasks)
            logger.info("Plugin %r: registered %d scheduled task(s)", name, len(tasks))
        except Exception:
            logger.exception("Plugin %r: get_scheduled_tasks() raised", name)

        _LOADED_PLUGINS.add(name)
        logger.info("Plugin %r loaded (v%s)", name, meta.get("version", "?"))


async def download_and_install_plugin(
    app: "Quart",
    name: str,
    download_url: str,
    expected_sha256: str,
) -> tuple[bool, str]:
    """Download a plugin zip, verify its checksum, and extract it to PLUGIN_DIR.

    Returns (success, message). Does NOT hot-reload — call hot_reload_plugin()
    after this to activate the plugin without a restart (new plugins only).
    """
    import httpx

    plugin_dir = Path(app.config["PLUGIN_DIR"]).resolve()

    # Ensure plugin_dir is on sys.path (may not be set yet if no plugins were
    # enabled at startup)
    plugin_dir_str = str(plugin_dir)
    if plugin_dir_str not in sys.path:
        sys.path.insert(0, plugin_dir_str)

    try:
        logger.info("Downloading plugin %r from %s", name, download_url)
        async with httpx.AsyncClient(timeout=60, follow_redirects=True) as client:
            resp = await client.get(download_url)
            resp.raise_for_status()
        content = resp.content
    except Exception as exc:
        msg = f"Download failed: {exc}"
        logger.error("Plugin %r: %s", name, msg)
        return False, msg

    # Verify checksum if provided
    if expected_sha256:
        actual = hashlib.sha256(content).hexdigest()
        if actual != expected_sha256.lower():
            msg = f"Checksum mismatch: expected {expected_sha256}, got {actual}"
            logger.error("Plugin %r: %s", name, msg)
            return False, msg

    # Extract the zip
    try:
        with tempfile.TemporaryDirectory() as tmpdir:
            zip_path = Path(tmpdir) / f"{name}.zip"
            zip_path.write_bytes(content)
            with zipfile.ZipFile(zip_path) as zf:
                members = zf.namelist()
                # Collect unique top-level names (dirs and loose files)
                top_level = {m.split("/")[0] for m in members}

                dest = plugin_dir / name
                dest.mkdir(parents=True, exist_ok=True)

                # Detect a single top-level directory to strip. Handles:
                #   name/               (clean plugin zip)
                #   name-v1.0.0/        (GitHub release tag archive)
                #   fabledscryer-plugins-main/name/  (whole-repo archive)
                # Strategy: if there is exactly one top-level item and it is a
                # directory whose name starts with the plugin name (case-insensitive),
                # strip it. Otherwise extract flat into dest.
                prefix: str | None = None
                if len(top_level) == 1:
                    candidate = list(top_level)[0]
                    if candidate.lower().startswith(name.lower()):
                        prefix = candidate + "/"

                for member in members:
                    if prefix:
                        if not member.startswith(prefix):
                            continue
                        rel = member[len(prefix):]
                    else:
                        rel = member
                    if not rel or rel.endswith("/"):
                        if rel:
                            (dest / rel.rstrip("/")).mkdir(parents=True, exist_ok=True)
                        continue
                    target = dest / rel
                    target.parent.mkdir(parents=True, exist_ok=True)
                    target.write_bytes(zf.read(member))
    except Exception as exc:
        msg = f"Extraction failed: {exc}"
        logger.error("Plugin %r: %s", name, msg)
        return False, msg

    # Run migrations for the newly installed plugin
    try:
        mdir = plugin_dir / name / "migrations"
        if mdir.exists():
            from fabledscryer.core.migration_runner import run_plugin_migrations
            run_plugin_migrations(app.config["DATABASE_URL"], [mdir])
    except Exception:
        logger.exception("Plugin %r: migration failed after install", name)
        return False, "Plugin extracted but migrations failed — check logs"

    logger.info("Plugin %r installed successfully", name)
    return True, "Installed"


def hot_reload_plugin(app: "Quart", name: str) -> tuple[bool, str]:
    """Activate a newly installed plugin without restarting the app.

    This works for plugins that were NOT previously loaded. For plugins that
    are already registered (blueprint already mounted), a full restart is
    required to pick up code changes — use restart_app() in that case.

    Returns (success, message).
    """
    import fabledscryer

    if name in _LOADED_PLUGINS:
        return False, "Plugin already loaded — restart required to apply updates"

    plugin_dir = Path(app.config["PLUGIN_DIR"]).resolve()
    plugin_path = plugin_dir / name

    if not plugin_path.exists():
        return False, f"Plugin directory {plugin_path} not found"

    yaml_path = plugin_path / "plugin.yaml"
    if not yaml_path.exists():
        return False, "plugin.yaml not found"

    try:
        with yaml_path.open() as f:
            meta = yaml.safe_load(f) or {}
    except Exception as exc:
        return False, f"Could not read plugin.yaml: {exc}"

    if meta.get("name") != name:
        return False, f"plugin.yaml name mismatch: expected {name!r}, got {meta.get('name')!r}"

    min_ver = meta.get("min_app_version")
    if min_ver:
        try:
            if Version(fabledscryer.__version__) < Version(min_ver):
                return False, (
                    f"Plugin requires fabledscryer>={min_ver} "
                    f"but running {fabledscryer.__version__}"
                )
        except Exception:
            return False, f"Invalid min_app_version: {min_ver!r}"

    # Merge config defaults with any stored user config
    defaults = meta.get("config", {})
    stored_cfg = app.config.get("PLUGINS", {}).get(name, {})
    user_overrides = {k: v for k, v in stored_cfg.items() if k != "enabled"}
    merged = {"enabled": True, **defaults, **user_overrides}
    app.config["PLUGINS"][name] = merged

    # Import
    try:
        module = importlib.import_module(name)
    except ImportError as exc:
        return False, f"Import failed: {exc}"

    if not hasattr(module, "setup") or not hasattr(module, "get_scheduled_tasks"):
        return False, "Plugin missing required exports (setup, get_scheduled_tasks)"

    try:
        module.setup(app)
    except Exception as exc:
        return False, f"setup() raised: {exc}"

    if hasattr(module, "get_blueprint"):
        try:
            bp = module.get_blueprint()
            app.register_blueprint(bp, url_prefix=f"/plugins/{name}")
        except Exception as exc:
            return False, f"Blueprint registration failed: {exc}"

    try:
        tasks = module.get_scheduled_tasks()
        # De-duplicate: skip tasks whose name already exists in the registry
        existing_names = {t.name for t in app._task_registry}
        new_tasks = [t for t in tasks if t.name not in existing_names]
        app._task_registry.extend(new_tasks)
    except Exception as exc:
        return False, f"get_scheduled_tasks() raised: {exc}"

    _LOADED_PLUGINS.add(name)
    logger.info("Plugin %r hot-reloaded (v%s)", name, meta.get("version", "?"))
    return True, f"Plugin activated (v{meta.get('version', '?')})"


def restart_app() -> None:
    """Replace the current process with a fresh copy (in-app restart).

    Used when a plugin update requires a full restart to take effect.
    In Docker the container will be relaunched by the restart policy.
    """
    logger.info("Initiating in-app restart via os.execv")
    os.execv(sys.executable, [sys.executable] + sys.argv)