BABA_YAGA/Assets/Scripts/UI/LobbyController.cs

using UnityEngine;
using UnityEngine.UIElements;
using System.Collections.Generic;
using System.Threading.Tasks;
using Fusion;
using System.Linq;

namespace Hallucinate.UI
{
    public class LobbyController : BaseUIController
    {
        private VisualTreeAsset _roomItemTemplate;
        private PlayerDataManager _playerDataManager;

        private VisualElement _joinContainer, _createContainer, _loungeContainer, _passOverlay;
        private TextField _roomIDInput, _roomNameInput, _roomPassInput;
        private Toggle _passToggle;
        private Label _createErrorLabel;
        private Button _confirmCreateBtn;
        private ScrollView _roomList;
        private TextField _joinPassInput;
        private Label _joinPassError;
        private SessionInfo _selectedSession;
        private Label _loungeRoomName;
        private Button _readyBtn, _startBtn;
        private Label _hostNameLabel, _hostStatusLabel;
        private VisualElement _hostChatBox;
        private Label _hostChatMessage;
        private Label _guestNameLabel, _guestStatusLabel;
        private VisualElement _guestChatBox;
        private Label _guestChatMessage;
        private TextField _chatInput;

        private VisualElement _hostAvatar, _guestAvatar;
        private VisualElement _transferHostOverlay;
        private Button _confirmTransferBtn, _closeTransferBtn;
        private PlayerRef _pendingTransferPlayer;

        private bool _isBusy = false;

        public override void Initialize(VisualElement uxmlRoot, UIManager manager)
        {
            base.Initialize(uxmlRoot, manager);
            _joinContainer = root.Q<VisualElement>("JoinContainer");
            _createContainer = root.Q<VisualElement>("CreateContainer");
            _loungeContainer = root.Q<VisualElement>("LoungeContainer");
            _passOverlay = root.Q<VisualElement>("PasswordOverlay");
            _roomIDInput = root.Q<TextField>("RoomIDInput");
            _roomNameInput = root.Q<TextField>("RoomNameInput");
            _roomPassInput = root.Q<TextField>("RoomPassInput");
            _passToggle = root.Q<Toggle>("PassToggle");
            _createErrorLabel = root.Q<Label>("CreateErrorLabel");
            _roomList = root.Q<ScrollView>("RoomList");
            _joinPassInput = root.Q<TextField>("JoinPassInput");
            _joinPassError = root.Q<Label>("JoinPassError");
            _loungeRoomName = root.Q<Label>("LoungeRoomName");
            _readyBtn = root.Q<Button>("ReadyBtn");
            _startBtn = root.Q<Button>("StartBtn");
            if (_startBtn != null) _startBtn.style.display = DisplayStyle.None; // Default to hidden
            _hostNameLabel = root.Q<Label>("HostName");
            _hostStatusLabel = root.Q<Label>("HostReadyStatus");
            _hostChatBox = root.Q<VisualElement>("HostChatBox");
            _hostChatMessage = root.Q<Label>("HostChatMessage");
            _guestNameLabel = root.Q<Label>("GuestName");
            _guestStatusLabel = root.Q<Label>("GuestReadyStatus");
            _guestChatBox = root.Q<VisualElement>("GuestChatBox");
            _guestChatMessage = root.Q<Label>("GuestChatMessage");
            _chatInput = root.Q<TextField>("ChatInput");

            _hostAvatar = root.Q<VisualElement>("HostAvatar");
            _guestAvatar = root.Q<VisualElement>("GuestAvatar");
            _transferHostOverlay = root.Q<VisualElement>("TransferHostOverlay");
            _confirmTransferBtn = root.Q<Button>("ConfirmTransferBtn");
            _closeTransferBtn = root.Q<Button>("CloseTransferBtn");

            root.Q<Button>("GoToCreateBtn").clicked += ShowCreate;
            root.Q<Button>("CancelCreateBtn").clicked += ShowJoin;
            root.Q<Button>("BackToMenuBtn").clicked += async () => { if (_isBusy) return; await uiManager.Pop(); };
            _confirmCreateBtn = root.Q<Button>("ConfirmCreateBtn");
            if (_confirmCreateBtn != null) _confirmCreateBtn.clicked += OnCreateRoomClicked;
            root.Q<Button>("ConfirmJoinBtn").clicked += OnConfirmPasswordClicked;
            root.Q<Button>("ClosePassBtn").clicked += () => { if(_passOverlay != null) _passOverlay.style.display = DisplayStyle.None; };
            root.Q<Button>("LeaveLoungeBtn").clicked += OnLeaveLoungeClicked;

            if (_readyBtn != null) _readyBtn.clicked += OnReadyClicked;
            if (_startBtn != null) _startBtn.clicked += OnStartClicked;
            if (_passToggle != null) _passToggle.RegisterValueChangedCallback(evt => { if (_roomPassInput != null) _roomPassInput.style.display = evt.newValue ? DisplayStyle.Flex : DisplayStyle.None; });
            if (_chatInput != null) _chatInput.RegisterCallback<KeyDownEvent>(OnChatKeyDown, TrickleDown.TrickleDown);
            
            if (_hostAvatar != null) _hostAvatar.RegisterCallback<ClickEvent>(evt => OnAvatarClicked(true));
            if (_guestAvatar != null) _guestAvatar.RegisterCallback<ClickEvent>(evt => OnAvatarClicked(false));
            if (_confirmTransferBtn != null) _confirmTransferBtn.clicked += OnConfirmTransferHost;
            if (_closeTransferBtn != null) _closeTransferBtn.clicked += () => { if(_transferHostOverlay != null) _transferHostOverlay.style.display = DisplayStyle.None; };

            if (LocalizationManager.Instance != null) { LocalizationManager.Instance.OnLanguageChanged += ApplyLocalization; ApplyLocalization(); }
            if (BasicSpawner.Instance != null) RegisterSpawnerEvents();
            else Invoke(nameof(RegisterSpawnerEvents), 0.1f);
        }

