acdream/src/AcDream.Bake/BakeRunner.cs

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using AcDream.Content;
using AcDream.Content.Pak;
using DatReaderWriter;
using DatReaderWriter.DBObjs;
using DatReaderWriter.Options;

namespace AcDream.Bake;

/// <summary>Options for one bake run (parsed from CLI args by Program, or built directly by tests).</summary>
public sealed record BakeOptions {
    public required string DatDir { get; init; }
    public required string OutPath { get; init; }
    public HashSet<uint>? IdFilter { get; init; }
    public HashSet<uint>? LandblockFilter { get; init; }
    public int Threads { get; init; } = System.Environment.ProcessorCount;
}

/// <summary>
/// The bake pipeline: enumerate ids -> sorted-batch parallel extraction ->
/// deterministic pak write. Public (rather than inline in Program) so the
/// dat-gated byte-reproducibility test can drive the REAL pipeline.
///
/// <para><b>Determinism + bounded memory (review finding 1):</b> the full id
/// list is built first and sorted by pak key; work proceeds in fixed-size
/// batches — extraction is parallel WITHIN a batch, then the batch's results
/// are sorted by key and written sequentially before the next batch starts.
/// Batches are contiguous key ranges, so the blob region ends up in global
/// key order regardless of thread scheduling, and memory is bounded by one
/// batch's decoded output instead of the whole game's. Side-staged
/// particle-preload meshes are deduped into a keyed map as each batch
/// drains (first instance wins — extraction output per id is
/// deterministic, so instance choice cannot affect bytes) and written after
/// all batches, sorted by key, skipping keys already written.</para>
/// </summary>
public static class BakeRunner {
    /// <summary>Ids extracted in parallel per batch before the sequential sorted write. Bounds peak memory (~512 decoded ObjectMeshData) while keeping the workers busy.</summary>
    private const int BatchSize = 512;

