using System.Collections;
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;

    [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(Transform enemyTransform)
    {
        movementJobHandle.Complete();
        transformAccessArray.Add(enemyTransform);
        
        activeEnemyCount++;
        if (UIManager.Instance != null) UIManager.Instance.UpdateEnemiesRemaining(activeEnemyCount);
    }

    public void RemoveEnemy(Transform enemyTransform)
    {
        movementJobHandle.Complete();
        for (int i = 0; i < transformAccessArray.length; i++)
        {
            if (transformAccessArray[i] == enemyTransform)
            {
                transformAccessArray.RemoveAtSwapBack(i);
                
                activeEnemyCount--;
                if (UIManager.Instance != null) UIManager.Instance.UpdateEnemiesRemaining(activeEnemyCount);
                
                CheckWaveClear(); 
                break;
            }
        }
    }

    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;
        
        targetPositions[0] = SamplePlayerAnimationController.Instance.transform.position;

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

        movementJobHandle = swarmJob.Schedule(transformAccessArray);
    }

    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);
    }
}