using System.Linq;
using System.Collections.Generic;
using Unity.Collections;
using UnityEngine.SceneManagement;
using UnityEditor;
using System.Runtime.InteropServices;
using Brick = UnityEngine.Rendering.ProbeBrickIndex.Brick;

namespace UnityEngine.Rendering
{
    class ProbeVolumeProfileInfo
    {
        public int simplificationLevels;
        public float minDistanceBetweenProbes;
        public Vector3 probeOffset;

        public int maxSubdivision => ProbeVolumeBakingSet.GetMaxSubdivision(simplificationLevels);
        public float minBrickSize => ProbeVolumeBakingSet.GetMinBrickSize(minDistanceBetweenProbes);
        public int cellSizeInBricks => ProbeVolumeBakingSet.GetCellSizeInBricks(simplificationLevels);
        public float cellSizeInMeters => (float)cellSizeInBricks * minBrickSize;

        public Vector3Int PositionToCell(Vector3 position) => Vector3Int.FloorToInt((position - probeOffset) / cellSizeInMeters);
    }

    public partial class AdaptiveProbeVolumes
    {
        static internal ProbeVolumeProfileInfo m_ProfileInfo = null;

        static void FindWorldBounds()
        {
            var prv = ProbeReferenceVolume.instance;
            prv.clearAssetsOnVolumeClear = true;

            var activeScene = SceneManager.GetActiveScene();
            var activeSet = ProbeVolumeBakingSet.GetBakingSetForScene(activeScene);

            bool hasFoundBounds = false;

            foreach (var sceneGUID in activeSet.sceneGUIDs)
            {
                var bakeData = activeSet.GetSceneBakeData(sceneGUID);
                if (bakeData.hasProbeVolume)
                {
                    if (hasFoundBounds)
                    {
                        globalBounds.Encapsulate(bakeData.bounds);
                    }
                    else
                    {
                        globalBounds = bakeData.bounds;
                        hasFoundBounds = true;
                    }
                }
            }

            ProbeReferenceVolume.instance.globalBounds = globalBounds;
        }

        internal static List<ProbeVolumePerSceneData> GetPerSceneDataList()
        {
            var fullPerSceneDataList = ProbeReferenceVolume.instance.perSceneDataList;
            if (!isBakingSceneSubset)
                return fullPerSceneDataList;

            List<ProbeVolumePerSceneData> usedPerSceneDataList = new ();
            foreach (var sceneData in fullPerSceneDataList)
            {
                if (partialBakeSceneList.Contains(ProbeReferenceVolume.GetSceneGUID(sceneData.gameObject.scene)))
                    usedPerSceneDataList.Add(sceneData);
            }
            return usedPerSceneDataList;
        }

        internal static List<ProbeVolume> GetProbeVolumeList()
        {
            var fullPvList = GameObject.FindObjectsByType<ProbeVolume>(FindObjectsSortMode.InstanceID);
            List<ProbeVolume> usedPVList;

            if (isBakingSceneSubset)
            {
                usedPVList = new List<ProbeVolume>();
                foreach (var pv in fullPvList)
                {
                    if (pv.isActiveAndEnabled && partialBakeSceneList.Contains(ProbeReferenceVolume.GetSceneGUID(pv.gameObject.scene)))
                        usedPVList.Add(pv);
                }
            }
            else
            {
                usedPVList = new List<ProbeVolume>(fullPvList);
            }

            return usedPVList;
        }

        static ProbeVolumeProfileInfo GetProfileInfoFromBakingSet(ProbeVolumeBakingSet set)
        {
            var result = new ProbeVolumeProfileInfo();
            result.minDistanceBetweenProbes = set.minDistanceBetweenProbes;
            result.simplificationLevels = set.simplificationLevels;
            result.probeOffset = set.probeOffset;
            return result;
        }

        static int PosToIndex(Vector3Int pos)
        {
            Vector3Int normalizedPos = pos - minCellPosition;
            return normalizedPos.z * (cellCount.x * cellCount.y) + normalizedPos.y * cellCount.x + normalizedPos.x;
        }

        static internal bool CanFreezePlacement()
        {
            if (!ProbeReferenceVolume.instance.supportLightingScenarios)
                return false;

            // Check if all the scene datas in the scene have a baking set, if not then we cannot enable this option.
            var sceneDataList = GetPerSceneDataList();
            if (sceneDataList.Count == 0)
                return false;

            foreach (var sceneData in sceneDataList)
            {
                if (sceneData.serializedBakingSet == null || sceneData.serializedBakingSet.GetSceneCellIndexList(sceneData.sceneGUID) == null)
                    return false;
            }

            return true;
        }

