using Unity.Burst;
using Unity.Collections;
using Unity.Jobs;
using Unity.Mathematics;

namespace SepCore.Simulation
{
    internal readonly struct EnemySeparationSettings
    {
        public EnemySeparationSettings(float cellSize, float maxRadius, float pushDamping, float maxStepScale,
            bool useTangentialInAttackRange, float pushSmoothing)
        {
            CellSize = cellSize;
            MaxRadius = maxRadius;
            PushDamping = pushDamping;
            MaxStepScale = maxStepScale;
            UseTangentialInAttackRange = useTangentialInAttackRange;
            PushSmoothing = pushSmoothing;
        }

        public float CellSize { get; }
        public float MaxRadius { get; }
        public float PushDamping { get; }
        public float MaxStepScale { get; }
        public bool UseTangentialInAttackRange { get; }
        public float PushSmoothing { get; }
    }

    internal static class EnemySeparationSystem
    {
        // 读写集：R EnemyTransform / EnemyMoveConfig / EnemyBehavior / EnemySeparation(PreviousPush)
        //         W EnemyTransform（经分离 scratch 由 Commit 写回）/ EnemySeparation(PreviousPush)
        public static JobHandle Schedule(EnemyStore enemies, NativeParallelMultiHashMap<long, int> buckets,
            NativeList<EnemyTransformSim> separatedTransforms, in EnemySeparationSettings settings,
            float3 playerPosition, JobHandle dependency)
        {
            int enemyCount = enemies.Count;
            if (enemyCount <= 0)
            {
                return dependency;
            }

            NativeArray<EnemyTransformSim> transformArray = enemies.Transform.Data.AsArray();
            NativeArray<EnemyMoveConfigSim> configArray = enemies.MoveConfig.Data.AsArray();
            NativeArray<EnemyBehaviorSim> behaviorArray = enemies.Behavior.Data.AsArray();
            NativeArray<EnemyTransformSim> separatedTransformArray = separatedTransforms.AsArray();
            NativeArray<EnemySeparationSim> separationArray = enemies.Separation.Data.AsArray();

            BuildEnemySeparationBucketsJob buildJob = new BuildEnemySeparationBucketsJob
            {
                Transforms = transformArray,
                MoveConfigs = configArray,
                Buckets = buckets.AsParallelWriter(),
                CellSize = settings.CellSize
            };
            JobHandle buildHandle = buildJob.Schedule(enemyCount, 64, dependency);
            EnemySeparationJob separationJob = new EnemySeparationJob
            {
                Transforms = transformArray,
                MoveConfigs = configArray,
                Behaviors = behaviorArray,
                Buckets = buckets,
                SeparatedTransforms = separatedTransformArray,
                Separations = separationArray,
                CellSize = settings.CellSize,
                MaxRadius = settings.MaxRadius,
                PlayerPosition = playerPosition,
                PushDamping = settings.PushDamping,
                MaxStepScale = settings.MaxStepScale,
                UseTangentialInAttackRange = settings.UseTangentialInAttackRange,
                PushSmoothing = settings.PushSmoothing
            };
            return separationJob.Schedule(enemyCount, 64, buildHandle);
        }

        public static void Commit(EnemyStore enemies, NativeList<EnemyTransformSim> separatedTransforms,
            int enemyCount)
        {
            for (int i = 0; i < enemyCount; i++)
            {
                enemies.Transform[i] = separatedTransforms[i];
            }
        }
    }

    [BurstCompile]
    internal struct BuildEnemySeparationBucketsJob : IJobParallelFor
    {
        [ReadOnly] public NativeArray<EnemyTransformSim> Transforms;
        [ReadOnly] public NativeArray<EnemyMoveConfigSim> MoveConfigs;
        public NativeParallelMultiHashMap<long, int>.ParallelWriter Buckets;
        public float CellSize;

