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;

        // Containers
        private VisualElement _joinContainer, _createContainer, _loungeContainer, _passOverlay;
        
        // Create Room Fields
        private TextField _roomIDInput, _roomNameInput, _roomPassInput;
        private Toggle _passToggle;
        
        // Join Room Fields
        private ScrollView _roomList;
        private TextField _joinPassInput;
        private Label _joinPassError;
        private SessionInfo _selectedSession;

        // Lounge Elements
        private Label _loungeRoomName;
        private Button _readyBtn, _startBtn;
        
        // Host Slot
        private Label _hostNameLabel, _hostStatusLabel;
        private VisualElement _hostChatBox;
        private Label _hostChatMessage;

        // Guest Slot
        private Label _guestNameLabel, _guestStatusLabel;
        private VisualElement _guestChatBox;
        private Label _guestChatMessage;

        // Chat Input
        private TextField _chatInput;

        public override void Initialize(VisualElement uxmlRoot, UIManager manager)
        {
            base.Initialize(uxmlRoot, manager);

            // Query Containers
            _joinContainer = root.Q<VisualElement>("JoinContainer");
            _createContainer = root.Q<VisualElement>("CreateContainer");
            _loungeContainer = root.Q<VisualElement>("LoungeContainer");
            _passOverlay = root.Q<VisualElement>("PasswordOverlay");

            // Create Room Fields
            _roomIDInput = root.Q<TextField>("RoomIDInput");
            _roomNameInput = root.Q<TextField>("RoomNameInput");
            _roomPassInput = root.Q<TextField>("RoomPassInput");
            _passToggle = root.Q<Toggle>("PassToggle");

            // Join Room Fields
            _roomList = root.Q<ScrollView>("RoomList");
            _joinPassInput = root.Q<TextField>("JoinPassInput");
            _joinPassError = root.Q<Label>("JoinPassError");

            // Lounge Elements
            _loungeRoomName = root.Q<Label>("LoungeRoomName");
            _readyBtn = root.Q<Button>("ReadyBtn");
            _startBtn = root.Q<Button>("StartBtn");

            // Host Slot
            _hostNameLabel = root.Q<Label>("HostName");
            _hostStatusLabel = root.Q<Label>("HostReadyStatus");
            _hostChatBox = root.Q<VisualElement>("HostChatBox");
            _hostChatMessage = root.Q<Label>("HostChatMessage");

            // Guest Slot
            _guestNameLabel = root.Q<Label>("GuestName");
            _guestStatusLabel = root.Q<Label>("GuestReadyStatus");
            _guestChatBox = root.Q<VisualElement>("GuestChatBox");
            _guestChatMessage = root.Q<Label>("GuestChatMessage");

            // Chat Input
            _chatInput = root.Q<TextField>("ChatInput");

            // Event Bindings
            root.Q<Button>("GoToCreateBtn").clicked += ShowCreate;
            root.Q<Button>("CancelCreateBtn").clicked += ShowJoin;
            root.Q<Button>("BackToMenuBtn").clicked += async () => await uiManager.Pop();
            root.Q<Button>("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);
            }

            // Đăng ký sự kiện từ Spawner
            if (BasicSpawner.Instance != null)
            {
                RegisterSpawnerEvents();
                _ = BasicSpawner.Instance.StartLobby();
            }
            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; };
            BasicSpawner.Instance.OnJoinStartedEvent += () => { };
            _ = BasicSpawner.Instance.StartLobby();
        }

        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 = Object.FindFirstObjectByType<NetworkRunner>();
                    if (runner != null)
                    {
                        PlayerDataManager.Instance.RPC_SendChatMessage(runner.LocalPlayer, msg);
                        _chatInput.value = "";
                        // Re-focus after clearing
                        _chatInput.Focus();
                    }
                }
            }
        }

        private void OnChatMessageReceived(PlayerRef sender, string message)
        {
            var runner = Object.FindFirstObjectByType<NetworkRunner>();
            if (runner == null) return;

            // Kiểm tra sender là Host hay Guest
            bool isHost = sender.PlayerId == 1; // Trong Host Mode, người tạo phòng luôn có ID 1

            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) // Chỉ ẩn nếu chưa có tin nhắn mới đè lên
                box.style.display = DisplayStyle.None;
        }

        public void SetRoomTemplate(VisualTreeAsset template) => _roomItemTemplate = template;

        public override async Task PlayTransitionIn()
        {
            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;
            _ = BasicSpawner.Instance?.StartLobby();
        }

        public void ShowCreate()
        {
            if (_joinContainer != null) _joinContainer.style.display = DisplayStyle.None;
            if (_createContainer != null) _createContainer.style.display = DisplayStyle.Flex;
        }

        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();
            
            _playerDataManager = Object.FindFirstObjectByType<PlayerDataManager>();
            if (_playerDataManager != null)
            {
                _playerDataManager.OnChatMessageReceived += OnChatMessageReceived;
            }
        }

        private async void OnCreateRoomClicked()
        {
            var spawner = BasicSpawner.Instance;
            if (spawner == null) return;

            string id = _roomIDInput != null && !string.IsNullOrEmpty(_roomIDInput.value) 
                ? _roomIDInput.value.Trim() 
                : Random.Range(1000, 9999).ToString();

            if (_roomIDInput != null) _roomIDInput.value = id;

            string name = _roomNameInput != null && !string.IsNullOrEmpty(_roomNameInput.value) 
                ? _roomNameInput.value 
                : $"Room {id}";
                
            string pass = (_passToggle != null && _passToggle.value && _roomPassInput != null) 
                ? _roomPassInput.value 
                : null;

            bool success = await spawner.StartHost(id, name, pass);
            if (success) ShowLounge(name);
        }

        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();
                
                // Hiển thị tên phòng thân thiện nếu có
                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}";
                
                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.clicked += () => OnRoomItemClicked(session);
                
                _roomList.Add(item);
            }
        }

        private async void OnRoomItemClicked(SessionInfo session)
        {
            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 (_selectedSession == null) return;
            string pass = _joinPassInput != null ? _joinPassInput.value : "";
            if (_passOverlay != null) _passOverlay.style.display = DisplayStyle.None;
            await JoinRoom(_selectedSession.Name, pass);
        }

        private async Task JoinRoom(string sessionName, string password)
        {
            if (BasicSpawner.Instance != null)
            {
                bool success = await BasicSpawner.Instance.StartClient(sessionName, password);
                if (success) ShowLounge(sessionName);
            }
        }

        private void OnReadyClicked()
        {
            var runner = Object.FindFirstObjectByType<NetworkRunner>();
            if (runner != null && _playerDataManager != null)
            {
                if (_playerDataManager.TryGetPlayerMetaData(runner.LocalPlayer, out var myData))
                {
                    _playerDataManager.RPC_SetReady(runner.LocalPlayer, !myData.IsReady);
                }
            }
        }

        private void OnStartClicked()
        {
            BasicSpawner.Instance?.StartGame();
        }

        private void OnLeaveLoungeClicked()
        {
            var runner = Object.FindFirstObjectByType<NetworkRunner>();
            if (runner != null)
            {
                runner.Shutdown();
            }
            if (_playerDataManager != null)
            {
                _playerDataManager.OnChatMessageReceived -= OnChatMessageReceived;
            }
            ShowJoin();
        }

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

        private void UpdateLoungeUI()
        {
            var runner = Object.FindFirstObjectByType<NetworkRunner>();
            if (runner == null) return;

            if (_playerDataManager == null)
            {
                _playerDataManager = Object.FindFirstObjectByType<PlayerDataManager>();
                if (_playerDataManager != null)
                {
                    _playerDataManager.OnChatMessageReceived += OnChatMessageReceived;
                }
            }

            if (_playerDataManager == null || !_playerDataManager.Object.IsValid) return;

            PlayerRef hostRef = PlayerRef.None;
            PlayerRef guestRef = PlayerRef.None;

            foreach (var p in runner.ActivePlayers)
            {
                if (runner.IsPlayerServer(p))
                {
                    hostRef = p;
                }
                else
                {
                    guestRef = p;
                }
            }

            // Update Room Name for Guest
            if (runner.SessionInfo != null && runner.SessionInfo.Properties.TryGetValue("rn", out var rnProp))
            {
                _loungeRoomName.text = rnProp.ToString().ToUpper();
            }

            // Update Host UI
            if (hostRef != PlayerRef.None && _playerDataManager.TryGetPlayerMetaData(hostRef, out var hostData))
            {
                _hostNameLabel.text = hostData.Name.ToString().ToUpper();
                _hostStatusLabel.text = hostData.IsReady ? "READY" : "NOT READY";
                _hostStatusLabel.style.color = hostData.IsReady ? Color.green : Color.red;
            }
            else if (hostRef != PlayerRef.None)
            {
                _hostNameLabel.text = "SYNCING...";
                _hostStatusLabel.text = "-";
            }

            // Update Guest UI
            if (guestRef != PlayerRef.None && _playerDataManager.TryGetPlayerMetaData(guestRef, out var guestData))
            {
                _guestNameLabel.text = guestData.Name.ToString().ToUpper();
                _guestStatusLabel.text = guestData.IsReady ? "READY" : "NOT READY";
                _guestStatusLabel.style.color = guestData.IsReady ? Color.green : Color.red;
            }
            else if (runner.ActivePlayers.Count() >= 2)
            {
                _guestNameLabel.text = "SYNCING...";
                _guestStatusLabel.text = "-";
            }
            else
            {
                _guestNameLabel.text = "WAITING...";
                _guestStatusLabel.text = "-";
                _guestStatusLabel.style.color = Color.gray;
            }

            // Start Button visibility logic
            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 isHost = runner.LocalPlayer == hostRef;

            if (_startBtn != null)
            {
                _startBtn.style.display = isHost ? DisplayStyle.Flex : DisplayStyle.None;
                _startBtn.SetEnabled(allReady && playerCount >= 2);
            }

            if (_readyBtn != null)
            {
                if (_playerDataManager.TryGetPlayerMetaData(runner.LocalPlayer, out var myData))
                {
                    // Style for Ready Button
                    if (myData.IsReady)
                    {
                        _readyBtn.text = "UNREADY";
                        _readyBtn.style.backgroundColor = new StyleColor(Color.green);
                        _readyBtn.style.color = new StyleColor(Color.black);
                    }
                    else
                    {
                        _readyBtn.text = "READY UP";
                        _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();
        }
    }
}