        static NativeList<Vector3> RunPlacement(ProbeVolumeProfileInfo profileInfo, ProbeReferenceVolume refVolume, ref bool canceledByUser)
        {
            Debug.Assert(profileInfo != null);

            // Overwrite loaded settings with data from profile. Note that the m_BakingSet.profile is already patched up if isFreezingPlacement
            float prevBrickSize = refVolume.MinBrickSize();
            int prevMaxSubdiv = refVolume.GetMaxSubdivision();
            Vector3 prevOffset = refVolume.ProbeOffset();
            refVolume.SetSubdivisionDimensions(profileInfo.minBrickSize, profileInfo.maxSubdivision, profileInfo.probeOffset);

            // All probes need to be baked only once for the whole batch and not once per cell
            // The reason is that the baker is not deterministic so the same probe position baked in two different cells may have different values causing seams artefacts.
            m_BakingBatch = new BakingBatch(cellCount, refVolume);

            // Run subdivision
            ProbeSubdivisionResult result;
            using (new BakingSetupProfiling(BakingSetupProfiling.Stages.BakeBricks))
                result = GetWorldSubdivision(ref canceledByUser);

            if (canceledByUser)
                return new NativeList<Vector3>(Allocator.Temp);

            // Compute probe positions
            NativeList<Vector3> positions;
            using (new BakingSetupProfiling(BakingSetupProfiling.Stages.ApplySubdivisionResults))
                positions = ApplySubdivisionResults(result, ref canceledByUser);

            // Restore loaded asset settings
            refVolume.SetSubdivisionDimensions(prevBrickSize, prevMaxSubdiv, prevOffset);

            return positions;
        }

        static ProbeSubdivisionResult GetWorldSubdivision(ref bool canceledByUser)
        {
            var perSceneDataList = GetPerSceneDataList();

            if (isFreezingPlacement)
                return GetBricksFromLoaded(perSceneDataList);

            var ctx = PrepareProbeSubdivisionContext(perSceneDataList);
            return BakeBricks(ctx, m_BakingBatch.contributors, showProgress: true, ref canceledByUser);
        }

        static NativeList<Vector3> ApplySubdivisionResults(ProbeSubdivisionResult results, ref bool canceledByUser)
        {
            int cellIdx = 0, freq = 10;
            BakingSetupProfiling.GetProgressRange(out float progress0, out float progress1);

            var positions = new NativeList<Vector3>(Allocator.Persistent);
            foreach ((var position, var bounds, var bricks) in results.cells)
            {
                if (cellIdx++ % freq == 0) // Don't refresh progress bar at every iteration because it's slow
                {
                    if (EditorUtility.DisplayCancelableProgressBar("Baking Probe Volumes", $"Subdividing cell {cellIdx} out of {results.cells.Count}", Mathf.Lerp(progress0, progress1, cellIdx / (float)results.cells.Count)))
                    {
                        canceledByUser = true;
                        return positions;
                    }
                }

                int positionStart = positions.Length;

                ConvertBricksToPositions(bricks, out var probePositions, out var brickSubdivLevels);
                if (!DeduplicateProbePositions(in probePositions, in brickSubdivLevels, m_BakingBatch, positions, out var probeIndices))
                    return new NativeList<Vector3>(Allocator.Persistent);

                BakingCell cell = new BakingCell()
                {
                    index = PosToIndex(position),
                    position = position,
                    bounds = bounds,
                    bricks = bricks,
                    probePositions = probePositions,
                    probeIndices = probeIndices,
                };

                m_BakingBatch.cells.Add(cell);
                m_BakingBatch.cellIndex2SceneReferences[cell.index] = new HashSet<string>(results.scenesPerCells[cell.position]);
            }

            return positions;
        }

        // We know that the current limitation on native containers is this. When an integer overflow bug (https://jira.unity3d.com/browse/UUM-113721) has been fixed, we can raise the limit
        // This and related work is tracked by https://jira.unity3d.com/browse/GFXLIGHT-1738
        static readonly long k_MaxNumberOfPositions = 67180350;

