using AnotherReplayReader.ReplayFile;
using System;
using System.Collections.Generic;
using System.Collections.Immutable;
using System.Linq;

namespace AnotherReplayReader.Utils
{
    /// <summary>
    /// Index of replay facts extracted from CommandChunk data.
    /// Used by AIAnalysisValidation to cross-reference LLM claims against
    /// actual replay operations.
    /// </summary>
    internal sealed class ReplayFactIndex
    {
        /// <summary>First time a UnitId was observed in any command.</summary>
        public ImmutableDictionary<uint, TimeSpan> UnitIdFirstObservedTime { get; }

        /// <summary>Special powers used by each UnitId.</summary>
        public ImmutableDictionary<uint, ImmutableHashSet<string>> UnitIdSpecialPowers { get; }

        /// <summary>UnitIds that appeared as builder ("建造者") in construction commands.</summary>
        public ImmutableHashSet<uint> BuilderUnitIds { get; }

        /// <summary>UnitIds that appeared as production structures ("出兵建筑").</summary>
        public ImmutableHashSet<uint> ProducerUnitIds { get; }

        /// <summary>Per player, per unit asset name, first production start time.</summary>
        public ImmutableDictionary<int, ImmutableDictionary<string, TimeSpan>> PlayerFirstProductionTime { get; }

        /// <summary>Per player, which UnitIds they have selected.</summary>
        public ImmutableDictionary<int, ImmutableHashSet<uint>> PlayerSelectedUnitIds { get; }

        /// <summary>Per player, UnitIds with strong ownership evidence (control group, builder/producer, repair, sell, power caster).</summary>
        public ImmutableDictionary<int, ImmutableHashSet<uint>> PlayerStrongOwnershipUnitIds { get; }

        /// <summary>Per player, UnitIds with weak ownership evidence (plain selection only).</summary>
        public ImmutableDictionary<int, ImmutableHashSet<uint>> PlayerWeakOwnershipUnitIds { get; }

        /// <summary>Per player, tech/protocol choices (0x24E).</summary>
        public ImmutableDictionary<int, ImmutableHashSet<string>> PlayerTechChoices { get; }

        public ReplayFactIndex(
            ImmutableDictionary<uint, TimeSpan> unitIdFirstObservedTime,
            ImmutableDictionary<uint, ImmutableHashSet<string>> unitIdSpecialPowers,
            ImmutableHashSet<uint> builderUnitIds,
            ImmutableHashSet<uint> producerUnitIds,
            ImmutableDictionary<int, ImmutableDictionary<string, TimeSpan>> playerFirstProductionTime,
            ImmutableDictionary<int, ImmutableHashSet<uint>> playerSelectedUnitIds,
            ImmutableDictionary<int, ImmutableHashSet<uint>> playerStrongOwnershipUnitIds,
            ImmutableDictionary<int, ImmutableHashSet<uint>> playerWeakOwnershipUnitIds,
            ImmutableDictionary<int, ImmutableHashSet<string>> playerTechChoices)
        {
            UnitIdFirstObservedTime = unitIdFirstObservedTime;
            UnitIdSpecialPowers = unitIdSpecialPowers;
            BuilderUnitIds = builderUnitIds;
            ProducerUnitIds = producerUnitIds;
            PlayerFirstProductionTime = playerFirstProductionTime;
            PlayerSelectedUnitIds = playerSelectedUnitIds;
            PlayerStrongOwnershipUnitIds = playerStrongOwnershipUnitIds;
            PlayerWeakOwnershipUnitIds = playerWeakOwnershipUnitIds;
            PlayerTechChoices = playerTechChoices;
        }