    public static int Run(BakeOptions options) {
        Console.WriteLine("acdream-bake");
        Console.WriteLine($"dat dir: {options.DatDir}");
        Console.WriteLine($"out:     {options.OutPath}");
        Console.WriteLine($"threads: {options.Threads}");
        Console.WriteLine();

        var sw = Stopwatch.StartNew();
        using var dats = new DatCollection(options.DatDir, DatAccessType.Read);
        // The unified AcDream.Content.DatCollectionAdapter (MP1b review
        // finding 7): same instance class the client and the equivalence
        // suite use — adapter drift between bake and runtime is impossible
        // by construction.
        using var datReaderWriter = new DatCollectionAdapter(dats);
        var extractorLogger = new ConsoleErrorLogger(nameof(MeshExtractor));

        // Thread-safe side-stage sink for particle-preload meshes MeshExtractor
        // emits mid-extraction (MP1a documented contract; extractor is shared
        // across the batch workers).
        var sideStaged = new ConcurrentQueue<ObjectMeshData>();
        var extractor = new MeshExtractor(datReaderWriter, extractorLogger, data => sideStaged.Enqueue(data));

        // ---- enumeration -----------------------------------------------------

        var gfxObjIds = dats.GetAllIdsOfType<GfxObj>().ToList();
        var setupIds = dats.GetAllIdsOfType<Setup>().ToList();
        var envCellIds = EnumerateEnvCellIds(dats, options.LandblockFilter);

        if (options.IdFilter is not null) {
            gfxObjIds = gfxObjIds.Where(options.IdFilter.Contains).ToList();
            setupIds = setupIds.Where(options.IdFilter.Contains).ToList();
            envCellIds = envCellIds.Where(options.IdFilter.Contains).ToList();
        }

        // Full work list sorted by pak key — the batching below preserves this
        // global order on disk (determinism precondition).
        var work = new List<(ulong Key, PakAssetType Type, uint FileId)>(gfxObjIds.Count + setupIds.Count + envCellIds.Count);
        work.AddRange(gfxObjIds.Select(id => (PakKey.Compose(PakAssetType.GfxObjMesh, id), PakAssetType.GfxObjMesh, id)));
        work.AddRange(setupIds.Select(id => (PakKey.Compose(PakAssetType.SetupMesh, id), PakAssetType.SetupMesh, id)));
        work.AddRange(envCellIds.Select(id => (PakKey.Compose(PakAssetType.EnvCellMesh, id), PakAssetType.EnvCellMesh, id)));
        work.Sort((a, b) => a.Key.CompareTo(b.Key));

        Console.WriteLine($"enumerated: {gfxObjIds.Count:N0} GfxObj, {setupIds.Count:N0} Setup, {envCellIds.Count:N0} EnvCell " +
            $"({work.Count:N0} total)");
        Console.WriteLine();

        // ---- sorted-batch extraction + streaming write -------------------------

        var header = new PakHeader {
            PortalIteration = (uint)dats.Portal.Iteration!.CurrentIteration,
            CellIteration = (uint)dats.Cell.Iteration!.CurrentIteration,
            HighResIteration = (uint)dats.HighRes.Iteration!.CurrentIteration,
            LanguageIteration = (uint)dats.Local.Iteration!.CurrentIteration,
            BakeToolVersion = 1,
        };

        var failures = new ConcurrentBag<(PakAssetType Type, uint FileId, string Reason)>();
        var writtenKeys = new HashSet<ulong>();
        var sideStagedByKey = new Dictionary<ulong, ObjectMeshData>();
        long completed = 0;
        int writtenCounts0 = 0, writtenCounts1 = 0, writtenCounts2 = 0; // GfxObj, Setup, EnvCell
        var lastProgressReport = Stopwatch.StartNew();

        using (var writer = new PakWriter(options.OutPath, header)) {
            for (int batchStart = 0; batchStart < work.Count; batchStart += BatchSize) {
                var batch = work.Skip(batchStart).Take(BatchSize).ToList();
                var batchResults = new ConcurrentBag<(ulong Key, PakAssetType Type, ObjectMeshData Data)>();

                Parallel.ForEach(
                    batch,
                    new ParallelOptions { MaxDegreeOfParallelism = options.Threads },
                    item => {
                        var (key, type, fileId) = item;
                        try {
                            // EnvCell entries store the cell's synthetic geometry mesh
                            // (bit 32 set — see MeshExtractor.PrepareMeshData's EnvCell
                            // branch); the cell's static objects are covered by their
                            // own GfxObj/Setup entries.
                            ulong extractorId = type == PakAssetType.EnvCellMesh ? fileId | 0x1_0000_0000UL : fileId;
                            // isSetup: matches the runtime's own request sites —
                            // WbMeshAdapter.IncrementRefCount/EnsureLoaded pass
                            // isSetup: false for every MeshRef id (incl. cell-geometry
                            // ids); only true Setup extraction passes true. (Review
                            // finding 10 — keep aligned with PakEquivalenceTests.)
                            bool isSetup = type == PakAssetType.SetupMesh;
                            var data = extractor.PrepareMeshData(extractorId, isSetup);
                            if (data is not null) {
                                batchResults.Add((key, type, data));
                            }
                            else {
                                failures.Add((type, fileId, "extractor returned null (no polygons or unresolvable id)"));
                            }
                        }
                        catch (Exception ex) {
                            // A malformed dat entry skips that id — never fatal to the
                            // bake, matching the runtime's own per-id behavior.
                            failures.Add((type, fileId, ex.Message));
                        }

                        Interlocked.Increment(ref completed);
                    });

                // Sequential sorted write: batch results in key order.
                foreach (var (key, type, data) in batchResults.OrderBy(r => r.Key)) {
                    writer.AddBlob(key, data);
                    writtenKeys.Add(key);
                    switch (type) {
                        case PakAssetType.GfxObjMesh: writtenCounts0++; break;
                        case PakAssetType.SetupMesh: writtenCounts1++; break;
                        case PakAssetType.EnvCellMesh: writtenCounts2++; break;
                    }
                }

                // Drain side-staged preloads per batch into the dedup map so the
                // queue never grows past one batch's emissions (memory bound);
                // written after all batches, sorted (see below).
                while (sideStaged.TryDequeue(out var staged)) {
                    uint fileId = (uint)(staged.ObjectId & 0xFFFFFFFFu);
                    ulong key = PakKey.Compose(PakAssetType.GfxObjMesh, fileId);
                    if (!sideStagedByKey.ContainsKey(key)) sideStagedByKey[key] = staged;
                }

                if (lastProgressReport.Elapsed.TotalSeconds >= 5 || batchStart + BatchSize >= work.Count) {
                    ReportProgress(completed, work.Count, failures.Count, sw.Elapsed);
                    lastProgressReport.Restart();
                }
            }

            // Side-staged preload meshes not already covered by a primary entry:
            // sorted by key for deterministic placement.
            int sideStagedWritten = 0, sideStagedDuped = 0;
            foreach (var (key, data) in sideStagedByKey.OrderBy(kv => kv.Key)) {
                if (writtenKeys.Contains(key)) { sideStagedDuped++; continue; }
                writer.AddBlob(key, data);
                writtenKeys.Add(key);
                sideStagedWritten++;
            }

            writer.Finish();
            sw.Stop();

            var outSize = new FileInfo(options.OutPath).Length;
            Console.WriteLine();
            Console.WriteLine("=== bake summary ===");
            Console.WriteLine($"  GfxObj baked:     {writtenCounts0:N0}");
            Console.WriteLine($"  Setup baked:      {writtenCounts1:N0}");
            Console.WriteLine($"  EnvCell baked:    {writtenCounts2:N0}");
            Console.WriteLine($"  side-staged baked (particle preload GfxObjs): {sideStagedWritten:N0} ({sideStagedDuped:N0} deduped)");
            Console.WriteLine($"  total blobs:      {writtenKeys.Count:N0}");
            Console.WriteLine($"  failures:         {failures.Count:N0}");
            Console.WriteLine($"  elapsed:          {sw.Elapsed.TotalSeconds:F1} s");
            Console.WriteLine($"  output size:      {outSize / 1024.0 / 1024.0:F1} MB");
            Console.WriteLine($"  output path:      {options.OutPath}");
        }

        if (!failures.IsEmpty) {
            Console.WriteLine();
            Console.WriteLine($"failures ({failures.Count}):");
            foreach (var (type, fileId, reason) in failures.OrderBy(f => f.FileId).Take(200)) {
                Console.WriteLine($"  {type,-12} 0x{fileId:X8}: {reason}");
            }
            if (failures.Count > 200) Console.WriteLine($"  ... and {failures.Count - 200} more");
        }

        return 0;
    }