        private void RegisterSpawnerEvents()
        {
            if (BasicSpawner.Instance == null) return;
            BasicSpawner.Instance.OnSessionListUpdatedEvent += UpdateRoomList;
            BasicSpawner.Instance.OnJoinFailedEvent += () => { if(_joinPassError != null) _joinPassError.style.display = DisplayStyle.Flex; };
        }

        private void OnChatKeyDown(KeyDownEvent evt)
        {
            if (evt.keyCode == KeyCode.Return || evt.keyCode == KeyCode.KeypadEnter)
            {
                evt.StopImmediatePropagation();
                evt.PreventDefault();
                string msg = _chatInput.value.Trim();
                if (!string.IsNullOrEmpty(msg) && PlayerDataManager.Instance != null)
                {
                    var runner = BasicSpawner.Instance.Runner;
                    if (runner != null) { PlayerDataManager.Instance.RPC_SendChatMessage(runner.LocalPlayer, msg); _chatInput.value = ""; _chatInput.Focus(); }
                }
            }
        }

        private void OnChatMessageReceived(PlayerRef sender, string message)
        {
            var runner = BasicSpawner.Instance.Runner;
            if (runner == null) return;
            var sortedPlayers = runner.ActivePlayers.OrderBy(p => p.PlayerId).ToList();
            bool isHost = sortedPlayers.Count > 0 && sender == sortedPlayers[0];
            if (isHost) ShowChatBubble(_hostChatBox, _hostChatMessage, message);
            else ShowChatBubble(_guestChatBox, _guestChatMessage, message);
        }

        private async void ShowChatBubble(VisualElement box, Label label, string msg)
        {
            if (box == null || label == null) return;
            label.text = msg;
            box.style.display = DisplayStyle.Flex;
            await Task.Delay(4000);
            if (label.text == msg) box.style.display = DisplayStyle.None;
        }

