Project-M/Assets/_Project/Scripts/Client/UI/WorldLauncher.cs

using System;
using System.Collections;
using ProjectM.Simulation;
using Unity.Entities;
using Unity.NetCode;
using UnityEngine;
using UnityEngine.SceneManagement;

namespace ProjectM.Client
{
    public enum SessionMode { Single, Host, Join }

    /// <summary>
    /// The SINGLE funnel for netcode world lifecycle driven by the menu, so a connection/handshake fix lives in
    /// one place. <see cref="StartSession"/> creates the right worlds (Single/Host = server+client, Join =
    /// client-only) via the public <c>ClientServerBootstrap.Create*World</c> helpers (which register the
    /// ServerWorld/ClientWorld statics), seeds the existing <see cref="ConnectionConfig"/> request component,
    /// optionally stages a save (Continue), then loads Game.unity worlds-first (the subscene-streaming
    /// invariant). <see cref="TeardownToMenu"/> autosaves (host), disposes all worlds, and returns to MainMenu.
    /// Every dispose/scene step runs at a frame boundary on <see cref="SessionRunner"/> — never in an ECS system.
    /// </summary>
    public static class WorldLauncher
    {
        public const ushort Port = 7979;
        const string GameScene = "Game";
        const string MenuScene = "MainMenu";
        const string Loopback = "127.0.0.1";

        public static bool Busy { get; private set; }

        public static void StartSession(SessionMode mode, string joinIp, bool loadSave)
        {
            if (Busy) return;
            Busy = true;
            SessionRunner.Run(StartRoutine(mode, joinIp, loadSave));
        }

        public static void TeardownToMenu()
        {
            if (Busy) return;
            Busy = true;
            SessionRunner.Run(TeardownRoutine());
        }

        static IEnumerator StartRoutine(SessionMode mode, string joinIp, bool loadSave)
        {
            // Dispose the idle menu world so a netcode world can own DefaultGameObjectInjectionWorld (the
            // subscene then streams into the netcode worlds, exactly as in the always-on bootstrap flow).
            World.DisposeAllWorlds();
            yield return null;

            World server = null;
            World client = ClientServerBootstrap.CreateClientWorld("ClientWorld");

            if (mode == SessionMode.Join)
            {
                Seed(client, ConnectionMode.Join, string.IsNullOrWhiteSpace(joinIp) ? Loopback : joinIp.Trim(), Port);
            }
            else
            {
                server = ClientServerBootstrap.CreateServerWorld("ServerWorld");
                string bind = mode == SessionMode.Single ? Loopback : "0.0.0.0"; // Single binds loopback (no firewall)
                Seed(server, ConnectionMode.Host, bind, Port);
                Seed(client, ConnectionMode.Join, Loopback, Port);
                if (loadSave) StagePendingSave(server);
            }

            World.DefaultGameObjectInjectionWorld = server ?? client;

            // Worlds exist -> loading Game.unity streams its SubScene into them.
            SceneManager.LoadScene(GameScene, LoadSceneMode.Single);
            Busy = false;
        }

        static IEnumerator TeardownRoutine()
        {
            var server = ClientServerBootstrap.ServerWorld;
            if (server is { IsCreated: true })
                TrySaveFromServer(server);

            yield return null;
            World.DisposeAllWorlds();
            yield return null;

            SceneManager.LoadScene(MenuScene, LoadSceneMode.Single);
            Busy = false;
        }

        static void Seed(World world, ConnectionMode mode, string address, ushort port)
        {
            if (world is not { IsCreated: true }) return;
            var em = world.EntityManager;
            using var q = em.CreateEntityQuery(ComponentType.ReadWrite<ConnectionConfig>());
            Entity e = q.IsEmptyIgnoreFilter ? em.CreateEntity(typeof(ConnectionConfig)) : q.GetSingletonEntity();
            em.SetComponentData(e, new ConnectionConfig
            {
                Mode = mode,
                Address = address,
                Port = port,
                Requested = true,
            });
        }

        static void StagePendingSave(World server)
        {
            var data = SaveService.Load();
            if (data == null) return;
            var em = server.EntityManager;
            var e = em.CreateEntity();
            em.AddComponentData(e, new PendingSave { GoalCharge = data.GoalCharge, GoalTarget = data.GoalTarget, HasData = 1 });
            var buf = em.AddBuffer<PendingSaveLedgerRow>(e);
            if (data.Ledger != null)
                foreach (var row in data.Ledger)
                    buf.Add(new PendingSaveLedgerRow { ItemId = (ushort)row.ItemId, Count = row.Count });

            // M7: stage player-built structures on a SEPARATE carrier (BaseRestoreSystem owns its lifecycle; the
            // PendingSave entity above is consumed + destroyed by CycleDirectorSpawnSystem at director spawn).
            if (data.Structures != null && data.Structures.Length > 0)
            {
                var se = em.CreateEntity();
                var sbuf = em.AddBuffer<PendingStructure>(se);
                foreach (var s in data.Structures)
                    sbuf.Add(new PendingStructure
                    {
                        Type = s.Type, CellX = s.CellX, CellZ = s.CellZ, Direction = s.Direction,
                        RemainingTicks = s.RemainingTicks, ConveyorResId = s.ConveyorResId, ConveyorCount = s.ConveyorCount,
                    });
                var iobuf = em.AddBuffer<PendingStructureIo>(se);
                if (data.StructureIo != null)
                    foreach (var io in data.StructureIo)
                        iobuf.Add(new PendingStructureIo { StructureIndex = io.StructureIndex, Slot = io.Slot, ResourceId = io.ResourceId, Count = io.Count });
            }
        }

        static void TrySaveFromServer(World server)
        {
            try
            {
                var em = server.EntityManager;
                em.CompleteAllTrackedJobs();
                using var q = em.CreateEntityQuery(ComponentType.ReadOnly<ResourceLedger>());
                if (q.IsEmptyIgnoreFilter) return;
                var dir = q.GetSingletonEntity();
                var goal = em.HasComponent<GoalProgress>(dir) ? em.GetComponentData<GoalProgress>(dir) : default;
                var buffer = em.GetBuffer<StorageEntry>(dir, true);
                var rows = new LedgerRow[buffer.Length];
                for (int i = 0; i < buffer.Length; i++)
                    rows[i] = new LedgerRow { ItemId = buffer[i].ItemId, Count = buffer[i].Count };

                // M7: also persist player-built structures (same shared scan as the autosave path).
                uint nowTick = 0;
                using (var tq = em.CreateEntityQuery(ComponentType.ReadOnly<NetworkTime>()))
                    if (!tq.IsEmptyIgnoreFilter)
                    {
                        var st = tq.GetSingleton<NetworkTime>().ServerTick;
                        if (st.IsValid) nowTick = st.TickIndexForValidTick;
                    }
                SaveStructureScan.Collect(em, nowTick, out var structures, out var structureIo);

                SaveService.Save(new SaveData
                {
                    GoalCharge = goal.Charge,
                    GoalTarget = goal.Target,
                    Ledger = rows,
                    Structures = structures,
                    StructureIo = structureIo,
                    SavedAtMs = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds(),
                });
            }
            catch (Exception e)
            {
                Debug.LogWarning($"[WorldLauncher] Quit-to-menu save skipped: {e.Message}");
            }
        }
    }
}