        public void Execute(int index)
        {
            if (!MoveConfigs[index].AvoidOverlap)
            {
                return;
            }

            float3 position = Transforms[index].Position;
            position.y = 0f;
            int cellX = SimulationCell.ToCell(position.x, CellSize);
            int cellZ = SimulationCell.ToCell(position.z, CellSize);
            Buckets.Add(SimulationCell.Key(cellX, cellZ), index);
        }
    }

    [BurstCompile]
    internal struct EnemySeparationJob : IJobParallelFor
    {
        [ReadOnly] public NativeArray<EnemyTransformSim> Transforms;
        [ReadOnly] public NativeArray<EnemyMoveConfigSim> MoveConfigs;
        [ReadOnly] public NativeArray<EnemyBehaviorSim> Behaviors;
        [ReadOnly] public NativeParallelMultiHashMap<long, int> Buckets;
        public NativeArray<EnemyTransformSim> SeparatedTransforms;
        public NativeArray<EnemySeparationSim> Separations;
        public float CellSize;
        public float MaxRadius;
        public float3 PlayerPosition;
        public float PushDamping;
        public float MaxStepScale;
        public bool UseTangentialInAttackRange;
        public float PushSmoothing;

        public void Execute(int index)
        {
            EnemyMoveConfigSim selfConfig = MoveConfigs[index];
            if (!selfConfig.AvoidOverlap)
            {
                SeparatedTransforms[index] = Transforms[index];
                return;
            }

            EnemyTransformSim self = Transforms[index];
            float3 candidate = self.Position;
            candidate.y = 0f;
            float3 original = candidate;
            float3 fallback =
                math.normalizesafe(new float3(self.Forward.x, 0f, self.Forward.z), new float3(1f, 0f, 0f));
            float selfRadius = selfConfig.BodyRadius > 0f ? selfConfig.BodyRadius : 0.45f;
            int iterations = selfConfig.SeparationIterations > 0 ? selfConfig.SeparationIterations : 1;
            int queryRange = math.max(1, (int)math.ceil((selfRadius + MaxRadius) / CellSize));

            for (int iter = 0; iter < iterations; iter++)
            {
                int cellX = SimulationCell.ToCell(candidate.x, CellSize);
                int cellZ = SimulationCell.ToCell(candidate.z, CellSize);
                float3 pushAccumulation = float3.zero;

                for (int dx = -queryRange; dx <= queryRange; dx++)
                {
                    for (int dz = -queryRange; dz <= queryRange; dz++)
                    {
                        long key = SimulationCell.Key(cellX + dx, cellZ + dz);
                        if (!Buckets.TryGetFirstValue(key, out int otherIndex,
                                out NativeParallelMultiHashMapIterator<long> iterator))
                        {
                            continue;
                        }

                        do
                        {
                            if (otherIndex == index)
                            {
                                continue;
                            }

                            EnemyMoveConfigSim otherConfig = MoveConfigs[otherIndex];
                            if (!otherConfig.AvoidOverlap)
                            {
                                continue;
                            }

                            float otherRadius = otherConfig.BodyRadius > 0f ? otherConfig.BodyRadius : 0.45f;
                            float minDistance = selfRadius + otherRadius;
                            float minDistanceSqr = minDistance * minDistance;

                            float3 otherPosition = Transforms[otherIndex].Position;
                            otherPosition.y = 0f;
                            float3 toSelf = candidate - otherPosition;
                            float sqrDistance = math.lengthsq(toSelf);

                            if (sqrDistance <= float.Epsilon)
                            {
                                float3 zeroDistanceAxis = GetZeroDistanceSeparationAxis(index, otherIndex);
                                float directionSign = index < otherIndex ? 1f : -1f;
                                pushAccumulation += zeroDistanceAxis * (selfRadius * 0.25f * directionSign);
                                continue;
                            }

                            if (sqrDistance >= minDistanceSqr)
                            {
                                continue;
                            }

                            float distance = math.sqrt(sqrDistance);
                            float penetration = minDistance - distance;
                            pushAccumulation += (toSelf / distance) * penetration;
                        } while (Buckets.TryGetNextValue(out otherIndex, ref iterator));
                    }
                }

                if (math.lengthsq(pushAccumulation) <= float.Epsilon)
                {
                    continue;
                }

                float3 resolvedPush = pushAccumulation * PushDamping;

                float maxStep = selfRadius * MaxStepScale;
                float pushLength = math.length(resolvedPush);
                if (pushLength > maxStep && pushLength > float.Epsilon)
                {
                    resolvedPush = resolvedPush / pushLength * maxStep;
                }

                candidate += resolvedPush;
            }

            float3 framePush = candidate - original;
            float2 previousPush2 = Separations[index].PreviousPush;
            float3 previousPush = new float3(previousPush2.x, 0f, previousPush2.y);
            float3 smoothedPush = SmoothSeparationPush(framePush, previousPush, PushSmoothing);

            if (UseTangentialInAttackRange && Behaviors[index].State == EnemyBehaviorSim.StateInAttackRange)
            {
                smoothedPush = ProjectToTangential(smoothedPush, PlayerPosition, original);
            }

            float maxTotalStep = selfRadius * MaxStepScale * iterations;
            float smoothedLength = math.length(smoothedPush);
            if (smoothedLength > maxTotalStep && smoothedLength > float.Epsilon)
            {
                smoothedPush = smoothedPush / smoothedLength * maxTotalStep;
            }

            float3 finalPosition = original + smoothedPush;
            Separations[index] = new EnemySeparationSim { PreviousPush = new float2(smoothedPush.x, smoothedPush.z) };
            self.Position = new float3(finalPosition.x, self.Position.y, finalPosition.z);
            if (math.lengthsq(smoothedPush) > float.Epsilon)
            {
                self.Forward = new float3(fallback.x, self.Forward.y, fallback.z);
            }

            SeparatedTransforms[index] = self;
        }