        private void ApplyLocalization() { if (LocalizationManager.Instance == null) return; }
        private string GetT(string key) => LocalizationManager.Instance != null ? LocalizationManager.Instance.GetLocalizedString(key) : key;
        public void SetRoomTemplate(VisualTreeAsset template) => _roomItemTemplate = template;
        
        public override async Task PlayTransitionIn() 
        { 
            _isBusy = false; 
            await base.PlayTransitionIn(); 
            ShowJoin(); 
        }

        public void ShowJoin()
        {
            if (_joinContainer != null) _joinContainer.style.display = DisplayStyle.Flex;
            if (_createContainer != null) _createContainer.style.display = DisplayStyle.None;
            if (_loungeContainer != null) _loungeContainer.style.display = DisplayStyle.None;
            var spawner = BasicSpawner.Instance;
            if (spawner != null && (spawner.Runner == null || !spawner.Runner.IsRunning)) _ = spawner.StartLobby();
        }

        public void ShowCreate()
        {
            if (_joinContainer != null) _joinContainer.style.display = DisplayStyle.None;
            if (_createContainer != null) _createContainer.style.display = DisplayStyle.Flex;
            if (_createErrorLabel != null) _createErrorLabel.style.display = DisplayStyle.None;
            if (_confirmCreateBtn != null) _confirmCreateBtn.SetEnabled(true);
            if (_roomIDInput != null) _roomIDInput.value = "ROOM_" + Random.Range(1000, 9999).ToString();
        }

        private void ShowLounge(string roomName)
        {
            if (_joinContainer != null) _joinContainer.style.display = DisplayStyle.None;
            if (_createContainer != null) _createContainer.style.display = DisplayStyle.None;
            if (_loungeContainer != null) _loungeContainer.style.display = DisplayStyle.Flex;
            if (_loungeRoomName != null) _loungeRoomName.text = roomName.ToUpper();
            var spawner = BasicSpawner.Instance;
            if (spawner != null && spawner.Runner != null && spawner.Runner.SessionInfo != null)
            {
                var loungeIdLabel = root.Q<Label>("LoungeID");
                if (loungeIdLabel != null) loungeIdLabel.text = GetT("LOBBY_ID_PREFIX") + spawner.Runner.SessionInfo.Name;
            }
            _playerDataManager = PlayerDataManager.Instance;
            if (_playerDataManager != null) _playerDataManager.OnChatMessageReceived += OnChatMessageReceived;
        }

        private async void OnCreateRoomClicked()
        {
            if (_isBusy) return;
            _isBusy = true;

            if (_confirmCreateBtn != null) _confirmCreateBtn.SetEnabled(false);
            if (_createErrorLabel != null) _createErrorLabel.style.display = DisplayStyle.None;
            
            var spawner = BasicSpawner.Instance;
            if (spawner == null) 
            { 
                ShowCreateError("System Error: Spawner missing. Please re-enter the Lobby."); 
                _isBusy = false;
                return; 
            }
            
            string id = (_roomIDInput != null && !string.IsNullOrEmpty(_roomIDInput.value)) ? _roomIDInput.value.Trim() : "ROOM_" + Random.Range(1000, 9999).ToString();
            string name = (_roomNameInput != null && !string.IsNullOrEmpty(_roomNameInput.value)) ? _roomNameInput.value.Trim() : id;
            string pass = (_passToggle != null && _passToggle.value && _roomPassInput != null) ? _roomPassInput.value : null;
            
            try 
            {
                bool success = await spawner.StartHost(id, name, pass);
                if (success) ShowLounge(name);
                else ShowCreateError("Failed to create room. ID might be taken.");
            } 
            catch (System.Exception ex) 
            { 
                ShowCreateError("Network Error: " + ex.Message); 
            }
            finally
            {
                _isBusy = false;
                if (_confirmCreateBtn != null) _confirmCreateBtn.SetEnabled(true);
            }
        }

