using Unity.Burst;
using Unity.Collections;
using Unity.Entities;
using Unity.Mathematics;
using Unity.NetCode;
using Unity.Transforms;

namespace ProjectM.Simulation
{
    /// <summary>
    /// LANTERN 4-socket predicted "fire". On the single fully-predicting pass of each tick, for every player
    /// it loops the 4 ability sockets: a socket whose Spark is set (AbilitySocket.SparkId != 0), whose own
    /// fire event (PlayerInput.Socket0..3) is set this tick, and whose per-socket SocketCooldown has elapsed
    /// fires that Spark. Runs in both worlds: the owning client predict-spawns a Projectile ghost (classified
    /// into the authoritative ghost by ProjectileClassificationSystem via Projectile.SpawnId), and the server
    /// spawns the replicated truth.
    ///
    /// The single AbilityRef/AbilityCooldown/EffectiveAbilityStats model was replaced by the socket kit:
    /// AbilitySocket (loadout), SocketCooldown (hot per-socket cooldown), EffectiveSocketStats (per-socket
    /// folded stats from StatRecomputeSystem). To stay under the 7-type SystemAPI.Query cap, the query holds
    /// only PlayerInput/PlayerFacing/LocalTransform/GhostOwner and reads the socket data by entity via
    /// BufferLookup/ComponentLookup (mirroring the BoonEffects lookup).
    ///
    /// SpawnId key (owner14 | socket2 | fireCount12 | fork4) reserves socket bits so two sockets firing the
    /// same-prefab projectile on one tick classify to DISTINCT ghosts (the review's NP-1/RS-1/DB-3 fix). The
    /// per-socket fire count comes from the replicated command buffer at ServerTick (not a local counter) so
    /// client and server agree. FORK fans extra predicted projectiles, each with a unique fork index in the
    /// low bits; PIERCE/CHAIN/PULL seed the server-only ProjectileEffectState.
    ///
    /// Determinism/idempotency: gated behind IsFirstTimeFullyPredictingTick so a rollback re-sim never
    /// double-spawns. No wall-clock, no System.Random. Predict-spawn is reserved for the Projectile archetype;
    /// Cone applies its effect server-only (cooldown predicted both worlds). Aoe/Hitscan/movement archetypes
    /// are handled elsewhere (steps 3/4) and fall through here. Auto-target stays server-only + gamepad-only.
    /// </summary>
    [UpdateInGroup(typeof(PredictedSimulationSystemGroup))]
    [UpdateAfter(typeof(PlayerAimSystem))]
    [BurstCompile]
    public partial struct AbilityFireSystem : ISystem
    {
        // Server-only knockback stamp for the Cone (guarded + boss-immune).
        ComponentLookup<KnockbackState> m_KnockbackLookup;
        ComponentLookup<BossState> m_BossLookup;
        // Owner-replicated mechanic-changer boons, read by the player entity.
        ComponentLookup<BoonEffects> m_BoonEffectsLookup;
        // LANTERN socket kit, read by the player entity so the fire query stays at 4 type args (7-arg cap).
        BufferLookup<AbilitySocket> m_SocketLookup;
        ComponentLookup<SocketCooldown> m_SocketCdLookup;
        BufferLookup<EffectiveSocketStats> m_EffSocketLookup;

        /// <summary>~9 degree gap between adjacent Split-Shot projectiles (tunable).</summary>
        const float k_ForkSpreadRad = 0.157f;

        [BurstCompile]
        public void OnCreate(ref SystemState state)
        {
            state.RequireForUpdate<AbilityDatabase>();
            state.RequireForUpdate<NetworkTime>();
            m_KnockbackLookup = state.GetComponentLookup<KnockbackState>(isReadOnly: false);
            m_BossLookup = state.GetComponentLookup<BossState>(isReadOnly: true);
            m_BoonEffectsLookup = state.GetComponentLookup<BoonEffects>(isReadOnly: true);
            m_SocketLookup = state.GetBufferLookup<AbilitySocket>(isReadOnly: true);
            m_SocketCdLookup = state.GetComponentLookup<SocketCooldown>(isReadOnly: false);
            m_EffSocketLookup = state.GetBufferLookup<EffectiveSocketStats>(isReadOnly: true);
        }

