EE2Clone/Assets/Scripts/Hybrid/SelectionManager.cs

using System.Collections.Generic;
using Unity.Entities;
using Unity.Mathematics;
using Unity.Physics;
using Unity.Transforms;
using UnityEngine;
using EE2Clone.Components;

namespace EE2Clone.Hybrid
{
    /// <summary>
    /// Handles click-select, shift-add, and box-drag selection of entities.
    /// Client-side only. Adds/removes SelectedTag on ECS entities.
    /// </summary>
    public class SelectionManager : MonoBehaviour
    {
        public static SelectionManager Instance { get; private set; }

        [Header("Settings")]
        [SerializeField] private LayerMask selectableLayers = ~0;
        [SerializeField] private float dragThreshold = 5f;

        public List<Entity> SelectedEntities { get; } = new();

        // Control groups: index 0-9
        private readonly List<Entity>[] _controlGroups = new List<Entity>[10];

        // Box selection state
        private bool _isDragging;
        private Vector2 _dragStart;

        // Exposed for SelectionBoxUI
        public bool IsDragging => _isDragging;
        public Vector2 DragStart => _dragStart;
        public Vector2 DragEnd { get; private set; }

        private Camera _mainCamera;

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

            for (int i = 0; i < 10; i++)
                _controlGroups[i] = new List<Entity>();
        }

        private void Start()
        {
            _mainCamera = Camera.main;
        }

        private void Update()
        {
            var input = RTSInputActions.Instance;
            if (input == null || _mainCamera == null) return;

            HandleSelection(input);
            HandleControlGroups(input);
        }

        private void HandleSelection(RTSInputActions input)
        {
            if (input.LeftClickPressed)
            {
                _dragStart = input.MousePosition;
                _isDragging = false;
            }

            if (input.LeftClickHeld)
            {
                DragEnd = input.MousePosition;
                if (Vector2.Distance(_dragStart, DragEnd) > dragThreshold)
                    _isDragging = true;
            }

            if (input.LeftClickReleased)
            {
                DragEnd = input.MousePosition;

                if (_isDragging)
                    BoxSelect(input.ShiftHeld);
                else
                    ClickSelect(input.MousePosition, input.ShiftHeld);

                _isDragging = false;
            }
        }

        private void ClickSelect(Vector2 screenPos, bool additive)
        {
            if (!additive)
                ClearSelection();

            var ray = _mainCamera.ScreenPointToRay(new Vector3(screenPos.x, screenPos.y, 0));

            // Raycast using Unity Physics (ECS)
            var world = World.DefaultGameObjectInjectionWorld;
            if (world == null || !world.IsCreated) return;

            var physicsWorld = world.EntityManager
                .CreateEntityQuery(typeof(PhysicsWorldSingleton))
                .GetSingleton<PhysicsWorldSingleton>();

            var rayInput = new RaycastInput
            {
                Start = ray.origin,
                End = ray.origin + ray.direction * 500f,
                Filter = new CollisionFilter
                {
                    BelongsTo = ~0u,
                    CollidesWith = ~0u,
                    GroupIndex = 0
                }
            };

            if (physicsWorld.CastRay(rayInput, out var hit))
            {
                var em = world.EntityManager;
                var entity = physicsWorld.Bodies[hit.RigidBodyIndex].Entity;

                if (em.HasComponent<UnitTag>(entity) || em.HasComponent<BuildingTag>(entity))
                {
                    if (additive && SelectedEntities.Contains(entity))
                    {
                        SelectedEntities.Remove(entity);
                        em.RemoveComponent<SelectedTag>(entity);
                    }
                    else if (!SelectedEntities.Contains(entity))
                    {
                        SelectedEntities.Add(entity);
                        em.AddComponent<SelectedTag>(entity);
                    }
                }
            }
        }

        private void BoxSelect(bool additive)
        {
            if (!additive)
                ClearSelection();

            var world = World.DefaultGameObjectInjectionWorld;
            if (world == null || !world.IsCreated) return;

            var em = world.EntityManager;

            var min = Vector2.Min(_dragStart, DragEnd);
            var max = Vector2.Max(_dragStart, DragEnd);

            // Query all unit entities and check if their screen position is within the box
            var query = em.CreateEntityQuery(typeof(UnitTag), typeof(LocalTransform));
            var entities = query.ToEntityArray(Unity.Collections.Allocator.Temp);
            var transforms = query.ToComponentDataArray<LocalTransform>(Unity.Collections.Allocator.Temp);

            for (int i = 0; i < entities.Length; i++)
            {
                var screenPos = _mainCamera.WorldToScreenPoint(transforms[i].Position);
                if (screenPos.z > 0 &&
                    screenPos.x >= min.x && screenPos.x <= max.x &&
                    screenPos.y >= min.y && screenPos.y <= max.y)
                {
                    var entity = entities[i];
                    if (!SelectedEntities.Contains(entity))
                    {
                        SelectedEntities.Add(entity);
                        if (!em.HasComponent<SelectedTag>(entity))
                            em.AddComponent<SelectedTag>(entity);
                    }
                }
            }

            entities.Dispose();
            transforms.Dispose();
        }

        public void ClearSelection()
        {
            var world = World.DefaultGameObjectInjectionWorld;
            if (world != null && world.IsCreated)
            {
                var em = world.EntityManager;
                foreach (var entity in SelectedEntities)
                {
                    if (em.Exists(entity) && em.HasComponent<SelectedTag>(entity))
                        em.RemoveComponent<SelectedTag>(entity);
                }
            }
            SelectedEntities.Clear();
        }

        private void HandleControlGroups(RTSInputActions input)
        {
            int group = input.ControlGroupPressed;
            if (group < 0) return;

            if (input.CtrlHeld)
            {
                // Assign current selection to control group
                _controlGroups[group].Clear();
                _controlGroups[group].AddRange(SelectedEntities);
            }
            else
            {
                // Recall control group
                ClearSelection();
                var world = World.DefaultGameObjectInjectionWorld;
                if (world == null || !world.IsCreated) return;
                var em = world.EntityManager;

                foreach (var entity in _controlGroups[group])
                {
                    if (em.Exists(entity))
                    {
                        SelectedEntities.Add(entity);
                        if (!em.HasComponent<SelectedTag>(entity))
                            em.AddComponent<SelectedTag>(entity);
                    }
                }
            }
        }

        public Entity[] GetSelectedEntitiesOfType<T>() where T : unmanaged, IComponentData
        {
            var world = World.DefaultGameObjectInjectionWorld;
            if (world == null || !world.IsCreated) return System.Array.Empty<Entity>();

            var em = world.EntityManager;
            var result = new List<Entity>();
            foreach (var entity in SelectedEntities)
            {
                if (em.Exists(entity) && em.HasComponent<T>(entity))
                    result.Add(entity);
            }
            return result.ToArray();
        }

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