        static bool DeduplicateProbePositions(in Vector3[] probePositions, in int[] brickSubdivLevel, BakingBatch batch,
            NativeList<Vector3> uniquePositions, out int[] indices)
        {
            long numberOfPositions = (long)probePositions.Length + batch.positionToIndex.Count;
            if (numberOfPositions > k_MaxNumberOfPositions)
            {
                Debug.LogError($"The number of Adaptive Probe Volume (APV) probes Unity generated exceeds the current system limit of {k_MaxNumberOfPositions} probes per Baking Set. Reduce density either by adjusting the general Probe Spacing in the Lighting window, or by modifying the Adaptive Probe Volumes in the scene to limit where the denser subdivision levels are used.");
                indices = null;

                return false;
            }

            indices = new int[probePositions.Length];
            int uniqueIndex = batch.positionToIndex.Count;

            for (int i = 0; i < probePositions.Length; i++)
            {
                var pos = probePositions[i];
                var brickSubdiv = brickSubdivLevel[i];
                int probeHash = batch.GetProbePositionHash(pos);

                if (batch.positionToIndex.TryGetValue(probeHash, out var index))
                {
                    indices[i] = index;
                    int oldBrickLevel = batch.uniqueBrickSubdiv[probeHash];
                    if (brickSubdiv < oldBrickLevel)
                        batch.uniqueBrickSubdiv[probeHash] = brickSubdiv;
                }
                else
                {
                    batch.positionToIndex[probeHash] = uniqueIndex;
                    indices[i] = uniqueIndex;
                    batch.uniqueBrickSubdiv[probeHash] = brickSubdiv;
                    uniquePositions.Add(pos);
                    uniqueIndex++;
                }
            }

            return true;
        }

        static ProbeSubdivisionResult GetBricksFromLoaded(List<ProbeVolumePerSceneData> dataList)
        {
            var result = new ProbeSubdivisionResult();

            // We read bricks from the asset rather than using the currently loaded cells.
            // This is because not all bricks from the previous bake are guaranteed to be currently loaded,
            // we may for example have hit the max brick count given the selected memory budget.
            ProbeVolumeStreamableAsset bricksDataAsset = m_BakingSet.cellBricksDataAsset;
            bricksDataAsset.EnsureAssetLoaded();
            using NativeArray<Brick> previousBricks = bricksDataAsset.asset.GetData<Brick>();

            foreach (var data in dataList)
            {
                var cellSize = m_ProfileInfo.minDistanceBetweenProbes * 3.0f * m_ProfileInfo.cellSizeInBricks;
                Vector3 cellDimensions = new Vector3(cellSize, cellSize, cellSize);

                // Loop through cells in asset, we need to be careful as there'll be duplicates.
                // As we go through the cells we fill ProbeSubdivisionResult as we go.
                var cells = m_BakingSet.GetSceneCellIndexList(data.sceneGUID);
                foreach (var cellIndex in cells)
                {
                    var cellDesc = m_BakingSet.GetCellDesc(cellIndex);
                    var cellPos = cellDesc.position;

                    if (!result.scenesPerCells.ContainsKey(cellPos))
                    {
                        result.scenesPerCells[cellPos] = new HashSet<string>();

                        var center = new Vector3((cellPos.x + 0.5f) * cellSize, (cellPos.y + 0.5f) * cellSize, (cellPos.z + 0.5f) * cellSize);

                        var cellStreamingDesc = bricksDataAsset.streamableCellDescs[cellIndex];
                        int bricksOffset = cellStreamingDesc.offset / bricksDataAsset.elementSize;
                        int bricksCount = Mathf.Min(cellStreamingDesc.elementCount, cellDesc.bricksCount);
                        Brick[] bricks = previousBricks.GetSubArray(bricksOffset, bricksCount).ToArray();

                        result.cells.Add((cellPos, new Bounds(center, cellDimensions), bricks));
                    }
                    result.scenesPerCells[cellPos].Add(data.sceneGUID);
                }
            }

            return result;
        }

        static internal ProbeSubdivisionContext PrepareProbeSubdivisionContext(List<ProbeVolumePerSceneData> perSceneDataList, bool liveContext = false)
        {
            // Prepare all the information in the scene for baking GI.
            Vector3 refVolOrigin = Vector3.zero; // TODO: This will need to be center of the world bounds.

            if (m_BakingSet == null)
            {
                if (perSceneDataList.Count == 0) return new ProbeSubdivisionContext();
                SetBakingContext(perSceneDataList);
            }

            var profileInfo = m_ProfileInfo;
            if (liveContext || m_ProfileInfo == null)
                profileInfo = GetProfileInfoFromBakingSet(m_BakingSet);

            ProbeSubdivisionContext ctx = new ProbeSubdivisionContext();
            ctx.Initialize(m_BakingSet, profileInfo, refVolOrigin);
            return ctx;
        }

