GameFlowManager.cs
using Sandbox;
using Sandbox.Network;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;

public enum OnlineFlowState
{
	Idle,
	Searching,
	Connecting,
	WaitingForOpponent,
	Lobby,
	InMatch,
	Error
}

/// <summary>
/// Owns session-level flow that does not belong inside RoundManager:
/// finding/hosting a 1v1 lobby, presenting a short pre-match lobby,
/// starting sandbox networking and leaving a network session cleanly.
///
/// Put this component on the same networked scene GameObject as RoundManager.
/// </summary>
public sealed partial class GameFlowManager : Component, Component.INetworkListener
{
	private const string LobbyKindKey = "omr";
	private const string LobbyStateKey = "st";
	private const string DuelLobbyValue = "duel";
	private const string SandboxLobbyValue = "sand";
	private const string BotLobbyValue = "bot";
	private const string SearchingLobbyValue = "q";
	private const string FullLobbyValue = "f";

	[Property]
	public RoundManager RoundManager { get; set; }

	[Property, Group( "Online" )]
	public float ConnectFallbackDelay { get; set; } = 5.0f;

	[Property, Sync( SyncFlags.FromHost )]
	public OnlineFlowState OnlineState { get; set; } =
		OnlineFlowState.Idle;

	[Property, Sync( SyncFlags.FromHost )]
	public float MatchStartsAt { get; set; }

	[Property, Sync( SyncFlags.FromHost )]
	public string ErrorMessage { get; set; } = "";

	private enum HostedLobbyKind
	{
		None,
		Duel,
		Sandbox,
		BotMatch
	}

	private HostedLobbyKind _hostedLobbyKind =
		HostedLobbyKind.None;

	private bool _lobbyDataPublished;
	private bool _searchInProgress;
	private int _searchRequestId;
	private float _connectFallbackAt;

	//
	// A host's local PlayerPrefab can survive Networking.Disconnect()
	// long enough to overlap the next PlayerPrefab spawned by
	// NetworkHelper. Keep an explicit teardown flag so the offline
	// menu is guaranteed to contain zero PlayerState objects before
	// another session is allowed to start.
	//
	private bool _cleanupPlayersWhenOffline;

	protected override void OnStart()
	{
		RoundManager ??=
			Components.Get<RoundManager>();
		_hadConnection = Networking.IsActive && !Networking.IsConnecting;

		// Offline is considered host by the engine. This gives the local
		// main-menu scene a clean starting state before any lobby exists.
		if (
			Networking.IsHost &&
			!Networking.IsActive
		)
		{
			ResetOnlineState();
		}
	}

	protected override void OnUpdate()
	{
		RoundManager ??=
			Components.Get<RoundManager>();

		if ( UpdateSessionLifecycle() )
			return;

		//
		// Networking.Disconnect() is not guaranteed to make every
		// session-owned object disappear in the same frame. Once the
		// network is fully inactive, purge any PlayerState roots left
		// behind by the previous lobby.
		//
		if (
			_cleanupPlayersWhenOffline &&
			!Networking.IsActive
		)
		{
			CleanupPlayerObjects();
			FinalizeOfflineMenuState();

			_cleanupPlayersWhenOffline =
				false;
			ReloadOfflineMenu();
			return;
		}

		if ( RoundManager is null )
			return;

		// A timed-out join returns to a clean menu rather than silently hosting
		// another isolated lobby. IsActive also includes connecting sessions.
		if (
			_connectionAttempt &&
			_connectFallbackAt > 0f &&
			RealTime.Now >= _connectFallbackAt
		)
		{
			FailConnection();
			return;
		}

		if ( !Networking.IsHost )
			return;

		PublishLobbyDataWhenReady();
		UpdateRendezvous();

		if (
			RoundManager.Mode == SessionMode.Menu
		)
		{
			UpdateOnlineLobby();
		}
		else if (
			RoundManager.Mode == SessionMode.Match &&
			OnlineState != OnlineFlowState.InMatch
		)
		{
			OnlineState =
				OnlineFlowState.InMatch;
		}
	}

	// ============================================================
	// ONLINE 1V1
	// ============================================================

