plugin-helper/src/plugin_helper/models.py

from __future__ import annotations

import tomllib
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Any


VALID_STRATEGIES = {"dll-to-plugins", "zip-to-pending", "bsipa-zip", "root-zip", "manual"}


@dataclass(frozen=True)
class Dependency:
    id: str
    constraint: str | None = None
    required: bool = True


@dataclass(frozen=True)
class RegistryPlugin:
    id: str
    name: str
    repo: str | None
    asset_patterns: tuple[str, ...] = ()
    install_strategy: str = "manual"
    category: str | None = None
    dependencies: tuple[Dependency, ...] = ()


@dataclass(frozen=True)
class Registry:
    plugins: dict[str, RegistryPlugin] = field(default_factory=dict)

    def get(self, plugin_id: str) -> RegistryPlugin | None:
        return self.plugins.get(plugin_id)


@dataclass(frozen=True)
class LockedPlugin:
    id: str
    repo: str | None
    tag: str | None
    asset: str | None
    sha256: str | None
    download_url: str | None = None
    install_strategy: str | None = None
    reason: str | None = None


@dataclass(frozen=True)
class Lockfile:
    beat_saber_version: str
    instance: str
    plugins: tuple[LockedPlugin, ...]


def _quote_toml_string(value: str) -> str:
    return '"' + value.replace("\\", "\\\\").replace('"', '\\"') + '"'


def lockfile_to_toml(lockfile: Lockfile) -> str:
    lines = [
        f"beat_saber_version = {_quote_toml_string(lockfile.beat_saber_version)}",
        f"instance = {_quote_toml_string(lockfile.instance)}",
        "",
    ]
    for plugin in lockfile.plugins:
        lines.append("[[plugins]]")
        lines.append(f"id = {_quote_toml_string(plugin.id)}")
        if plugin.repo is not None:
            lines.append(f"repo = {_quote_toml_string(plugin.repo)}")
        if plugin.tag is not None:
            lines.append(f"tag = {_quote_toml_string(plugin.tag)}")
        if plugin.asset is not None:
            lines.append(f"asset = {_quote_toml_string(plugin.asset)}")
        if plugin.download_url is not None:
            lines.append(f"download_url = {_quote_toml_string(plugin.download_url)}")
        if plugin.sha256 is not None:
            lines.append(f"sha256 = {_quote_toml_string(plugin.sha256)}")
        if plugin.install_strategy is not None:
            lines.append(f"install_strategy = {_quote_toml_string(plugin.install_strategy)}")
        if plugin.reason is not None:
            lines.append(f"reason = {_quote_toml_string(plugin.reason)}")
        lines.append("")
    return "\n".join(lines)


def write_lockfile(path: Path, lockfile: Lockfile) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(lockfile_to_toml(lockfile), encoding="utf-8")


def replace_locked_plugins(lockfile: Lockfile, replacements: dict[str, LockedPlugin]) -> Lockfile:
    return replace(
        lockfile,
        plugins=tuple(replacements.get(plugin.id, plugin) for plugin in lockfile.plugins),
    )


def _load_toml(path: Path) -> dict[str, Any]:
    with path.open("rb") as handle:
        return tomllib.load(handle)


def _registry_plugin_from_item(item: dict[str, Any], source: Path) -> RegistryPlugin:
    dependencies = tuple(
        Dependency(
            id=dep["id"],
            constraint=dep.get("constraint"),
            required=dep.get("required", True),
        )
        for dep in item.get("dependencies", [])
    )
    strategy = item.get("install_strategy", "manual")
    if strategy not in VALID_STRATEGIES:
        raise ValueError(f"{source}: invalid install_strategy for {item['id']}: {strategy}")
    return RegistryPlugin(
        id=item["id"],
        name=item.get("name", item["id"]),
        repo=item.get("repo"),
        asset_patterns=tuple(item.get("asset_patterns", [])),
        install_strategy=strategy,
        category=item.get("category"),
        dependencies=dependencies,
    )


def _add_registry_plugin(plugins: dict[str, RegistryPlugin], plugin: RegistryPlugin, source: Path) -> None:
    if plugin.id in plugins:
        raise ValueError(f"{source}: duplicate plugin id: {plugin.id}")
    plugins[plugin.id] = plugin


def load_registry(path: Path) -> Registry:
    if not path.exists():
        legacy_path = path.with_suffix(".toml")
        if path.name == "plugins" and legacy_path.exists():
            return load_registry(legacy_path)
        return Registry()
    plugins: dict[str, RegistryPlugin] = {}
    if path.is_dir():
        for item_path in sorted(path.glob("*.toml")):
            data = _load_toml(item_path)
            _add_registry_plugin(plugins, _registry_plugin_from_item(data, item_path), item_path)
        return Registry(plugins)

    data = _load_toml(path)
    for item in data.get("plugins", []):
        _add_registry_plugin(plugins, _registry_plugin_from_item(item, path), path)
    return Registry(plugins)


def load_lockfile(path: Path) -> Lockfile:
    data = _load_toml(path)
    plugins = tuple(
        LockedPlugin(
            id=item["id"],
            repo=item.get("repo"),
            tag=item.get("tag"),
            asset=item.get("asset"),
            sha256=item.get("sha256"),
            download_url=item.get("download_url"),
            install_strategy=item.get("install_strategy"),
            reason=item.get("reason"),
        )
        for item in data.get("plugins", [])
    )
    return Lockfile(
        beat_saber_version=data["beat_saber_version"],
        instance=data.get("instance", data["beat_saber_version"]),
        plugins=plugins,
    )