using System.Collections.Generic;
using Cysharp.Threading.Tasks;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;
using UnityEngine;

namespace DanieleMarotta.RiversongCodeShowcase
{
    public abstract class ChunkMeshGenerator
    {
        private readonly List<ChunkGenerationJobState> _pendingJobs = new();

        private JobHandle _dependency;

        protected ChunkMeshGenerator(World world, GameConfig config)
        {
            World = world;
            Config = config;
        }

        protected World World { get; }

        protected GameConfig Config { get; }

        public void InitializeChunk(TerrainChunk chunk)
        {
            for (var lod = 0; lod < GetLodCount(); lod++) InitializeChunk(chunk, lod);
        }

        private void InitializeChunk(TerrainChunk chunk, int lod)
        {
            var childName = GetGameObjectName(chunk.Root.name, chunk.Coords, lod);

            var child = new GameObject(childName);
            child.transform.SetParent(chunk.Root.transform);

            var mesh = new Mesh { name = childName };
            mesh.MarkDynamic();

            var meshFilter = child.AddComponent<MeshFilter>();
            meshFilter.sharedMesh = mesh;

            var renderer = child.AddComponent<MeshRenderer>();

            InitializeChunk(chunk, lod, child, mesh, renderer);
        }

        protected abstract int GetLodCount();

        protected abstract string GetGameObjectName(string chunkName, int2 chunkCoords, int lod);

        protected abstract void InitializeChunk(TerrainChunk chunk, int lod, GameObject gameObject, Mesh mesh, Renderer renderer);

        public void BeginGeneratingChunks(List<TerrainChunk> chunks)
        {
            var chunkCoords = new NativeArray<int2>(chunks.Count, Allocator.Persistent, NativeArrayOptions.UninitializedMemory);
            for (var i = 0; i < chunks.Count; i++) chunkCoords[i] = chunks[i].Coords;

            for (var lod = 0; lod < GetLodCount(); lod++)
            {
                var jobState = CreateJobState(lod, chunks);
                _pendingJobs.Add(jobState);

                var jobHandle = ScheduleJob(jobState, chunkCoords);
                _dependency = JobHandle.CombineDependencies(_dependency, jobHandle);
            }

            chunkCoords.Dispose(_dependency);
        }

        private ChunkGenerationJobState CreateJobState(int lod, List<TerrainChunk> chunks)
        {
            var jobState = new ChunkGenerationJobState { Lod = lod };

            jobState.Chunks.AddRange(chunks);

            foreach (var chunk in chunks)
            {
                var mesh = GetMesh(chunk, lod);
                jobState.Meshes.Add(mesh);
            }

            jobState.MeshDataArray = Mesh.AllocateWritableMeshData(jobState.Meshes);

            return jobState;
        }

        protected abstract Mesh GetMesh(TerrainChunk chunk, int lod);

        protected abstract JobHandle ScheduleJob(ChunkGenerationJobState jobState, NativeArray<int2> chunkCoords);

        public async UniTask Wait()
        {
            while (!_dependency.IsCompleted) await UniTask.NextFrame();

            _dependency.Complete();
            _dependency = default;
        }

        public void EndGeneratingChunks()
        {
            foreach (var jobState in _pendingJobs) OnCompleted(jobState);
            _pendingJobs.Clear();
        }

        private void OnCompleted(ChunkGenerationJobState jobState)
        {
            Mesh.ApplyAndDisposeWritableMeshData(jobState.MeshDataArray, jobState.Meshes);

            foreach (var mesh in jobState.Meshes)
            {
                mesh.RecalculateBounds();
                mesh.UploadMeshData(false);
            }

            OnCompleted(jobState.Lod, jobState.Chunks);

            jobState.Clear();
        }

        protected abstract void OnCompleted(int lod, List<TerrainChunk> chunks);
    }
}