	public void FindMatch()
	{
		if (
			RoundManager is null ||
			RoundManager.Mode != SessionMode.Menu ||
			_searchInProgress || _rendezvousInProgress || _leaving || Networking.IsConnecting
		)
			return;

		if ( !PrepareForNewSession() )
			return;
		if ( Networking.IsActive && !Networking.IsHost )
			return;

		OMRSceneFlow.SessionNotice = "";
		RoundManager.SetMatchBotEnabled( false );
		// The editor's second instance uses loopback, not Steam discovery.
		// Enable explicitly with omr_local_testing 1 in that instance.
		if ( LocalTestingEnabled && !Application.IsEditor )
		{
			JoinLocalTestHost();
			return;
		}

		_ = FindMatchAsync();
	}

	private async Task FindMatchAsync()
	{
		int requestId =
			++_searchRequestId;

		_searchInProgress =
			true;

		ErrorMessage =
			"";

		OnlineState =
			OnlineFlowState.Searching;

		MatchStartsAt =
			0f;

		try
		{
			// This fallback preserves editor workflows where a NetworkHelper
			// has already created a lobby. For real discovery, set that
			// component's Start Server property to false (see setup notes).
			if ( Networking.IsActive )
			{
				if ( Networking.IsHost )
				{
					UseCurrentSessionAsDuelLobby();
				}

				return;
			}

			var lobbies = await QueryDuelLobbies();

			if (
				requestId !=
					_searchRequestId
			)
				return;

			LobbyInformation candidate =
				lobbies
					.Where( IsJoinableDuelLobby )
					.OrderBy( x => x.OwnerId )
					.FirstOrDefault();

			if ( candidate.OwnerId != 0 )
			{
				Log.Info(
					$"MATCHMAKING | Joining lobby hosted by {candidate.OwnerId} | Members {candidate.Members}/{candidate.MaxMembers}"
				);

				BeginConnection( candidate.LobbyId.ToString() );

				return;
			}

			CreateDuelLobby();
		}
		catch ( Exception exception )
		{
			if (
				requestId !=
					_searchRequestId
			)
				return;

			ErrorMessage =
				exception.Message;

			OnlineState =
				OnlineFlowState.Error;

			Log.Warning(
				$"MATCHMAKING | Search failed | {exception.Message}"
			);
		}
		finally
		{
			if (
				requestId ==
					_searchRequestId
			)
			{
				_searchInProgress =
					false;
			}
		}
	}

	private bool IsJoinableDuelLobby(
		LobbyInformation lobby
	)
	{
		if (
			lobby.OwnerId == 0 ||
			lobby.LobbyId == 0 ||
			lobby.OwnerId == LocalOwnerId ||
			lobby.IsFull ||
			lobby.MaxMembers > 2
		)
			return false;

		return
			lobby.Get(
				LobbyKindKey,
				""
			) == DuelLobbyValue &&
			lobby.Get(
				LobbyStateKey,
				""
			) == SearchingLobbyValue;
	}

	private void CreateDuelLobby()
	{
		_searchRequestId++;
		_searchInProgress =
			false;

		LobbyConfig config =
			new()
			{
				Name = "One More Round 1v1",
				MaxPlayers = 2,
				Privacy = LobbyPrivacy.Public,
				Hidden = false,
				DestroyWhenHostLeaves = true
			};

		Networking.CreateLobby(
			config
		);

		_hostedLobbyKind =
			HostedLobbyKind.Duel;

		_lobbyDataPublished =
			false;

		_connectFallbackAt =
			0f;

		OnlineState =
			OnlineFlowState.WaitingForOpponent;
		_nextRendezvousAt = RealTime.Now + 4f;

		MatchStartsAt =
			0f;

		Log.Info(
			"MATCHMAKING | No suitable lobby found | Hosting 1v1 lobby"
		);
	}

