UnityTowerDefense/Assets/_Project/Scripts/Net/SessionFlow.cs

// Assets/_Project/Scripts/Net/SessionFlow.cs
using Unity.Netcode;
using UnityEngine;
using UnityEngine.SceneManagement;

namespace TD.Net
{
    /// <summary>
    /// Bootstrap script — drop one on a GameObject in the MainMenu scene so it
    /// runs once when the app boots. Subscribes to NGO disconnect callbacks
    /// and routes the local peer back to the MainMenu scene whenever its
    /// connection ends (host left, server kicked it, transport error, etc.).
    /// </summary>
    /// <remarks>
    /// <para><b>Why this exists.</b> Per Option A (see Project_Roadmap.md §1.7),
    /// host disconnect tears down the session for everyone. Each client needs
    /// to recover gracefully and return to its own main menu. Without this,
    /// a client would be stuck in the Lobby or Match scene with no working
    /// network connection after the host quits.</para>
    ///
    /// <para><b>Lifecycle.</b> The NetworkManager itself is marked
    /// DontDestroyOnLoad by NGO once it spawns. This script is also marked
    /// DontDestroyOnLoad so it survives scene transitions and its subscription
    /// stays alive across MainMenu → Lobby → Match → back.</para>
    ///
    /// <para><b>Server side.</b> When the host (server) shuts down, the
    /// <c>OnServerStopped</c> callback fires on the host too. For the host
    /// that's fine — they've initiated the shutdown deliberately (e.g., from
    /// the "Return to Main Menu" button) and being sent to MainMenu is what
    /// they wanted.</para>
    /// </remarks>
    public class SessionFlow : MonoBehaviour
    {
        // ----- Singleton --------------------------------------------------

        public static SessionFlow Instance { get; private set; }

        private void Awake()
        {
            if (Instance != null && Instance != this)
            {
                Destroy(gameObject);
                return;
            }
            Instance = this;
            DontDestroyOnLoad(gameObject);
        }

        private void Start()
        {
            HookCallbacks();
        }

        private void OnDestroy()
        {
            UnhookCallbacks();
            if (Instance == this) Instance = null;
        }

        // ----- NGO callback wiring ---------------------------------------

        private bool callbacksHooked;

        private void HookCallbacks()
        {
            var nm = NetworkManager.Singleton;
            if (nm == null)
            {
                // NetworkManager may not be alive yet on the very first frame
                // depending on script execution order — retry next frame.
                Invoke(nameof(HookCallbacks), 0.1f);
                return;
            }
            if (callbacksHooked) return;

            nm.OnClientDisconnectCallback += HandleClientDisconnect;
            nm.OnServerStopped             += HandleServerStopped;
            callbacksHooked = true;
        }

        private void UnhookCallbacks()
        {
            if (!callbacksHooked) return;
            var nm = NetworkManager.Singleton;
            if (nm != null)
            {
                nm.OnClientDisconnectCallback -= HandleClientDisconnect;
                nm.OnServerStopped             -= HandleServerStopped;
            }
            callbacksHooked = false;
        }

        // ----- Disconnect handlers ---------------------------------------

        // Fires when ANY client disconnects on the server, and when the LOCAL
        // client gets disconnected on a client. We only care about the latter:
        // when the local peer loses its connection (host left, transport error,
        // explicit Shutdown call), route back to the main menu.
        private void HandleClientDisconnect(ulong clientId)
        {
            var nm = NetworkManager.Singleton;
            if (nm == null) return;

            // On the server, this fires when other clients disconnect — ignore
            // those, the server keeps running. We only act when the LOCAL
            // client's connection ends.
            if (clientId != nm.LocalClientId) return;

            ReturnToMainMenu();
        }

        // Fires on the server when the server shuts down (including the host).
        // The host's own client also gets HandleClientDisconnect — but
        // OnServerStopped is the canonical "server is gone" signal and is
        // safer to act on for the host path.
        private void HandleServerStopped(bool isHost)
        {
            ReturnToMainMenu();
        }

        // ----- Scene transition ------------------------------------------

        // Returns the local peer to the MainMenu scene if they're not already
        // there. Uses Unity's local SceneManager (NOT NGO's networked scene
        // manager) because by the time this fires the network connection is
        // already gone.
        private static void ReturnToMainMenu()
        {
            if (SceneManager.GetActiveScene().name == SceneNames.MainMenu)
                return;

            Debug.Log("[SessionFlow] Connection ended — returning to main menu.");
            SceneManager.LoadScene(SceneNames.MainMenu);
        }
    }
}