        static internal ProbeSubdivisionResult BakeBricks(ProbeSubdivisionContext ctx, in GIContributors contributors, bool showProgress, ref bool canceledByUser)
        {
            var result = new ProbeSubdivisionResult();

            if (ctx.probeVolumes.Count == 0)
                return result;

            int cellIdx = 0, freq = 100;
            BakingSetupProfiling.GetProgressRange(out float progress0, out float progress1);

            using (var gpuResources = ProbePlacement.AllocateGPUResources(ctx.probeVolumes.Count, ctx.profile))
            {
                // subdivide all the cells and generate brick positions
                foreach (var cell in ctx.cells)
                {
                    if (showProgress && cellIdx++ % freq == 0)  // Don't refresh progress bar at every iteration because it's slow
                    {
                        if (EditorUtility.DisplayCancelableProgressBar("Generating Probe Volume Bricks", $"Processing cell {cellIdx} out of {ctx.cells.Count}", Mathf.Lerp(progress0, progress1, cellIdx / (float)ctx.cells.Count)))
                        {
                            canceledByUser = true;
                            return new ProbeSubdivisionResult();
                        }
                    }

                    var scenesInCell = new HashSet<string>();

                    // Calculate overlaping probe volumes to avoid unnecessary work
                    var overlappingProbeVolumes = new List<(ProbeVolume component, ProbeReferenceVolume.Volume volume, Bounds bounds)>();
                    foreach (var probeVolume in ctx.probeVolumes)
                    {
                        if (ProbeVolumePositioning.OBBAABBIntersect(probeVolume.volume, cell.bounds, probeVolume.bounds))
                        {
                            overlappingProbeVolumes.Add(probeVolume);
                            scenesInCell.Add(ProbeReferenceVolume.GetSceneGUID(probeVolume.component.gameObject.scene));
                        }
                    }

                    // Calculate valid renderers to avoid unnecessary work (a renderer needs to overlap a probe volume and match the layer)
                    var filteredContributors = contributors.Filter(ctx.bakingSet, cell.bounds, overlappingProbeVolumes);

                    if (filteredContributors.Count == 0 && !overlappingProbeVolumes.Any(v => v.component.fillEmptySpaces))
                        continue;

                    var bricks = ProbePlacement.SubdivideCell(cell.position, cell.bounds, ctx, gpuResources, filteredContributors, overlappingProbeVolumes);
                    if (bricks.Length == 0)
                        continue;

                    foreach (var renderer in filteredContributors.renderers)
                        scenesInCell.Add(ProbeReferenceVolume.GetSceneGUID(renderer.component.gameObject.scene));
                    foreach (var terrain in filteredContributors.terrains)
                        scenesInCell.Add(ProbeReferenceVolume.GetSceneGUID(terrain.component.gameObject.scene));

                    result.cells.Add((cell.position, cell.bounds, bricks));
                    result.scenesPerCells[cell.position] = scenesInCell;
                }
            }

            return result;
        }

        static void ModifyProfileFromLoadedData(ProbeVolumeBakingSet bakingSet)
        {
            m_ProfileInfo.simplificationLevels = bakingSet.bakedSimplificationLevels;
            m_ProfileInfo.minDistanceBetweenProbes = bakingSet.bakedMinDistanceBetweenProbes;
            m_ProfileInfo.probeOffset = bakingSet.bakedProbeOffset;
            globalBounds = bakingSet.globalBounds;
        }

        // Converts brick information into positional data at kBrickProbeCountPerDim * kBrickProbeCountPerDim * kBrickProbeCountPerDim resolution
        internal static void ConvertBricksToPositions(Brick[] bricks, out Vector3[] outProbePositions, out int[] outBrickSubdiv)
        {
            int posIdx = 0;
            float scale = ProbeReferenceVolume.instance.MinBrickSize() / ProbeBrickPool.kBrickCellCount;
            Vector3 offset = ProbeReferenceVolume.instance.ProbeOffset();

            outProbePositions = new Vector3[bricks.Length * ProbeBrickPool.kBrickProbeCountTotal];
            outBrickSubdiv = new int[bricks.Length * ProbeBrickPool.kBrickProbeCountTotal];

            foreach (var b in bricks)
            {
                int brickSize = ProbeReferenceVolume.CellSize(b.subdivisionLevel);
                Vector3Int brickOffset = b.position * ProbeBrickPool.kBrickCellCount;

                for (int z = 0; z < ProbeBrickPool.kBrickProbeCountPerDim; z++)
                {
                    for (int y = 0; y < ProbeBrickPool.kBrickProbeCountPerDim; y++)
                    {
                        for (int x = 0; x < ProbeBrickPool.kBrickProbeCountPerDim; x++)
                        {
                            var probeOffset = brickOffset + new Vector3Int(x, y, z) * brickSize;

                            outProbePositions[posIdx] = offset + (Vector3)probeOffset * scale;
                            outBrickSubdiv[posIdx] = b.subdivisionLevel;

                            posIdx++;
                        }
                    }
                }
            }
        }
    }
}