        public static ReplayFactIndex Build(
            ImmutableArray<(TimeSpan Time, ImmutableArray<CommandChunk> Commands)> timeline,
            IReadOnlyDictionary<uint, string> stringHashTable)
        {
            var unitFirstObserved = new Dictionary<uint, TimeSpan>();
            var unitSpecialPowers = new Dictionary<uint, HashSet<string>>();
            var builderUnits = new HashSet<uint>();
            var producerUnits = new HashSet<uint>();
            var playerFirstProduction = new Dictionary<int, Dictionary<string, TimeSpan>>();
            var playerSelected = new Dictionary<int, HashSet<uint>>();
            var playerStrongOwnership = new Dictionary<int, HashSet<uint>>();
            var playerWeakOwnership = new Dictionary<int, HashSet<uint>>();
            var playerTechChoices = new Dictionary<int, HashSet<string>>();
            // 编队号在同一局内可能被不同玩家复用，因此键必须包含玩家。
            var controlGroups = new Dictionary<long, HashSet<uint>>();

            foreach (var (time, commands) in timeline)
            {
                foreach (var command in commands)
                {
                    ProcessCommand(time, command, stringHashTable,
                        unitFirstObserved, unitSpecialPowers,
                        builderUnits, producerUnits,
                        playerFirstProduction, playerSelected,
                        playerStrongOwnership, playerWeakOwnership,
                        playerTechChoices, controlGroups);
                }
            }

            return new ReplayFactIndex(
                unitFirstObserved.ToImmutableDictionary(),
                unitSpecialPowers.ToImmutableDictionary(
                    kv => kv.Key, kv => kv.Value.ToImmutableHashSet()),
                builderUnits.ToImmutableHashSet(),
                producerUnits.ToImmutableHashSet(),
                playerFirstProduction.ToImmutableDictionary(
                    kv => kv.Key, kv => kv.Value.ToImmutableDictionary()),
                playerSelected.ToImmutableDictionary(
                    kv => kv.Key, kv => kv.Value.ToImmutableHashSet()),
                playerStrongOwnership.ToImmutableDictionary(
                    kv => kv.Key, kv => kv.Value.ToImmutableHashSet()),
                playerWeakOwnership.ToImmutableDictionary(
                    kv => kv.Key, kv => kv.Value.ToImmutableHashSet()),
                playerTechChoices.ToImmutableDictionary(
                    kv => kv.Key, kv => kv.Value.ToImmutableHashSet()));
        }