	private void UseCurrentSessionAsDuelLobby()
	{
		// Networking.MaxPlayers is read-only once a session exists.
		// Sessions created by this manager use LobbyConfig.MaxPlayers = 2.
		// When reusing an editor/NetworkHelper session, we simply reuse
		// its existing capacity and keep the OMR game flow limited to
		// the two ready PlayerState slots handled by RoundManager.
		_hostedLobbyKind =
			HostedLobbyKind.Duel;

		_lobbyDataPublished =
			false;

		OnlineState =
			OnlineFlowState.WaitingForOpponent;

		MatchStartsAt =
			0f;

		Log.Info(
			"MATCHMAKING | Using existing host session as 1v1 lobby"
		);
	}

	private void UpdateOnlineLobby()
	{
		if (
			OnlineState !=
				OnlineFlowState.WaitingForOpponent &&
			OnlineState !=
				OnlineFlowState.Lobby
		)
			return;

		int connectedPlayers =
			GetReadyPlayerCount();

		if (
			OnlineState ==
				OnlineFlowState.WaitingForOpponent
		)
		{
			if ( connectedPlayers < 2 )
				return;

			OnlineState =
				OnlineFlowState.Lobby;

			MatchStartsAt =
				0f;

			RoundManager.BeginPreMatchLobby();

			PublishDuelState(
				FullLobbyValue
			);

			Log.Info(
				"MATCHMAKING | OPPONENT FOUND | PRE-MATCH LOBBY"
			);

			return;
		}

		if ( connectedPlayers < 2 )
		{
			RoundManager.EndPreMatchLobby();

			OnlineState =
				OnlineFlowState.WaitingForOpponent;

			MatchStartsAt =
				0f;

			PublishDuelState(
				SearchingLobbyValue
			);

			return;
		}

		if (
			RoundManager.ArePreMatchPlayersReady()
		)
		{
			OnlineState =
				OnlineFlowState.InMatch;

			MatchStartsAt =
				0f;

			RoundManager.StartMatchMode();
		}
	}

	public void CancelMatchmaking()
	{
		DisconnectToOfflineMenu();
	}

	// ============================================================
	// PRACTICE SESSION
	// ============================================================

	public void StartBotMatchSession()
	{
		if ( RoundManager is null || !Networking.IsHost || _leaving || Networking.IsConnecting )
			return;
		if ( !PrepareForNewSession() )
			return;

		CancelSearchWork();
		OMRSceneFlow.SessionNotice = "";
		ResetOnlineState();

		if ( !Networking.IsActive )
		{
			LobbyConfig config = new()
			{
				Name = "One More Round Bot Match",
				MaxPlayers = 1,
				Privacy = LobbyPrivacy.Public,
				Hidden = true,
				DestroyWhenHostLeaves = true
			};
			Networking.CreateLobby( config );
		}

		_hostedLobbyKind = HostedLobbyKind.BotMatch;
		_lobbyDataPublished = false;
		RoundManager.SetMatchBotEnabled( true );
		RoundManager.StartMatchMode();
	}

	public void StartSandboxSession()
	{
		if (
			RoundManager is null ||
			!Networking.IsHost || _leaving || Networking.IsConnecting
		)
			return;

		if ( !PrepareForNewSession() )
			return;
		CancelSearchWork();
		OMRSceneFlow.SessionNotice = "";

		ResetOnlineState();

		if ( !Networking.IsActive )
		{
			LobbyConfig config =
				new()
				{
					Name = "One More Round Sandbox",
					MaxPlayers = 2,
					Privacy = LobbyPrivacy.Public,
					Hidden = true,
					DestroyWhenHostLeaves = true
				};

			Networking.CreateLobby(
				config
			);
		}
		// If a session already exists (for example editor testing through
		// NetworkHelper), its MaxPlayers value cannot be changed at runtime.
		// Newly-created Sandbox lobbies above are explicitly capped at 2.
		_hostedLobbyKind =
			HostedLobbyKind.Sandbox;

		_lobbyDataPublished =
			false;

		RoundManager.StartSandboxMode();
	}