        [BurstCompile]
        public void OnUpdate(ref SystemState state)
        {
            var networkTime = SystemAPI.GetSingleton<NetworkTime>();
            if (!networkTime.IsFirstTimeFullyPredictingTick)
                return;
            var serverTick = networkTime.ServerTick;
            if (!serverTick.IsValid)
                return;

            var dbEntity = SystemAPI.GetSingletonEntity<AbilityDatabase>();
            var abilityPrefabs = SystemAPI.GetBuffer<AbilityPrefabElement>(dbEntity);
            var abilityDb = SystemAPI.GetSingleton<AbilityDatabase>();
            ref var adb = ref abilityDb.Value.Value;

            bool isServer = state.WorldUnmanaged.IsServer();
            m_KnockbackLookup.Update(ref state);
            m_BossLookup.Update(ref state);
            m_BoonEffectsLookup.Update(ref state);
            m_SocketLookup.Update(ref state);
            m_SocketCdLookup.Update(ref state);
            m_EffSocketLookup.Update(ref state);

            // Server-only LIVING-enemy target set (auto-target assist + Cone cleave), collected once.
            var candidatePositions = new NativeList<float3>(Allocator.Temp);
            var coneTargets = new NativeList<Entity>(Allocator.Temp);
            var coneTargetPos = new NativeList<float3>(Allocator.Temp);
            if (isServer)
            {
                foreach (var (tx, th, te) in
                         SystemAPI.Query<RefRO<LocalTransform>, RefRO<Health>>().WithAll<EnemyTag>().WithEntityAccess())
                {
                    if (th.ValueRO.Current <= 0f) continue; // corpses are neither aim magnets nor cleave targets
                    candidatePositions.Add(tx.ValueRO.Position);
                    coneTargets.Add(te); coneTargetPos.Add(tx.ValueRO.Position);
                }
            }
            var candidates = candidatePositions.AsArray();

            var ecb = new EntityCommandBuffer(state.WorldUpdateAllocator);

            foreach (var (input, facing, xform, owner, entity) in
                     SystemAPI.Query<RefRO<PlayerInput>, RefRO<PlayerFacing>, RefRO<LocalTransform>, RefRO<GhostOwner>>()
                         .WithAll<Simulate>().WithDisabled<Dead>()
                         .WithEntityAccess())
            {
                if (!m_SocketLookup.HasBuffer(entity) || !m_SocketCdLookup.HasComponent(entity) || !m_EffSocketLookup.HasBuffer(entity))
                    continue;
                var sockets = m_SocketLookup[entity];
                var effSockets = m_EffSocketLookup[entity];
                var cd = m_SocketCdLookup[entity]; // struct copy; written back after mutation

                BoonEffects bfx = m_BoonEffectsLookup.HasComponent(entity) ? m_BoonEffectsLookup[entity] : default;
                bool pull = (bfx.Flags & BoonFlag.KnockToPull) != 0;

                // Replicated command buffer at this tick (per-socket fire count for the SpawnId + scheme for aim assist).
                var inputBuffer = SystemAPI.GetBuffer<InputBufferData<PlayerInput>>(entity);
                bool haveApplied = inputBuffer.GetDataAtTick(serverTick, out var applied);

                bool cdDirty = false;
                int socketCount = math.min(SocketId.Count, sockets.Length);
                for (int sk = 0; sk < socketCount; sk++)
                {
                    byte sparkId = sockets[sk].SparkId;
                    if (sparkId == 0) continue;                       // empty socket
                    if (!input.ValueRO.GetSocket(sk).IsSet) continue; // this socket not fired this tick

                    // Per-socket cooldown gate (0 = ready).
                    uint nextFireRaw = cd.Get(sk);
                    if (nextFireRaw != 0)
                    {
                        var nextTick = new NetworkTick(nextFireRaw);
                        if (nextTick.IsValid && nextTick.IsNewerThan(serverTick)) continue;
                    }

                    EffectiveSocketStats es = sk < effSockets.Length ? effSockets[sk] : default;
                    byte archetype = adb.TryGetAbility(sparkId, out var adef) ? adef.Archetype : (byte)AbilityArchetype.Projectile;

                    // CONE: no projectile ghost. Predict the cooldown on both worlds; apply server-only cleave.
                    if (archetype == (byte)AbilityArchetype.Cone)
                    {
                        if (isServer)
                        {
                            float2 cFace = facing.ValueRO.Direction;
                            cFace = math.lengthsq(cFace) < 1e-6f ? new float2(0f, 1f) : math.normalize(cFace);
                            float cRange = math.max(0.1f, es.Range);
                            float cCosHalf = math.cos(math.clamp(es.AutoTargetConeRadians, 0.01f, 3.14159f));
                            uint cStamp = TickUtil.NonZero(serverTick.TickIndexForValidTick);
                            for (int ci = 0; ci < coneTargets.Length; ci++)
                            {
                                if (!MeleeConeMath.InCone(xform.ValueRO.Position, cFace, cRange, cCosHalf, coneTargetPos[ci]))
                                    continue;
                                ecb.AppendToBuffer(coneTargets[ci], new DamageEvent
                                {
                                    Amount = es.Damage,
                                    SourceNetworkId = owner.ValueRO.NetworkId,
                                    SourceTick = cStamp,
                                });
                                KnockbackUtil.Stamp(ref m_KnockbackLookup, m_BossLookup, coneTargets[ci],
                                    xform.ValueRO.Position, coneTargetPos[ci], cFace, Tuning.KnockbackSpeed,
                                    TickUtil.NonZero(serverTick.TickIndexForValidTick + (uint)math.max(1, Tuning.KnockbackDurationTicks)), pull);
                            }
                        }
                        cd.Set(sk, TickUtil.NonZero(serverTick.TickIndexForValidTick + (uint)math.max(1, es.CooldownTicks)));
                        cdDirty = true;
                        continue;
                    }
                    // Aoe/Hitscan (steps 3) and movement/Blink (step 4) are not dispatched here yet.
                    if (archetype != (byte)AbilityArchetype.Projectile)
                        continue;

                    // Resolve the projectile ghost prefab for this Spark id.
                    Entity prefab = Entity.Null;
                    for (int i = 0; i < abilityPrefabs.Length; i++)
                    {
                        if (abilityPrefabs[i].Id == sparkId) { prefab = abilityPrefabs[i].Prefab; break; }
                    }
                    if (prefab == Entity.Null) continue;
                    if (!haveApplied) continue;

                    uint socketFireCount = applied.InternalInput.GetSocket(sk).Count;

                    float2 rawAim = facing.ValueRO.Direction;
                    rawAim = math.lengthsq(rawAim) < 1e-6f ? new float2(0f, 1f) : math.normalize(rawAim);

                    // Client fires along raw aim; only the server applies the gamepad auto-target assist.
                    float2 dir = rawAim;
                    if (isServer && es.AutoTargetRange > 0f && applied.InternalInput.Scheme == InputSchemeId.Gamepad)
                    {
                        dir = AutoTarget.Resolve(xform.ValueRO.Position, rawAim, es.AutoTargetRange, es.AutoTargetConeRadians, candidates);
                    }

                    byte pierce = bfx.Pierce;
                    byte chain = bfx.Chain;
                    byte projFlags = pull ? ProjectileEffectFlag.Pull : (byte)0;
                    int shots = 1 + math.min((int)bfx.Fork, 8);

                    for (int s = 0; s < shots; s++)
                    {
                        float offset = (s - (shots - 1) * 0.5f) * k_ForkSpreadRad;
                        math.sincos(offset, out float sa, out float ca);
                        float2 sdir = math.normalize(new float2(dir.x * ca - dir.y * sa, dir.x * sa + dir.y * ca));

                        // SpawnId: owner(14) | socket(2) | fireCount(12) | fork(4) -- see ProjectileSpawnId.
                        uint spawnId = ProjectileSpawnId.Pack(owner.ValueRO.NetworkId, sk, socketFireCount, s);

                        var projectile = ecb.Instantiate(prefab);
                        float3 planarDir = new float3(sdir.x, 0f, sdir.y);
                        float3 spawnPos = xform.ValueRO.Position + planarDir * 0.6f;
                        spawnPos.y = xform.ValueRO.Position.y;
                        quaternion rot = quaternion.LookRotationSafe(planarDir, math.up());

                        ecb.SetComponent(projectile, LocalTransform.FromPositionRotation(spawnPos, rot));
                        ecb.SetComponent(projectile, new GhostOwner { NetworkId = owner.ValueRO.NetworkId });
                        ecb.SetComponent(projectile, new Projectile
                        {
                            Direction = sdir,
                            SpawnId = spawnId,
                            Speed = es.ProjectileSpeed,
                            Damage = es.Damage,
                            Range = es.Range,
                            DistanceTravelled = 0f,
                        });
                        ecb.SetComponent(projectile, new ProjectileEffectState
                        {
                            PierceRemaining = pierce,
                            ChainRemaining = chain,
                            Flags = projFlags,
                        });
                    }

                    cd.Set(sk, TickUtil.NonZero(serverTick.TickIndexForValidTick + (uint)math.max(1, es.CooldownTicks)));
                    cdDirty = true;
                }

                if (cdDirty) m_SocketCdLookup[entity] = cd;
            }

            ecb.Playback(state.EntityManager);
            ecb.Dispose();
            candidatePositions.Dispose();
            coneTargets.Dispose();
            coneTargetPos.Dispose();
        }
    }
}