unity_agent/transport/unity_transport.py

"""Unity project file writer.

The :class:`UnityTransport` is the *only* component in the system that touches
the filesystem. Tools generate strings (C# code, YAML scene files, JSON
manifests) and hand them to the transport, which writes them to the correct
location inside the project directory and maintains an inventory of every file
it has written.

The directory layout produced matches what Unity expects when you click
"Open Project" in the Unity Hub::

    ProjectRoot/
        Assets/
            Scripts/<file>.cs
            Scenes/<scene>.unity
        Packages/
            manifest.json
        ProjectSettings/
            ProjectSettings.asset
            ProjectVersion.txt
"""

from __future__ import annotations

import hashlib
import json
import logging
import os
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional

from unity_agent.config import Settings

log = logging.getLogger("unity_agent.transport")


# --------------------------------------------------------------------------- #
# .meta file generation
# --------------------------------------------------------------------------- #
def make_meta(guid: Optional[str] = None) -> str:
    """Return a valid Unity ``.meta`` file body for the given GUID.

    Unity uses ``.meta`` files to track assets across machines. Every meta
    file needs a unique 32-hex-character GUID. We generate one from a hash
    of the current time + random bytes when the caller does not supply one.
    """
    if guid is None:
        seed = f"{time.time_ns()}-{os.urandom(8).hex()}".encode()
        guid = hashlib.md5(seed).hexdigest()[:32]
    return (
        "fileFormatVersion: 2\n"
        f"guid: {guid}\n"
        "MonoImporter:\n"
        "  externalObjects: {}\n"
        "  serializedVersion: 2\n"
        "  defaultReferences: []\n"
        "  executionOrder: 0\n"
        "  icon: {instanceID: 0}\n"
        "  userData: \n"
        "  assetBundleName: \n"
        "  assetBundleVariant: \n"
    )


# --------------------------------------------------------------------------- #
# File record
# --------------------------------------------------------------------------- #
@dataclass
class WrittenFile:
    relative_path: str  # path relative to project root
    absolute_path: Path
    size_bytes: int
    kind: str  # "csharp", "scene", "config", "meta", "other"


@dataclass
class TransportStats:
    files: List[WrittenFile] = field(default_factory=list)

    def add(self, wf: WrittenFile) -> None:
        self.files.append(wf)

    @property
    def total_bytes(self) -> int:
        return sum(f.size_bytes for f in self.files)

    def by_kind(self, kind: str) -> List[WrittenFile]:
        return [f for f in self.files if f.kind == kind]


