#!/usr/bin/env python3
"""Stage source and restore files for a Linux Agent image build."""

from __future__ import annotations

import argparse
import os
import shutil
import subprocess
import xml.etree.ElementTree as ET
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
FIXED_MTIME = 946684800  # 2000-01-01T00:00:00Z
BACKEND = ROOT / "backend"
STATIC_MSQUIC = ROOT / "packaging/native/msquic"
IGNORED_PARTS = {".git", ".release", "artifacts", "bin", "obj", "TestResults", "tests"}
GLOBAL_BACKEND_FILES = {
    "Directory.Build.props",
    "Directory.Build.targets",
    "Directory.Packages.props",
    "global.json",
    "NuGet.Config",
    "NuGet.config",
}
# The container stages backend as a standalone tree. Copy the repository-wide
# SDK and package versions into it so restore/publish use the same configuration.
ROOT_BUILD_FILES = {"Directory.Packages.props", "global.json"}
CONTAINERFILES = {
    "agent": ROOT / "packaging/container/Agent.NativeBuild.Containerfile",
}


def restore_inputs() -> set[Path]:
    """Only the two executables, their project references and ancestor build inputs."""
    pending = [BACKEND / "agent/src/MoShiro.Agent/MoShiro.Agent.csproj",
               BACKEND / "fleet/tools/fleet-enroll/fleet-enroll.csproj"]
    projects: set[Path] = set()
    inputs: set[Path] = set()
    while pending:
        project = pending.pop().resolve()
        if project in projects:
            continue
        projects.add(project)
        inputs.add(project)
        for reference in ET.parse(project).iter("ProjectReference"):
            pending.append((project.parent / reference.attrib["Include"].replace("\\", "/")).resolve())
        directory = project.parent
        while directory.is_relative_to(BACKEND):
            inputs.update(p for p in directory.iterdir() if p.is_file() and (
                p.name in GLOBAL_BACKEND_FILES or p.name == "packages.lock.json" or p.suffix in {".props", ".targets"}))
            directory = directory.parent
    return inputs


def copy_for_restore(path: Path, relative: Path, target: Path) -> None:
    """The restore layer must not change when only the build version does: CI rewrites <Version> in the workspace
    Directory.Build.props (ci/version.sh), which used to invalidate the cached dependency restore on every commit.
    Restore only needs the committed props (package versions); the source copy still carries the real version."""
    if path.name == "Directory.Build.props":
        committed = subprocess.run(
            ["git", "show", f"HEAD:backend/{relative.as_posix()}"], cwd=ROOT, capture_output=True, check=False
        )
        if committed.returncode == 0:
            target.write_bytes(committed.stdout)
            return
    shutil.copy2(path, target)


def stage(component: str, runtime_identifier: str, destination: Path) -> None:
    if destination.exists():
        shutil.rmtree(destination)
    restore = destination / "restore/backend"
    source = destination / "source/backend"
    restore.mkdir(parents=True)
    source.mkdir(parents=True)
    shutil.copy2(CONTAINERFILES[component], destination / "Containerfile")
    inputs = restore_inputs()

    for path in BACKEND.rglob("*"):
        if not path.is_file() or IGNORED_PARTS.intersection(path.parts):
            continue
        relative = path.relative_to(BACKEND)
        if path in inputs:
            target = restore / relative
            target.parent.mkdir(parents=True, exist_ok=True)
            copy_for_restore(path, relative, target)
        target = source / relative
        target.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(path, target)

    for name in ROOT_BUILD_FILES:
        for directory in (restore, source):
            shutil.copy2(ROOT / name, directory / name)

    if runtime_identifier.startswith("linux-musl-"):
        static_library = STATIC_MSQUIC / runtime_identifier / "libmsquic.a"
        if not static_library.is_file():
            raise FileNotFoundError(f"static MsQuic archive is missing for {runtime_identifier}: {static_library}")
        staged_library = destination / "source/packaging/native/msquic" / runtime_identifier / "libmsquic.a"
        staged_library.parent.mkdir(parents=True, exist_ok=True)
        shutil.copy2(static_library, staged_library)
    else:
        (destination / "source/packaging/native/msquic").mkdir(parents=True)
        (destination / "source/packaging/native/msquic/.keep").touch()

    # A plain traversal project forwards the RID to each executable. SDK solution
    # builds reject a solution-wide RuntimeIdentifier (NETSDK1134).
    traversal = ('<Project DefaultTargets="Build">\n'
                 '  <ItemGroup>\n'
                 '    <NodeProject Include="agent/src/MoShiro.Agent/MoShiro.Agent.csproj" />\n'
                 '    <NodeProject Include="fleet/tools/fleet-enroll/fleet-enroll.csproj" />\n'
                 '  </ItemGroup>\n'
                 '  <Target Name="Restore">\n'
                 '    <MSBuild Projects="@(NodeProject)" Targets="Restore" BuildInParallel="false"\n'
                 '      Properties="RuntimeIdentifier=$(RuntimeIdentifier);SelfContained=$(SelfContained);PublishAot=$(PublishAot);UseAotAnywhere=$(UseAotAnywhere);StaticMsQuic=$(StaticMsQuic);StaticExecutable=$(StaticExecutable);RuntimeFrameworkVersion=$(RuntimeFrameworkVersion);InvariantGlobalization=$(InvariantGlobalization)" />\n'
                 '  </Target>\n'
                 '  <Target Name="Build">\n'
                 '    <MSBuild Projects="@(NodeProject)" Targets="Build" BuildInParallel="false"\n'
                 '      Properties="RuntimeIdentifier=$(RuntimeIdentifier);SelfContained=$(SelfContained);PublishAot=$(PublishAot);UseAotAnywhere=$(UseAotAnywhere);StaticMsQuic=$(StaticMsQuic);StaticExecutable=$(StaticExecutable);RuntimeFrameworkVersion=$(RuntimeFrameworkVersion);InvariantGlobalization=$(InvariantGlobalization);Version=$(Version);UseAppHost=$(UseAppHost);StripSymbols=$(StripSymbols);DebugType=$(DebugType);DebugSymbols=$(DebugSymbols)" />\n'
                 '  </Target>\n'
                 '</Project>\n')
    for directory in (restore, source):
        (directory / "Agent.Native.proj").write_text(traversal, encoding="utf-8")

    normalize_mtimes(destination / "restore")


def normalize_mtimes(directory: Path) -> None:
    """The Buildah layer cache keys COPY layers on file metadata as well as content. Files copied from a fresh checkout carry a new
    mtime every run, so the restore layer (and the ~70 s of NuGet downloads under it) never hit the cache. Pin the mtimes."""
    for path in [directory, *directory.rglob("*")]:
        os.utime(path, (FIXED_MTIME, FIXED_MTIME))


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("component", choices=tuple(CONTAINERFILES))
    parser.add_argument("runtime_identifier", choices=("linux-x64", "linux-arm64", "linux-musl-arm64"))
    parser.add_argument("destination", type=Path)
    args = parser.parse_args()
    stage(args.component, args.runtime_identifier, args.destination)


if __name__ == "__main__":
    main()