using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using UnityEngine.Jobs;
using Unity.Jobs;
using Unity.Burst;
using Unity.Mathematics;
using Unity.Collections;
using Synty.AnimationBaseLocomotion.Samples;

// --- WAVE DATA STRUCTURES ---
[System.Serializable]
public class WaveSpawnData
{
    public GameObject enemyPrefab;
    public int count;
}

[System.Serializable]
public class WaveDefinition
{
    public string waveName = "New Wave";
    public WaveSpawnData[] enemyGroups;
}

public class EnemySwarmManager : MonoBehaviour
{
    public static EnemySwarmManager Instance;
    private List<EnemyAI> activeEnemiesList = new List<EnemyAI>();

    [Header("Swarm Settings")]
    public float enemySpeed = 3.5f;

    [Header("Wave Management")]
    public WaveDefinition[] waves;
    public Transform[] spawnPoints;
    public float timeBetweenWaves = 5.0f;
    [Tooltip("Delay between individual enemy spawns to prevent lag spikes")]
    public float spawnStagger = 0.05f;

    private int currentWaveIndex = 0;
    private int activeEnemyCount = 0;
    private bool isSpawning = false;
    private bool gameEnded = false;

    // Burst Compiler variables
    private TransformAccessArray transformAccessArray;
    private NativeArray<float3> targetPositions;
    private JobHandle movementJobHandle;

    void Awake()
    {
        Instance = this;
        transformAccessArray = new TransformAccessArray(1000);
        //targetPositions = new NativeArray<float3>(1, Allocator.Persistent);
    }

    void Start()
    {
        if (waves.Length > 0)
        {
            UIManager.Instance.UpdateWaveProgress(1, waves.Length);
            UIManager.Instance.UpdateEnemiesRemaining(0);
            StartCoroutine(WaveCooldownRoutine());
        }
        else
        {
            Debug.LogWarning("No waves defined in the Swarm Manager!");
        }
    }

    // --- WAVE LOGIC ---

    private IEnumerator WaveCooldownRoutine()
    {
        isSpawning = true; // Prevent accidental wave-clears during cooldown
        Debug.Log($"<color=cyan>Next wave starting in {timeBetweenWaves} seconds...</color>");

        yield return new WaitForSeconds(timeBetweenWaves);

        StartCoroutine(SpawnWaveRoutine(waves[currentWaveIndex]));
    }

    private IEnumerator SpawnWaveRoutine(WaveDefinition wave)
    {
        Debug.Log($"<color=yellow>Starting {wave.waveName}</color>");
        isSpawning = true;

        // Iterate through each type of enemy requested in this wave
        foreach (WaveSpawnData spawnGroup in wave.enemyGroups)
        {
            for (int i = 0; i < spawnGroup.count; i++)
            {
                // Pick a random spawn point from the array
                Transform randomSpawnPoint = spawnPoints[UnityEngine.Random.Range(0, spawnPoints.Length)];

                // Add a tiny random offset so they don't spawn perfectly inside each other
                Vector3 spawnPos = randomSpawnPoint.position + new Vector3(
                    UnityEngine.Random.Range(-2f, 2f),
                    0f,
                    UnityEngine.Random.Range(-2f, 2f)
                );

                Instantiate(spawnGroup.enemyPrefab, spawnPos, randomSpawnPoint.rotation);

                // Yield briefly to spread the Instantiation load across multiple frames
                yield return new WaitForSeconds(spawnStagger);
            }
        }

        isSpawning = false;

        // Failsafe: Check if the player somehow killed them all before the spawning even finished
        CheckWaveClear();
    }

    public void RegisterEnemy(EnemyAI enemy) // <-- Changed parameter from Transform to EnemyAI
    {
        movementJobHandle.Complete();
        transformAccessArray.Add(enemy.transform);
        activeEnemiesList.Add(enemy); // Keep a parallel list on the main thread

        activeEnemyCount++;
        if (UIManager.Instance != null) UIManager.Instance.UpdateEnemiesRemaining(activeEnemyCount);
    }