	/// <summary>
	/// Leaves active Sandbox gameplay and returns to its setup UI
	/// without disconnecting from the current lobby.
	///
	/// This keeps the existing PlayerState objects alive so loadout
	/// and arena changes can be applied to the same host/client
	/// session instead of tearing the lobby down and immediately
	/// creating another one.
	/// </summary>
	public void ReturnToSandboxSetup()
	{
		if ( !Networking.IsHost || _leaving )
			return;
		CancelSearchWork();
		_searchRequestId++;
		_searchInProgress =
			false;

		_connectFallbackAt =
			0f;

		ResetOnlineState();

		if (
			RoundManager is null ||
			!Networking.IsHost
		)
			return;

		RoundManager.ReturnToMenu();

		if ( Networking.IsActive )
		{
			_hostedLobbyKind =
				HostedLobbyKind.Sandbox;

			_lobbyDataPublished =
				false;
		}

		Log.Info(
			"SANDBOX | Returned to setup without disconnecting"
		);
	}

	// ============================================================
	// LEAVING SESSIONS
	// ============================================================

	public void ReturnToMainMenu()
	{
		DisconnectToOfflineMenu();
	}

	/// <summary>
	/// Fully tears down the current network session and guarantees
	/// that no PlayerState/viewmodel from it survives into the next
	/// lobby.
	///
	/// The host explicitly destroys the current player network roots
	/// before disconnecting. We also perform a second local cleanup
	/// after Networking.IsActive becomes false, which covers the host
	/// player's special case and any late local teardown.
	/// </summary>
	private void DisconnectToOfflineMenu()
	{
		if ( _leaving )
			return;

		_leaving = true;

		CancelSearchWork();

		_hadConnection = false;
		_connectionAttempt = false;
		_pendingConnectionTarget = null;

		_searchRequestId++;

		_searchInProgress =
			false;

		_connectFallbackAt =
			0f;

		bool isActiveHost =
			Networking.IsHost &&
			Networking.IsActive;

		//
		// If this is an intentional host shutdown, tell remote clients FIRST.
		// The RPC is reliable + SendImmediate, and the host remains online for
		// a short grace period. Clients therefore get a chance to execute their
		// own normal disconnect before S&box destroys the lobby.
		//
		if ( isActiveHost )
		{
			NotifyHostClosing();

			_hostDisconnectPending =
				true;

			_hostDisconnectAt =
				RealTime.Now + 0.35f;
		}

		if (
			RoundManager is not null &&
			Networking.IsHost
		)
		{
			RoundManager.ReturnToMenu();

			//
			// The host owns session teardown, so destroy every current
			// player avatar while authority is still valid.
			//
			CleanupPlayerObjects();
		}

		_cleanupPlayersWhenOffline =
			true;

		//
		// Clients (and offline callers) can disconnect immediately.
		// The active host disconnects later in UpdateSessionLifecycle(),
		// after remote clients leave or the 350 ms safety deadline expires.
		//
		if (
			Networking.IsActive &&
			!isActiveHost
		)
		{
			Networking.Disconnect();
		}

		// Finish cleanup and reload on a later update, after disconnect settles.

		_hostedLobbyKind =
			HostedLobbyKind.None;

		_lobbyDataPublished =
			false;

		ResetOnlineState();
	}

	/// <summary>
	/// Before creating/joining a new lobby, make sure a previous
	/// disconnect has actually finished and that the scene contains
	/// no stale PlayerPrefab instances.
	/// </summary>
	private bool PrepareForNewSession()
	{
		if ( _leaving || Networking.IsConnecting || _connectionAttempt || _pendingConnectionTarget is not null )
			return false;
		if ( !_cleanupPlayersWhenOffline )
		{
			if ( !Networking.IsActive )
			{
				CleanupPlayerObjects();
			}

			return true;
		}

		if ( Networking.IsActive )
		{
			Log.Warning(
				"SESSION | Previous disconnect is still finishing"
			);

			return false;
		}

		CleanupPlayerObjects();

		_cleanupPlayersWhenOffline =
			false;

		return true;
	}

