ProgressionMods/progression_backend/workshop.py

from __future__ import annotations

import re
import shutil
import subprocess
import time
from math import ceil
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path

RIMWORLD_APP_ID = "294100"
FAILED_ITEM_PATTERN = re.compile(
    r"(?:ERROR!\s*)?Downloading item\s+(\d+)\s+failed", re.IGNORECASE
)


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def directory_size(path: Path) -> int:
    if not path.exists():
        return 0
    total = 0
    for entry in path.rglob("*"):
        try:
            if entry.is_file():
                total += entry.stat().st_size
        except OSError:
            continue
    return total


def find_about_xml(mod_path: Path) -> Path | None:
    direct = mod_path / "About" / "About.xml"
    if direct.is_file():
        return direct
    try:
        for candidate in mod_path.rglob("*.xml"):
            if (
                candidate.name.lower() == "about.xml"
                and candidate.parent.name.lower() == "about"
            ):
                return candidate
    except OSError:
        return None
    return None


def _child_text(root: ET.Element, name: str) -> str | None:
    wanted = name.lower()
    for child in root:
        if child.tag.lower() == wanted and child.text:
            return child.text.strip()
    return None


@dataclass(frozen=True)
class ParsedMod:
    package_id: str
    display_name: str
    declared_steam_app_id: str | None


def parse_about_xml(mod_path: Path) -> ParsedMod:
    about_xml = find_about_xml(mod_path)
    if not about_xml:
        raise ValueError("About/About.xml not found")
    try:
        root = ET.parse(about_xml).getroot()
    except ET.ParseError as error:
        raise ValueError(f"About.xml is invalid XML: {error}") from error

    package_id = _child_text(root, "packageId")
    if not package_id:
        raise ValueError("packageId not found in About.xml")
    display_name = _child_text(root, "name") or package_id
    return ParsedMod(
        package_id=package_id.lower(),
        display_name=display_name,
        declared_steam_app_id=_child_text(root, "steamAppId"),
    )


@dataclass(frozen=True)
class SteamCmdBatchResult:
    workshop_ids: list[str]
    started_at: str
    finished_at: str
    elapsed_seconds: float
    exit_code: int | None
    status: str
    failed_ids: set[str]
    output_tail: str


class SteamCmdRunner:
    def __init__(
        self,
        executable: str,
        content_dir: Path,
        timeout_seconds: int,
    ):
        self.executable = executable
        self.content_dir = content_dir
        self.timeout_seconds = timeout_seconds

    def executable_path(self) -> str | None:
        configured = Path(self.executable)
        if configured.is_file():
            return str(configured.resolve())
        return shutil.which(self.executable)

    def available(self) -> bool:
        return self.executable_path() is not None

    def run_batch(
        self,
        workshop_ids: list[str],
        timeout_seconds: float | None = None,
    ) -> SteamCmdBatchResult:
        executable = self.executable_path()
        if not executable:
            raise FileNotFoundError(
                f"SteamCMD was not found at '{self.executable}'. Set STEAMCMD_PATH."
            )
        invalid = [item for item in workshop_ids if not item.isdigit()]
        if invalid:
            raise ValueError(f"Invalid Workshop ids: {', '.join(invalid)}")

        command = [executable, "+login", "anonymous"]
        for workshop_id in workshop_ids:
            command.extend(
                [
                    "+workshop_download_item",
                    RIMWORLD_APP_ID,
                    workshop_id,
                    "validate",
                ]
            )
        command.append("+quit")

        started_at = utc_now()
        started = time.monotonic()
        exit_code: int | None = None
        status = "completed"
        output = ""
        timeout = self.timeout_seconds
        if timeout_seconds is not None:
            timeout = max(1, min(timeout, ceil(timeout_seconds)))
        try:
            completed = subprocess.run(
                command,
                cwd=str(Path(executable).resolve().parent),
                capture_output=True,
                text=True,
                encoding="utf-8",
                errors="replace",
                timeout=timeout,
                check=False,
                shell=False,
            )
            exit_code = completed.returncode
            output = f"{completed.stdout}\n{completed.stderr}".strip()
            if completed.returncode != 0:
                status = "failed"
        except subprocess.TimeoutExpired as error:
            status = "timeout"
            stdout = (
                error.stdout.decode(errors="replace")
                if isinstance(error.stdout, bytes)
                else error.stdout
            )
            stderr = (
                error.stderr.decode(errors="replace")
                if isinstance(error.stderr, bytes)
                else error.stderr
            )
            output = f"{stdout or ''}\n{stderr or ''}".strip()

        failed_ids = set(FAILED_ITEM_PATTERN.findall(output))
        for workshop_id in workshop_ids:
            if not (self.content_dir / workshop_id).is_dir():
                failed_ids.add(workshop_id)
        if not failed_ids:
            # SteamCMD's bootstrapper can return 1 after updating and relaunching even
            # though every requested item was installed successfully.
            status = "completed"
        elif len(failed_ids) < len(workshop_ids):
            status = "partial"
        else:
            status = "failed"

        return SteamCmdBatchResult(
            workshop_ids=workshop_ids,
            started_at=started_at,
            finished_at=utc_now(),
            elapsed_seconds=round(time.monotonic() - started, 3),
            exit_code=exit_code,
            status=status,
            failed_ids=failed_ids,
            output_tail=output[-8000:],
        )