        private void ShowCreateError(string message)
        {
            if (_createErrorLabel != null) { _createErrorLabel.text = message; _createErrorLabel.style.display = DisplayStyle.Flex; }
            if (_confirmCreateBtn != null) _confirmCreateBtn.SetEnabled(true);
        }

        private void UpdateRoomList(List<SessionInfo> sessions)
        {
            if (_roomList == null) return;
            _roomList.Clear();
            foreach (var session in sessions)
            {
                if (_roomItemTemplate == null) continue;
                var item = _roomItemTemplate.Instantiate();
                string displayName = session.Name;
                if (session.Properties.TryGetValue("rn", out var rnProp)) displayName = rnProp;
                item.Q<Label>("RoomName").text = displayName;
                item.Q<Label>("PlayerCount").text = $"{session.PlayerCount}/{session.MaxPlayers}";
                var statusBadge = item.Q<Label>("StatusBadge");
                if (statusBadge != null) statusBadge.text = GetT("ROOM_STATUS_WAITING");
                bool needsPass = session.Properties.ContainsKey("pw"); 
                var lockIcon = item.Q<Label>("LockIcon");
                if (lockIcon != null) lockIcon.style.display = needsPass ? DisplayStyle.Flex : DisplayStyle.None;
                var joinBtn = item.Q<Button>("JoinBtn");
                if (joinBtn != null) { joinBtn.text = GetT("ROOM_JOIN_BTN"); joinBtn.clicked += () => OnRoomItemClicked(session); }
                _roomList.Add(item);
            }
        }

        private async void OnRoomItemClicked(SessionInfo session)
        {
            if (_isBusy) return;
            
            bool needsPass = session.Properties.ContainsKey("pw");
            if (needsPass)
            {
                _selectedSession = session;
                if (_passOverlay != null) _passOverlay.style.display = DisplayStyle.Flex;
                if (_joinPassError != null) _joinPassError.style.display = DisplayStyle.None;
                if (_joinPassInput != null) _joinPassInput.value = "";
            }
            else await JoinRoom(session.Name, null);
        }

        private async void OnConfirmPasswordClicked()
        {
            if (_isBusy || _selectedSession == null) return;
            _isBusy = true;
            
            string pass = _joinPassInput != null ? _joinPassInput.value : "";
            if (_passOverlay != null) _passOverlay.style.display = DisplayStyle.None;
            await JoinRoom(_selectedSession.Name, pass);
            _isBusy = false;
        }

        private async Task JoinRoom(string sessionName, string password) 
        { 
            if (BasicSpawner.Instance != null) 
            { 
                if (await BasicSpawner.Instance.StartClient(sessionName, password)) ShowLounge(sessionName); 
            } 
        }
        
        private void OnReadyClicked() 
        { 
            if (_isBusy) return;
            var runner = BasicSpawner.Instance.Runner; 
            if (runner != null && _playerDataManager != null && _playerDataManager.TryGetPlayerMetaData(runner.LocalPlayer, out var myData)) 
                _playerDataManager.RPC_SetReady(runner.LocalPlayer, !myData.IsReady); 
        }
        
        private void OnStartClicked() 
        { 
            if (_isBusy) return;
            BasicSpawner.Instance?.StartGame(); 
        }

        private void OnAvatarClicked(bool isHostSlot)
        {
            if (_playerDataManager == null || _playerDataManager.Leader != BasicSpawner.Instance.Runner.LocalPlayer) return;

            var runner = BasicSpawner.Instance.Runner;
            var sortedPlayers = runner.ActivePlayers.OrderBy(p => p.PlayerId).ToList();
            
            PlayerRef clickedPlayer = PlayerRef.None;
            if (isHostSlot && sortedPlayers.Count > 0) clickedPlayer = sortedPlayers[0];
            else if (!isHostSlot && sortedPlayers.Count > 1) clickedPlayer = sortedPlayers[1];

            if (clickedPlayer != PlayerRef.None && clickedPlayer != runner.LocalPlayer)
            {
                _pendingTransferPlayer = clickedPlayer;
                if (_transferHostOverlay != null) _transferHostOverlay.style.display = DisplayStyle.Flex;
            }
        }