        private static float3 SmoothSeparationPush(float3 framePush, float3 previousPush, float pushSmoothing)
        {
            float frameLengthSqr = math.lengthsq(framePush);
            float previousLengthSqr = math.lengthsq(previousPush);

            if (frameLengthSqr <= float.Epsilon)
            {
                return float3.zero;
            }

            if (previousLengthSqr <= float.Epsilon || pushSmoothing <= 0f)
            {
                return framePush;
            }

            float frameLength = math.sqrt(frameLengthSqr);
            float previousLength = math.sqrt(previousLengthSqr);
            float3 frameDirection = framePush / frameLength;
            float3 previousDirection = previousPush / previousLength;
            float directionAlignment = math.dot(frameDirection, previousDirection);

            if (directionAlignment >= 0.35f)
            {
                return framePush;
            }

            float directionalFactor = math.saturate((0.35f - directionAlignment) / 1.35f);
            float smoothingStrength = pushSmoothing * directionalFactor;
            return math.lerp(framePush, previousPush, smoothingStrength);
        }

        private static float3 ProjectToTangential(float3 push, float3 playerPosition, float3 currentPosition)
        {
            if (math.lengthsq(push) <= float.Epsilon)
            {
                return push;
            }

            float3 toPlayer = playerPosition - currentPosition;
            float toPlayerSqr = math.lengthsq(toPlayer);
            if (toPlayerSqr <= float.Epsilon)
            {
                return push;
            }

            float3 radialDirection = toPlayer / math.sqrt(toPlayerSqr);
            float radialOffset = math.dot(push, radialDirection);
            return push - radialDirection * radialOffset;
        }

        private static float3 GetZeroDistanceSeparationAxis(int index, int otherIndex)
        {
            int lowIndex = math.min(index, otherIndex);
            int highIndex = math.max(index, otherIndex);
            uint pairHash = (uint)(lowIndex * 73856093) ^ (uint)(highIndex * 19349663);

            float axisX = (pairHash & 1023u) / 511.5f - 1f;
            float axisZ = ((pairHash >> 10) & 1023u) / 511.5f - 1f;
            float3 axis = new float3(axisX, 0f, axisZ);
            return math.normalizesafe(axis, new float3(1f, 0f, 0f));
        }
    }
}