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; }

        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)
        {
            UnitIdFirstObservedTime = unitIdFirstObservedTime;
            UnitIdSpecialPowers = unitIdSpecialPowers;
            BuilderUnitIds = builderUnitIds;
            ProducerUnitIds = producerUnitIds;
            PlayerFirstProductionTime = playerFirstProductionTime;
            PlayerSelectedUnitIds = playerSelectedUnitIds;
        }

        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>>();

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

            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()));
        }

        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)
        {
            var player = command.PlayerIndex;
            var cmdId = command.CommandId;

            switch (cmdId)
            {
                // select unit(s): 0x1F5
                case 0x1F5:
                    RecordSelectUnit(time, command, player, unitFirstObserved, playerSelected);
                    break;

                // special power (no target): 0x1FE
                case 0x1FE:
                    RecordSpecialPower(time, command, stringHashTable, unitFirstObserved, unitSpecialPowers);
                    break;

                // special power (target position): 0x1FF
                case 0x1FF:
                    RecordSpecialPower(time, command, stringHashTable, unitFirstObserved, unitSpecialPowers);
                    break;

                // special power (target position and angle): 0x200
                case 0x200:
                    RecordSpecialPower(time, command, stringHashTable, unitFirstObserved, unitSpecialPowers);
                    break;

                // special power (target unit): 0x201
                case 0x201:
                    RecordSpecialPower(time, command, stringHashTable, unitFirstObserved, unitSpecialPowers);
                    break;

                // special power (one or more targets): 0x232
                case 0x232:
                    RecordSpecialPower(time, command, stringHashTable, unitFirstObserved, unitSpecialPowers);
                    break;

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

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

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

                // sell building: 0x20A
                case 0x20A:
                    RecordObjectReference(time, command, unitFirstObserved);
                    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)
        {
            // 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);
                    }
                    else
                    {
                        foreach (var id in (uint[])entry.Value)
                        {
                            TryRecordFirstObserved(id, time, unitFirstObserved);
                            RecordPlayerSelection(player, id, playerSelected);
                        }
                    }
                }
            }
        }

        private static void RecordSpecialPower(
            TimeSpan time,
            CommandChunk command,
            IReadOnlyDictionary<uint, string> stringHashTable,
            Dictionary<uint, TimeSpan> unitFirstObserved,
            Dictionary<uint, HashSet<string>> unitSpecialPowers)
        {
            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
                        var hash = unchecked((uint)(int)entry.Value);
                        powerName = stringHashTable.TryGetValue(hash, out var name)
                            ? name
                            : $"Hash_{hash: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);
            }
        }

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

            foreach (var entry in command.Data)
            {
                switch (entry.Type)
                {
                    case CommandArgumentType.ObjectId or CommandArgumentType.ObjectId_2
                        when entry.Count == 1 && producerId is null:
                        producerId = (uint)entry.Value;
                        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);
            }

            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,
            Dictionary<uint, TimeSpan> unitFirstObserved,
            HashSet<uint> builderUnits)
        {
            // Data[0]: ObjectId (builder)
            // Data[1]: AsciiString (building name)
            RecordBuilder(time, command, unitFirstObserved, builderUnits);
        }

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

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

        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 TryRecordFirstObserved(uint unitId, TimeSpan time, Dictionary<uint, TimeSpan> unitFirstObserved)
        {
            if (unitId == 0) return;
            if (unitFirstObserved.ContainsKey(unitId)) return;
            unitFirstObserved[unitId] = time;
        }
    }
}