        private static void ProcessCommand(
            TimeSpan time,
            CommandChunk command,
            IReadOnlyDictionary<uint, string> stringHashTable,
            Dictionary<uint, TimeSpan> unitFirstObserved,
            Dictionary<uint, HashSet<string>> unitSpecialPowers,
            HashSet<uint> builderUnits,
            HashSet<uint> producerUnits,
            Dictionary<int, Dictionary<string, TimeSpan>> playerFirstProduction,
            Dictionary<int, HashSet<uint>> playerSelected,
            Dictionary<int, HashSet<uint>> playerStrongOwnership,
            Dictionary<int, HashSet<uint>> playerWeakOwnership,
            Dictionary<int, HashSet<string>> playerTechChoices,
            Dictionary<long, HashSet<uint>> controlGroups)
        {
            var player = command.PlayerIndex;
            var cmdId = command.CommandId;

            switch (cmdId)
            {
                // select unit(s): 0x1F5
                case 0x1F5:
                // 选择相同单位（W）：ObjectId 语义与选择相同，作为弱所有权
                case 0x1F6:
                // 选择所有单位（Q）：若能解析出 UnitId，同样作为弱所有权
                case 0x22A:
                    RecordSelectUnit(time, command, player, unitFirstObserved, playerSelected, playerWeakOwnership);
                    break;

                // 从选择中移除单位：与选择类似，仅弱所有权
                case 0x1F9:
                    RecordObjectReferenceWithOwnership(time, command, player, unitFirstObserved, playerWeakOwnership);
                    break;

                // special power (no target): 0x1FE —— 布局确凿，ObjectId 是施法者
                case 0x1FE:
                // special power (target position and angle): 0x200 —— 布局确凿，ObjectId 是施法者
                case 0x200:
                    RecordSpecialPower(time, command, player, stringHashTable,
                        unitFirstObserved, unitSpecialPowers, playerStrongOwnership);
                    break;

                // special power (target position): 0x1FF —— ObjectId 语义待核实，只记录"出现过"
                case 0x1FF:
                // special power (target unit): 0x201 —— ObjectId 可能是目标
                case 0x201:
                // special power (one or more targets): 0x232 —— ObjectId 可能是目标
                case 0x232:
                    RecordObjectReference(time, command, unitFirstObserved);
                    break;

                // start production: 0x205
                case 0x205:
                    RecordProduction(time, command, player, unitFirstObserved,
                        producerUnits, playerFirstProduction, playerStrongOwnership);
                    break;

                // start construction: 0x207
                case 0x207:
                    RecordConstruction(time, command, player, unitFirstObserved, builderUnits, playerStrongOwnership);
                    break;

                // place building: 0x209
                case 0x209:
                    RecordPlaceBuilding(time, command, player, unitFirstObserved, builderUnits, playerStrongOwnership);
                    break;

                // sell building: 0x20A
                case 0x20A:
                    RecordObjectReferenceWithOwnership(time, command, player, unitFirstObserved, playerStrongOwnership);
                    break;

                // 开始/停止维修建筑：只能维修己方建筑 → 强所有权
                case 0x228:
                case 0x229:
                    RecordObjectReferenceWithOwnership(time, command, player, unitFirstObserved, playerStrongOwnership);
                    break;

                // 命令矿车交矿 / 让矿车去采矿：只能命令己方矿车 → 强所有权
                case 0x212:
                case 0x248:
                    RecordObjectReferenceWithOwnership(time, command, player, unitFirstObserved, playerStrongOwnership);
                    break;

                // 创建编队：编队成员几乎确定是己方单位 → 强所有权
                case 0x1FA:
                    RecordControlGroupCreate(time, command, player, unitFirstObserved,
                        playerStrongOwnership, controlGroups);
                    break;

                // 选择编队 / 将编队加入选择：通过编队状态解析成员 → 强所有权
                case 0x1FB:
                case 0x1FC:
                    RecordControlGroupSelect(time, command, player, unitFirstObserved,
                        playerStrongOwnership, controlGroups);
                    break;

                // 选择协议：全局生效，无 UnitId
                case 0x24E:
                    RecordTechChoice(time, command, player, playerTechChoices);
                    break;

                // move: 0x214
                case 0x214:
                // attack move: 0x215
                case 0x215:
                    // These commands operate on currently selected units.
                    // The target is a position, not a UnitId.
                    break;
            }
        }

        private static void RecordSelectUnit(
            TimeSpan time,
            CommandChunk command,
            int player,
            Dictionary<uint, TimeSpan> unitFirstObserved,
            Dictionary<int, HashSet<uint>> playerSelected,
            Dictionary<int, HashSet<uint>> playerWeakOwnership)
        {
            // Data layout for 0x1F5:
            // Data[0]: Bool (isReplace), if count > 0 the rest are ObjectIds
            // Data[1..]: ObjectIds of selected units
            foreach (var entry in command.Data)
            {
                if (entry.Type is CommandArgumentType.ObjectId or CommandArgumentType.ObjectId_2)
                {
                    if (entry.Count == 1)
                    {
                        var unitId = (uint)entry.Value;
                        TryRecordFirstObserved(unitId, time, unitFirstObserved);
                        RecordPlayerSelection(player, unitId, playerSelected);
                        RecordPlayerOwnership(player, unitId, playerWeakOwnership);
                    }
                    else
                    {
                        foreach (var id in (uint[])entry.Value)
                        {
                            TryRecordFirstObserved(id, time, unitFirstObserved);
                            RecordPlayerSelection(player, id, playerSelected);
                            RecordPlayerOwnership(player, id, playerWeakOwnership);
                        }
                    }
                }
            }
        }