# --------------------------------------------------------------------------- #
# Transport
# --------------------------------------------------------------------------- #
class UnityTransport:
    """Filesystem writer for Unity projects.

    Parameters
    ----------
    settings:
        Global :class:`Settings`. The transport uses ``output_dir`` as the
        parent for all projects.
    project_name:
        Optional project name. If provided, all paths are resolved relative
        to ``<output_dir>/<project_name>``. If ``None`` you must call
        :meth:`open_project` before writing any files.
    """

    ASSETS_DIR = "Assets"
    SCRIPTS_DIR = "Assets/Scripts"
    SCENES_DIR = "Assets/Scenes"
    PACKAGES_DIR = "Packages"
    PROJECT_SETTINGS_DIR = "ProjectSettings"

    def __init__(self, settings: Settings, project_name: Optional[str] = None) -> None:
        self.settings = settings
        self.project_name = project_name
        self.stats = TransportStats()
        self._project_root: Optional[Path] = None
        if project_name is not None:
            self.open_project(project_name)

    # ------------------------------------------------------------------ #
    # Project lifecycle
    # ------------------------------------------------------------------ #
    def open_project(self, project_name: str) -> Path:
        """Create / open a project directory and return its absolute path."""
        self.project_name = project_name
        self._project_root = self.settings.project_path(project_name)
        if self.settings.auto_create_dirs:
            for sub in (
                self.ASSETS_DIR,
                self.SCRIPTS_DIR,
                self.PACKAGES_DIR,
                self.PROJECT_SETTINGS_DIR,
            ):
                (self._project_root / sub).mkdir(parents=True, exist_ok=True)
        log.info("Opened Unity project at %s", self._project_root)
        return self._project_root

    @property
    def project_root(self) -> Path:
        if self._project_root is None:
            raise RuntimeError("No project open. Call open_project() first.")
        return self._project_root

    # ------------------------------------------------------------------ #
    # Low-level write
    # ------------------------------------------------------------------ #
    def _write(
        self,
        relative_path: str,
        content: str,
        kind: str,
        *,
        write_meta: bool = False,
    ) -> WrittenFile:
        full = self.project_root / relative_path
        full.parent.mkdir(parents=True, exist_ok=True)
        if full.exists() and not self.settings.overwrite:
            raise FileExistsError(full)
        full.write_text(content, encoding="utf-8")
        wf = WrittenFile(
            relative_path=relative_path,
            absolute_path=full,
            size_bytes=len(content.encode("utf-8")),
            kind=kind,
        )
        self.stats.add(wf)
        if write_meta:
            meta_path = relative_path + ".meta"
            meta_content = make_meta()
            meta_full = self.project_root / meta_path
            meta_full.write_text(meta_content, encoding="utf-8")
            self.stats.add(
                WrittenFile(
                    relative_path=meta_path,
                    absolute_path=meta_full,
                    size_bytes=len(meta_content),
                    kind="meta",
                )
            )
        return wf

    # ------------------------------------------------------------------ #
    # High-level helpers
    # ------------------------------------------------------------------ #
    def write_csharp(self, filename: str, code: str) -> WrittenFile:
        """Write a C# script to ``Assets/Scripts/`` (with .meta)."""
        if not filename.endswith(".cs"):
            filename += ".cs"
        rel = f"{self.SCRIPTS_DIR}/{filename}"
        return self._write(rel, code, "csharp", write_meta=True)

    def write_scene(self, scene_name: str, yaml: str) -> WrittenFile:
        if not scene_name.endswith(".unity"):
            scene_name += ".unity"
        rel = f"{self.SCENES_DIR}/{scene_name}"
        return self._write(rel, yaml, "scene", write_meta=True)

    def write_config(self, relative_path: str, content: str) -> WrittenFile:
        return self._write(relative_path, content, "config", write_meta=False)

    def write_raw(self, relative_path: str, content: str, kind: str = "other") -> WrittenFile:
        return self._write(relative_path, content, kind, write_meta=False)

    # ------------------------------------------------------------------ #
    # Manifest / settings
    # ------------------------------------------------------------------ #
    def write_packages_manifest(self, dependencies: Optional[Dict[str, str]] = None) -> WrittenFile:
        deps = dependencies or {
            "com.unity.ai.navigation": "1.1.5",
            "com.unity.cinemachine": "2.9.7",
            "com.unity.collab-proxy": "2.2.0",
            "com.unity.ide.rider": "3.0.27",
            "com.unity.ide.visualstudio": "2.0.22",
            "com.unity.inputsystem": "1.7.0",
            "com.unity.postprocessing": "3.4.0",
            "com.unity.textmeshpro": "3.0.6",
            "com.unity.ugui": "1.0.0",
            "com.unity.modules.ai": "1.0.0",
            "com.unity.modules.animation": "1.0.0",
            "com.unity.modules.audio": "1.0.0",
            "com.unity.modules.imageconversion": "1.0.0",
            "com.unity.modules.particlesystem": "1.0.0",
            "com.unity.modules.physics": "1.0.0",
            "com.unity.modules.physics2d": "1.0.0",
            "com.unity.modules.terrain": "1.0.0",
            "com.unity.modules.terrainphysics": "1.0.0",
            "com.unity.modules.ui": "1.0.0",
        }
        manifest = {
            "dependencies": deps,
        }
        body = json.dumps(manifest, indent=2)
        return self.write_config(f"{self.PACKAGES_DIR}/manifest.json", body)

    def write_project_version(self) -> WrittenFile:
        return self.write_config(
            f"{self.PROJECT_SETTINGS_DIR}/ProjectVersion.txt",
            f"m_EditorVersion: {self.settings.unity_version}\n",
        )

    def write_project_settings(self) -> WrittenFile:
        body = (
            "%YAML 1.1\n"
            "%TAG !u! tag:unity3d.com,2011:\n"
            "--- !u!129 &1\n"
            "PlayerSettings:\n"
            "  m_ObjectHideFlags: 0\n"
            "  serializedVersion: 26\n"
            f"  companyName: {self.settings.company_name}\n"
            f"  productName: {self.settings.product_name}\n"
            "  defaultCursor: {instanceID: 0}\n"
            "  m_Version: 1\n"
        )
        return self.write_config(f"{self.PROJECT_SETTINGS_DIR}/ProjectSettings.asset", body)

    def write_assembly_definition(self, name: str = "UnityAgent.Scripts") -> WrittenFile:
        asmdef = {
            "name": name,
            "rootNamespace": "UnityAgent",
            "references": [],
            "includePlatforms": [],
            "excludePlatforms": [],
            "allowUnsafeCode": False,
            "autoReferenced": True,
            "defineConstraints": [],
            "versionDefines": [],
        }
        body = json.dumps(asmdef, indent=2)
        if not name.endswith(".asmdef"):
            name += ".asmdef"
        return self._write(f"{self.SCRIPTS_DIR}/{name}", body, "config", write_meta=True)

    # ------------------------------------------------------------------ #
    # Reporting
    # ------------------------------------------------------------------ #
    def summary(self) -> str:
        lines = [
            f"Unity project: {self.project_root}",
            f"Total files written: {len(self.stats.files)}",
            f"Total size: {self.stats.total_bytes:,} bytes",
        ]
        kinds = {}
        for f in self.stats.files:
            kinds[f.kind] = kinds.get(f.kind, 0) + 1
        for k, v in sorted(kinds.items()):
            lines.append(f"  {k}: {v}")
        return "\n".join(lines)