    public void RemoveEnemy(EnemyAI enemy) // <-- Changed parameter from Transform to EnemyAI
    {
        movementJobHandle.Complete();
        int index = activeEnemiesList.IndexOf(enemy);

        if (index >= 0)
        {
            // Swap-back deletion keeps arrays perfectly aligned and highly performant
            transformAccessArray.RemoveAtSwapBack(index);

            activeEnemiesList[index] = activeEnemiesList[activeEnemiesList.Count - 1];
            activeEnemiesList.RemoveAt(activeEnemiesList.Count - 1);

            activeEnemyCount--;
            if (UIManager.Instance != null) UIManager.Instance.UpdateEnemiesRemaining(activeEnemyCount);

            CheckWaveClear();
        }
    }

    private void CheckWaveClear()
    {
        // Don't trigger a clear if we are currently spawning, or if the game is over
        if (isSpawning || gameEnded) return;

        if (activeEnemyCount <= 0)
        {
            Debug.Log($"<color=green>Wave {currentWaveIndex + 1} Cleared!</color>");
            currentWaveIndex++;

            if (currentWaveIndex < waves.Length)
            {
                // Update UI for the upcoming wave
                UIManager.Instance.UpdateWaveProgress(currentWaveIndex + 1, waves.Length);
                StartCoroutine(WaveCooldownRoutine());
            }
            else
            {
                gameEnded = true;
                Debug.Log("<color=gold>VICTORY! All waves cleared!</color>");
                // TODO: Trigger Victory UI
            }
        }
    }

    // --- BURST COMPILER MOVEMENT ---

    /*void Update()
    {
        if (SamplePlayerAnimationController.Instance == null || transformAccessArray.length == 0) return;
        
        // --- NEW TARGETING LOGIC ---
        // Feed every single enemy's unique target into the NativeArray before the Job runs
        for (int i = 0; i < activeEnemiesList.Count; i++)
        {
            if (activeEnemiesList[i].currentTarget != null)
                targetPositions[i] = activeEnemiesList[i].currentTarget.position;
            else
                targetPositions[i] = SamplePlayerAnimationController.Instance.transform.position; // Fallback
        }

        SwarmJob swarmJob = new SwarmJob
        {
            targetPos = targetPositions,
            speed = enemySpeed,
            dt = Time.deltaTime
        };

        movementJobHandle = swarmJob.Schedule(transformAccessArray);
    }*/

    /*[BurstCompile]
    struct SwarmJob : IJobParallelForTransform
    {
        [ReadOnly] public NativeArray<float3> targetPos;
        public float speed;
        public float dt;

        public void Execute(int index, TransformAccess transform)
        {
            float3 currentPos = transform.position;
            // The Job now looks up the SPECIFIC target for this specific index!
            float3 dir = targetPos[index] - currentPos; 
            
            if (math.lengthsq(dir) > 4.0f) 
            {
                dir = math.normalize(dir);
                transform.position = currentPos + (dir * speed * dt);
                
                quaternion targetRot = quaternion.LookRotationSafe(dir, math.up());
                transform.rotation = targetRot; 
            }
        }
    }*/

    /*void LateUpdate()
    {
        movementJobHandle.Complete();
    }*/

    /*void OnDestroy()
    {
        movementJobHandle.Complete();
        if (transformAccessArray.isCreated) transformAccessArray.Dispose();
        if (targetPositions.IsCreated) targetPositions.Dispose();
    }*/

    /* [BurstCompile]
     struct SwarmJob : IJobParallelForTransform
     {
         [ReadOnly] public NativeArray<float3> targetPos;
         public float speed;
         public float dt;

         public void Execute(int index, TransformAccess transform)
         {
             float3 currentPos = transform.position;
             float3 dir = targetPos[0] - currentPos;

             if (math.lengthsq(dir) > 4.0f) 
             {
                 dir = math.normalize(dir);
                 transform.position = currentPos + (dir * speed * dt);

                 quaternion targetRot = quaternion.LookRotationSafe(dir, math.up());
                 transform.rotation = targetRot; 
             }
         }
     }*/
    public void PauseEnemyMovement(Transform enemyTransform)
    {
        movementJobHandle.Complete();
        for (int i = 0; i < transformAccessArray.length; i++)
        {
            if (transformAccessArray[i] == enemyTransform)
            {
                transformAccessArray.RemoveAtSwapBack(i);
                break;
            }
        }
    }

    public void ResumeEnemyMovement(Transform enemyTransform)
    {
        movementJobHandle.Complete();
        transformAccessArray.Add(enemyTransform);
    }
}