        private static void RecordSpecialPower(
            TimeSpan time,
            CommandChunk command,
            int player,
            IReadOnlyDictionary<uint, string> stringHashTable,
            Dictionary<uint, TimeSpan> unitFirstObserved,
            Dictionary<uint, HashSet<string>> unitSpecialPowers,
            Dictionary<int, HashSet<uint>> playerStrongOwnership)
        {
            string? powerName = null;
            var unitIds = new List<uint>();

            foreach (var entry in command.Data)
            {
                switch (entry.Type)
                {
                    case CommandArgumentType.Int32 when powerName is null:
                    {
                        // First Int32 is the special power hash ID；
                        // 同类型参数可能被打包成数组（首个元素是 hash）
                        var hash = entry.Count == 1 && entry.Value is int singleInt
                            ? unchecked((uint)singleInt)
                            : entry.Value is int[] ints && ints.Length > 0
                                ? unchecked((uint)ints[0])
                                : (uint?)null;
                        if (hash is { } hashValue)
                        {
                            powerName = stringHashTable.TryGetValue(hashValue, out var name)
                                ? name
                                : $"Hash_{hashValue:X8}";
                        }
                        break;
                    }
                    case CommandArgumentType.ObjectId or CommandArgumentType.ObjectId_2:
                    {
                        if (entry.Count == 1)
                        {
                            var id = (uint)entry.Value;
                            if (id != 0) unitIds.Add(id);
                        }
                        else
                        {
                            foreach (var id in (uint[])entry.Value)
                            {
                                if (id != 0) unitIds.Add(id);
                            }
                        }
                        break;
                    }
                }
            }

            if (powerName is null || unitIds.Count == 0)
            {
                return;
            }

            foreach (var unitId in unitIds)
            {
                TryRecordFirstObserved(unitId, time, unitFirstObserved);
                if (!unitSpecialPowers.TryGetValue(unitId, out var powers))
                {
                    powers = new HashSet<string>();
                    unitSpecialPowers[unitId] = powers;
                }
                powers.Add(powerName);
                RecordPlayerOwnership(player, unitId, playerStrongOwnership);
            }
        }

        private static void RecordProduction(
            TimeSpan time,
            CommandChunk command,
            int player,
            Dictionary<uint, TimeSpan> unitFirstObserved,
            HashSet<uint> producerUnits,
            Dictionary<int, Dictionary<string, TimeSpan>> playerFirstProduction,
            Dictionary<int, HashSet<uint>> playerStrongOwnership)
        {
            uint? producerId = null;
            string? unitName = null;

            foreach (var entry in command.Data)
            {
                switch (entry.Type)
                {
                    case CommandArgumentType.ObjectId or CommandArgumentType.ObjectId_2
                        when producerId is null:
                        if (entry.Count == 1)
                        {
                            var id = (uint)entry.Value;
                            if (id != 0) producerId = id;
                        }
                        else if (entry.Value is uint[] ids)
                        {
                            foreach (var id in ids)
                            {
                                if (id != 0)
                                {
                                    producerId = id;
                                    break;
                                }
                            }
                        }
                        break;
                    case CommandArgumentType.AsciiString or CommandArgumentType.UnicodeString
                        when unitName is null:
                        unitName = entry.Value.ToString() ?? string.Empty;
                        break;
                }
            }

            if (producerId.HasValue)
            {
                TryRecordFirstObserved(producerId.Value, time, unitFirstObserved);
                producerUnits.Add(producerId.Value);
                RecordPlayerOwnership(player, producerId.Value, playerStrongOwnership);
            }

            if (!string.IsNullOrWhiteSpace(unitName))
            {
                if (!playerFirstProduction.TryGetValue(player, out var perPlayer))
                {
                    perPlayer = new Dictionary<string, TimeSpan>();
                    playerFirstProduction[player] = perPlayer;
                }
                if (!perPlayer.ContainsKey(unitName))
                {
                    perPlayer[unitName] = time;
                }
            }
        }

        private static void RecordConstruction(
            TimeSpan time,
            CommandChunk command,
            int player,
            Dictionary<uint, TimeSpan> unitFirstObserved,
            HashSet<uint> builderUnits,
            Dictionary<int, HashSet<uint>> playerStrongOwnership)
        {
            // Data[0]: ObjectId (builder)
            // Data[1]: AsciiString (building name)
            RecordBuilder(time, command, player, unitFirstObserved, builderUnits, playerStrongOwnership);
        }