    private static void ReportProgress(long done, int total, int failures, TimeSpan elapsed) {
        double rate = elapsed.TotalSeconds > 0 ? done / elapsed.TotalSeconds : 0;
        double etaSeconds = rate > 0 ? (total - done) / rate : 0;
        Console.WriteLine($"[{elapsed:hh\\:mm\\:ss}] baked {done:N0}/{total:N0}, failures={failures:N0}, " +
            $"elapsed={elapsed.TotalSeconds:F0}s, ETA={etaSeconds:F0}s");
    }

    /// <summary>
    /// Enumerates EnvCell ids by walking the cell dat's LandBlockInfo entries
    /// (0xFFFE low 16 bits) and, for each, the NumCells-derived cell id range —
    /// GetAllIdsOfType&lt;T&gt;() does not cover cell-dat range-based types
    /// (DatReaderWriter's documented limitation; see the low-16-bit bucketing
    /// idiom in src/AcDream.Cli/Program.cs's CountCellByLow16, and the
    /// firstCellId/NumCells hydration idiom in GameWindow.BuildPhysicsDatBundle
    /// / BuildInteriorEntitiesForStreaming).
    /// </summary>
    private static List<uint> EnumerateEnvCellIds(DatCollection dats, HashSet<uint>? landblockFilter) {
        var landblockInfoIds = new List<uint>();
        foreach (var file in dats.Cell.Tree) {
            if ((file.Id & 0xFFFFu) != 0xFFFEu) continue;
            uint landblockId = file.Id & 0xFFFF0000u;
            if (landblockFilter is not null && !landblockFilter.Contains(landblockId)) continue;
            landblockInfoIds.Add(file.Id);
        }

        var envCellIds = new List<uint>();
        foreach (var lbInfoId in landblockInfoIds) {
            if (!dats.Cell.TryGet<LandBlockInfo>(lbInfoId, out var lbInfo) || lbInfo is null) continue;
            if (lbInfo.NumCells == 0) continue;

            uint landblockId = lbInfoId & 0xFFFF0000u;
            uint firstCellId = landblockId | 0x0100u;
            for (uint offset = 0; offset < lbInfo.NumCells; offset++) {
                envCellIds.Add(firstCellId + offset);
            }
        }
        return envCellIds;
    }
}