        private void OnConfirmTransferHost()
        {
            if (_pendingTransferPlayer != PlayerRef.None && _playerDataManager != null)
            {
                _playerDataManager.RPC_TransferLeader(_pendingTransferPlayer);
                if (_transferHostOverlay != null) _transferHostOverlay.style.display = DisplayStyle.None;
            }
        }
        
        private async void OnLeaveLoungeClicked() 
        { 
            if (_isBusy) return;
            _isBusy = true;
            
            var runner = BasicSpawner.Instance.Runner; 
            if (runner != null) await runner.Shutdown(false); 
            if (_playerDataManager != null) _playerDataManager.OnChatMessageReceived -= OnChatMessageReceived; 
            ShowJoin(); 
            
            _isBusy = false;
        }

        public override void Update() { if (_loungeContainer != null && _loungeContainer.style.display == DisplayStyle.Flex) UpdateLoungeUI(); }

        private void UpdateLoungeUI()
        {
            var spawner = BasicSpawner.Instance;
            if (spawner == null) return;
            var runner = spawner.Runner;
            if (runner == null) return;

            // 1. PlayerDataManager Sync
            if (_playerDataManager == null || _playerDataManager.Object == null || !_playerDataManager.Object.IsValid)
            {
                _playerDataManager = PlayerDataManager.Instance;
                if (_playerDataManager != null) 
                {
                    _playerDataManager.OnChatMessageReceived += OnChatMessageReceived;
                }
                else return; // Still waiting for synchronization
            }

            // 2. Identify Players
            var sortedPlayers = runner.ActivePlayers.OrderBy(p => p.PlayerId).ToList();
            PlayerRef hostRef = sortedPlayers.Count > 0 ? sortedPlayers[0] : PlayerRef.None;
            PlayerRef guestRef = sortedPlayers.Count > 1 ? sortedPlayers[1] : PlayerRef.None;
            PlayerRef leaderRef = _playerDataManager.Leader;

            // 3. Strict Visibility Check (Leader Only)
            bool isLeader = runner.LocalPlayer == leaderRef;
            if (_startBtn != null) 
            {
                _startBtn.style.display = isLeader ? DisplayStyle.Flex : DisplayStyle.None;
            }

            if (runner.SessionInfo != null && runner.SessionInfo.Properties.TryGetValue("rn", out var rnProp)) _loungeRoomName.text = rnProp.ToString().ToUpper();

            // Host Display
            if (hostRef != PlayerRef.None && _playerDataManager.TryGetPlayerMetaData(hostRef, out var hostData))
            {
                _hostNameLabel.text = hostData.Name.ToString().ToUpper();
                string readyStatus = hostData.IsReady ? GetT("LOBBY_READY") : GetT("LOBBY_NOT_READY");
                string roleLabel = (hostRef == leaderRef) ? GetT("LOBBY_HOST_LABEL") : GetT("LOBBY_PLAYER_LABEL");
                _hostStatusLabel.text = $"{roleLabel} - {readyStatus}";
                _hostStatusLabel.style.color = hostData.IsReady ? Color.green : Color.red;
                
                // Highlight leader avatar
                if (_hostAvatar != null)
                {
                    float width = (hostRef == leaderRef) ? 2f : 0f;
                    _hostAvatar.style.borderTopWidth = width;
                    _hostAvatar.style.borderBottomWidth = width;
                    _hostAvatar.style.borderLeftWidth = width;
                    _hostAvatar.style.borderRightWidth = width;
                }
            }
            else { _hostNameLabel.text = GetT("LOBBY_SYNCING"); _hostStatusLabel.text = "-"; }

            // Guest Display
            if (guestRef != PlayerRef.None && _playerDataManager.TryGetPlayerMetaData(guestRef, out var guestData))
            {
                _guestNameLabel.text = guestData.Name.ToString().ToUpper();
                string readyStatus = guestData.IsReady ? GetT("LOBBY_READY") : GetT("LOBBY_NOT_READY");
                string roleLabel = (guestRef == leaderRef) ? GetT("LOBBY_HOST_LABEL") : GetT("LOBBY_PLAYER_LABEL");
                _guestStatusLabel.text = $"{roleLabel} - {readyStatus}";
                _guestStatusLabel.style.color = guestData.IsReady ? Color.green : Color.red;

                // Highlight leader avatar
                if (_guestAvatar != null)
                {
                    float width = (guestRef == leaderRef) ? 2f : 0f;
                    _guestAvatar.style.borderTopWidth = width;
                    _guestAvatar.style.borderBottomWidth = width;
                    _guestAvatar.style.borderLeftWidth = width;
                    _guestAvatar.style.borderRightWidth = width;
                    _guestAvatar.style.borderTopColor = Color.white;
                    _guestAvatar.style.borderBottomColor = Color.white;
                    _guestAvatar.style.borderLeftColor = Color.white;
                    _guestAvatar.style.borderRightColor = Color.white;
                }
            }
            else if (runner.ActivePlayers.Count() >= 2) { _guestNameLabel.text = GetT("LOBBY_SYNCING"); _guestStatusLabel.text = "-"; }
            else 
            { 
                _guestNameLabel.text = GetT("LOBBY_WAITING_LABEL"); 
                _guestStatusLabel.text = "-"; 
                _guestStatusLabel.style.color = Color.gray; 
                if (_guestAvatar != null)
                {
                    _guestAvatar.style.borderTopWidth = 0;
                    _guestAvatar.style.borderBottomWidth = 0;
                    _guestAvatar.style.borderLeftWidth = 0;
                    _guestAvatar.style.borderRightWidth = 0;
                }
            }

            // 4. Start Button Logic (Leader Only)
            if (_startBtn != null && isLeader)
            {
                bool allReady = true;
                int playerCount = 0;
                foreach (var p in runner.ActivePlayers) 
                { 
                    playerCount++; 
                    if (_playerDataManager.TryGetPlayerMetaData(p, out var data)) 
                    { 
                        if (!data.IsReady) allReady = false; 
                    } 
                    else allReady = false; 
                }

                bool canStart = allReady && playerCount >= 2;
                _startBtn.SetEnabled(canStart);
                _startBtn.text = GetT("LOBBY_START_BTN");

                if (canStart)
                {
                    _startBtn.style.backgroundColor = new StyleColor(Color.green);
                    _startBtn.style.color = new StyleColor(Color.black);
                }
                else
                {
                    _startBtn.style.backgroundColor = new StyleColor(new Color(0.2f, 0.2f, 0.2f, 0.8f));
                    _startBtn.style.color = new StyleColor(new Color(1f, 1f, 1f, 0.5f));
                }
            }

            // Ready Button Logic (Everyone)
            if (_readyBtn != null)
            {
                if (_playerDataManager.TryGetPlayerMetaData(runner.LocalPlayer, out var myData))
                {
                    if (myData.IsReady) 
                    { 
                        _readyBtn.text = GetT("LOBBY_UNREADY_BTN"); 
                        _readyBtn.style.backgroundColor = new StyleColor(Color.green); 
                        _readyBtn.style.color = new StyleColor(Color.black); 
                    }
                    else 
                    { 
                        _readyBtn.text = GetT("LOBBY_READY_BTN"); 
                        _readyBtn.style.backgroundColor = new StyleColor(new Color(0.2f, 0.2f, 0.2f, 0.8f)); 
                        _readyBtn.style.color = new StyleColor(Color.white); 
                    }
                }
            }
        }

        private async void Invoke(string methodName, float delay) { await Task.Delay((int)(delay * 1000)); if (methodName == nameof(RegisterSpawnerEvents)) RegisterSpawnerEvents(); }
    }
}