        private static void RecordPlaceBuilding(
            TimeSpan time,
            CommandChunk command,
            int player,
            Dictionary<uint, TimeSpan> unitFirstObserved,
            HashSet<uint> builderUnits,
            Dictionary<int, HashSet<uint>> playerStrongOwnership)
        {
            // Data[0]: ObjectId (builder)
            // Data[1]: AsciiString (building name)
            // Data[2]: Int32 (count)
            // Data[3]: Vector3 (position)
            // Data[4]: Float32 (angle)
            RecordBuilder(time, command, player, unitFirstObserved, builderUnits, playerStrongOwnership);
        }

        private static void RecordBuilder(
            TimeSpan time,
            CommandChunk command,
            int player,
            Dictionary<uint, TimeSpan> unitFirstObserved,
            HashSet<uint> builderUnits,
            Dictionary<int, HashSet<uint>> playerStrongOwnership)
        {
            foreach (var entry in command.Data)
            {
                if (entry.Type is not (CommandArgumentType.ObjectId or CommandArgumentType.ObjectId_2))
                {
                    continue;
                }
                if (entry.Count == 1)
                {
                    var id = (uint)entry.Value;
                    if (id != 0)
                    {
                        TryRecordFirstObserved(id, time, unitFirstObserved);
                        builderUnits.Add(id);
                        RecordPlayerOwnership(player, id, playerStrongOwnership);
                    }
                    return; // only first ObjectId is the builder
                }
                if (entry.Value is uint[] ids)
                {
                    foreach (var id in ids)
                    {
                        if (id != 0)
                        {
                            TryRecordFirstObserved(id, time, unitFirstObserved);
                            builderUnits.Add(id);
                            RecordPlayerOwnership(player, id, playerStrongOwnership);
                            return;
                        }
                    }
                }
            }
        }

        private static void RecordControlGroupCreate(
            TimeSpan time,
            CommandChunk command,
            int player,
            Dictionary<uint, TimeSpan> unitFirstObserved,
            Dictionary<int, HashSet<uint>> playerStrongOwnership,
            Dictionary<long, HashSet<uint>> controlGroups)
        {
            // Data[0]: Int32 编队号；Data[1..]: 成员 ObjectId
            int? groupNumber = null;
            var members = new List<uint>();
            foreach (var entry in command.Data)
            {
                if (groupNumber is null && entry.Type == CommandArgumentType.Int32)
                {
                    groupNumber = entry.Count == 1 && entry.Value is int singleInt
                        ? singleInt
                        : entry.Value is int[] ints && ints.Length > 0
                            ? ints[0]
                            : (int?)null;
                    continue;
                }
                if (entry.Type is CommandArgumentType.ObjectId or CommandArgumentType.ObjectId_2)
                {
                    if (entry.Count == 1)
                    {
                        var id = (uint)entry.Value;
                        if (id != 0) members.Add(id);
                    }
                    else
                    {
                        foreach (var id in (uint[])entry.Value)
                        {
                            if (id != 0) members.Add(id);
                        }
                    }
                }
            }

            if (groupNumber is null || members.Count == 0)
            {
                return;
            }
            var group = new HashSet<uint>(members);
            controlGroups[ControlGroupKey(player, groupNumber.Value)] = group;
            foreach (var id in members)
            {
                TryRecordFirstObserved(id, time, unitFirstObserved);
                RecordPlayerOwnership(player, id, playerStrongOwnership);
            }
        }

        private static void RecordControlGroupSelect(
            TimeSpan time,
            CommandChunk command,
            int player,
            Dictionary<uint, TimeSpan> unitFirstObserved,
            Dictionary<int, HashSet<uint>> playerStrongOwnership,
            Dictionary<long, HashSet<uint>> controlGroups)
        {
            foreach (var entry in command.Data)
            {
                if (entry.Type != CommandArgumentType.Int32)
                {
                    continue;
                }
                var groupNumber = entry.Count == 1 && entry.Value is int singleInt
                    ? singleInt
                    : entry.Value is int[] ints && ints.Length > 0
                        ? ints[0]
                        : (int?)null;
                if (groupNumber is null)
                {
                    continue;
                }
                if (controlGroups.TryGetValue(ControlGroupKey(player, groupNumber.Value), out var members))
                {
                    foreach (var id in members)
                    {
                        TryRecordFirstObserved(id, time, unitFirstObserved);
                        RecordPlayerOwnership(player, id, playerStrongOwnership);
                    }
                }
            }
        }