	/// <summary>
	/// Removes all PlayerState network roots currently present in the
	/// local scene. This is only used while leaving a session or while
	/// offline immediately before starting another one.
	/// </summary>
	private void CleanupPlayerObjects()
	{
		if ( Scene is null )
			return;

		PlayerState[] players =
			Scene
				.GetAllComponents<PlayerState>()
				.ToArray();

		if ( players.Length == 0 )
			return;

		HashSet<GameObject> roots =
			new();

		foreach (
			PlayerState player in
				players
		)
		{
			if (
				player is null ||
				player.GameObject is null ||
				player.GameObject.IsDestroyed
			)
				continue;

			GameObject root =
				FindPlayerNetworkRoot(
					player.GameObject
				);

			if (
				root is null ||
				root.IsDestroyed
			)
				continue;

			roots.Add(
				root
			);
		}

		foreach (
			GameObject root in
				roots
		)
		{
			Log.Info(
				$"SESSION CLEANUP | Destroying stale player object '{root.Name}'"
			);

			root.Destroy();
		}
	}

	private void FinalizeOfflineMenuState()
	{
		if (
			Networking.IsActive ||
			RoundManager is null
		)
			return;

		//
		// After a client disconnects there is no host left to update
		// the synchronized round state. Locally returning these two
		// values to Menu/Waiting makes the offline UI visible again.
		//
		RoundManager.Mode =
			SessionMode.Menu;

		RoundManager.Phase =
			RoundPhase.Waiting;
	}

	private static GameObject FindPlayerNetworkRoot(
		GameObject playerObject
	)
	{
		if (
			playerObject is null ||
			playerObject.IsDestroyed
		)
			return playerObject;

		GameObject current =
			playerObject;

		while ( current is not null )
		{
			if ( current.IsNetworkRoot )
				return current;

			//
			// Never treat the scene/root container as a disposable player object.
			//
			// An offline or accidentally scene-authored PlayerPrefab is not a
			// network root. The old fallback climbed all the way to the top-level
			// scene object and returned it, so CleanupPlayerObjects() could destroy
			// the gameplay scene itself while preparing a Sandbox session.
			//
			if ( current.Parent is null )
				break;

			current =
				current.Parent;
		}

		//
		// PlayerState lives on the PlayerPrefab root in OMR. If no network root
		// exists (offline/editor/stale object), destroy only that player object.
		//
		return playerObject;
	}

	// ============================================================
	// LOBBY DATA
	// ============================================================

	private void PublishLobbyDataWhenReady()
	{
		if (
			_lobbyDataPublished ||
			!Networking.IsActive ||
			!Networking.IsHost
		)
			return;

		// MaxPlayers is configured when CreateLobby is called. The current
		// Networking.MaxPlayers property is read-only and is only queried.
		switch ( _hostedLobbyKind )
		{
			case HostedLobbyKind.Duel:
				Networking.SetData( "hdn", "0" );
				Networking.SetData(
					LobbyKindKey,
					DuelLobbyValue
				);

				Networking.SetData(
					LobbyStateKey,
					OnlineState == OnlineFlowState.Lobby
						? FullLobbyValue
						: SearchingLobbyValue
				);
				break;

			case HostedLobbyKind.Sandbox:
				Networking.SetData( "hdn", "1" );
				Networking.SetData( LobbyKindKey, SandboxLobbyValue );
				Networking.SetData( LobbyStateKey, "open" );
				break;

			case HostedLobbyKind.BotMatch:
				Networking.SetData( "hdn", "1" );
				Networking.SetData( LobbyKindKey, BotLobbyValue );
				Networking.SetData( LobbyStateKey, "bot" );
				break;

			default:
				return;
		}

		_lobbyDataPublished =
			true;
	}

	private void PublishDuelState(
		string state
	)
	{
		if (
			!Networking.IsHost ||
			!Networking.IsActive
		)
			return;

		Networking.SetData(
			LobbyStateKey,
			state
		);
	}

	// ============================================================
	// HELPERS
	// ============================================================

	private int GetReadyPlayerCount()
	{
		if ( Scene is null )
			return 0;

		return Scene
			.GetAllComponents<PlayerState>()
			.Count( x => x.IsReady );
	}

	private void ResetOnlineState()
	{
		OnlineState =
			OnlineFlowState.Idle;

		MatchStartsAt =
			0f;

		ErrorMessage =
			"";
	}
}