        private static long ControlGroupKey(int player, int group) =>
            ((long)player << 32) | (uint)group;

        private static void RecordTechChoice(
            TimeSpan time,
            CommandChunk command,
            int player,
            Dictionary<int, HashSet<string>> playerTechChoices)
        {
            foreach (var entry in command.Data)
            {
                if (entry.Type is CommandArgumentType.AsciiString or CommandArgumentType.UnicodeString)
                {
                    var tech = entry.Count == 1
                        ? entry.Value.ToString()
                        : entry.Value is string[] strings
                            ? strings.FirstOrDefault(s => !string.IsNullOrWhiteSpace(s))
                            : null;
                    if (string.IsNullOrWhiteSpace(tech))
                    {
                        continue;
                    }
                    if (!playerTechChoices.TryGetValue(player, out var set))
                    {
                        set = new HashSet<string>();
                        playerTechChoices[player] = set;
                    }
                    set.Add(tech);
                    return;
                }
            }
        }

        private static void RecordObjectReferenceWithOwnership(
            TimeSpan time,
            CommandChunk command,
            int player,
            Dictionary<uint, TimeSpan> unitFirstObserved,
            Dictionary<int, HashSet<uint>> playerOwnership)
        {
            foreach (var entry in command.Data)
            {
                if (entry.Type is not (CommandArgumentType.ObjectId or CommandArgumentType.ObjectId_2))
                {
                    continue;
                }
                if (entry.Count == 1)
                {
                    var id = (uint)entry.Value;
                    if (id == 0) continue;
                    TryRecordFirstObserved(id, time, unitFirstObserved);
                    RecordPlayerOwnership(player, id, playerOwnership);
                }
                else
                {
                    foreach (var id in (uint[])entry.Value)
                    {
                        if (id == 0) continue;
                        TryRecordFirstObserved(id, time, unitFirstObserved);
                        RecordPlayerOwnership(player, id, playerOwnership);
                    }
                }
            }
        }

        private static void RecordObjectReference(
            TimeSpan time,
            CommandChunk command,
            Dictionary<uint, TimeSpan> unitFirstObserved)
        {
            foreach (var entry in command.Data)
            {
                if (entry.Type is CommandArgumentType.ObjectId or CommandArgumentType.ObjectId_2)
                {
                    if (entry.Count == 1)
                    {
                        TryRecordFirstObserved((uint)entry.Value, time, unitFirstObserved);
                    }
                    else
                    {
                        foreach (var id in (uint[])entry.Value)
                        {
                            TryRecordFirstObserved(id, time, unitFirstObserved);
                        }
                    }
                }
            }
        }

        private static void RecordPlayerSelection(int player, uint unitId, Dictionary<int, HashSet<uint>> playerSelected)
        {
            if (!playerSelected.TryGetValue(player, out var set))
            {
                set = new HashSet<uint>();
                playerSelected[player] = set;
            }
            set.Add(unitId);
        }

        private static void RecordPlayerOwnership(
            int player,
            uint unitId,
            Dictionary<int, HashSet<uint>> playerOwnership)
        {
            if (!playerOwnership.TryGetValue(player, out var set))
            {
                set = new HashSet<uint>();
                playerOwnership[player] = set;
            }
            set.Add(unitId);
        }

        private static void TryRecordFirstObserved(uint unitId, TimeSpan time, Dictionary<uint, TimeSpan> unitFirstObserved)
        {
            if (unitId == 0) return;
            if (unitFirstObserved.ContainsKey(unitId)) return;
            unitFirstObserved[unitId] = time;
        }
    }
}