#include "AiSystem.h"

#include <optional>
#include <unordered_map>
#include <vector>

#include <QVector2D>

#include "Building.h"
#include "BuildingSystem.h"
#include "BuildingType.h"
#include "BuildingId.h"
#include "EntityAdmin.h"
#include "FactionComponent.h"
#include "HealthComponent.h"
#include "HomeReturnBehaviorComponent.h"
#include "HqProxyComponent.h"
#include "ModuleOwnerComponent.h"
#include "MovementIntentComponent.h"
#include "PositionComponent.h"
#include "RallyBehaviorComponent.h"
#include "RepairBehaviorComponent.h"
#include "RepairToolComponent.h"
#include "SalvageBehaviorComponent.h"
#include "SalvageCargoComponent.h"
#include "ScrapSystem.h"
#include "SensorRangeComponent.h"
#include "ShipIdentityComponent.h"
#include "StationBodyComponent.h"
#include "ThreatResponseBehaviorComponent.h"

// ---------------------------------------------------------------------------
// tickHomeReturnBehavior  (priority 4)
// ---------------------------------------------------------------------------

void AiSystem::tickHomeReturnBehavior(EntityAdmin& admin)
{
    admin.forEach<HomeReturnBehaviorComponent, HealthComponent, MovementIntentComponent>(
        [](entt::entity /*e*/, const HomeReturnBehaviorComponent& homeReturnBehavior,
           const HealthComponent& h, MovementIntentComponent& intent)
        {
            if (h.hp / h.maxHp < homeReturnBehavior.retreatHpFraction)
            {
                if (4 > intent.priority)
                {
                    intent = MovementIntentComponent{4, homeReturnBehavior.homePos};
                }
            }
        });
}

// ---------------------------------------------------------------------------
// tickThreatResponseBehavior  (priority 3)
// ---------------------------------------------------------------------------

void AiSystem::tickThreatResponseBehavior(EntityAdmin& admin, const BuildingSystem& buildings)
{
    // Snapshot all combatant entities for target acquisition.
    struct CombatantInfo
    {
        entt::entity entity;
        QVector2D position;
        bool isEnemy;
        bool isStation;
    };
    std::vector<CombatantInfo> combatants;

    admin.forEach<PositionComponent, FactionComponent, ShipIdentityComponent>(
        [&combatants](entt::entity e, const PositionComponent& pos,
                      const FactionComponent& f, const ShipIdentityComponent& /*si*/)
        {
            combatants.push_back({e, pos.value, f.isEnemy, false});
        });

    admin.forEach<PositionComponent, FactionComponent, StationBodyComponent>(
        [&combatants](entt::entity e, const PositionComponent& pos,
                      const FactionComponent& f, const StationBodyComponent& /*sb*/)
        {
            combatants.push_back({e, pos.value, f.isEnemy, true});
        });

    admin.forEach<PositionComponent, FactionComponent, HqProxyComponent>(
        [&combatants](entt::entity e, const PositionComponent& pos,
                      const FactionComponent& f, const HqProxyComponent& /*hq*/)
        {
            combatants.push_back({e, pos.value, f.isEnemy, true});
        });

    admin.forEach<ThreatResponseBehaviorComponent, PositionComponent, FactionComponent,
                  SensorRangeComponent, MovementIntentComponent>(
        [&](entt::entity e, ThreatResponseBehaviorComponent& threatResponseBehavior,
            PositionComponent& pos, FactionComponent& faction,
            SensorRangeComponent& sensor, MovementIntentComponent& intent)
        {
            const float range = sensor.value;

            // Validate current target.
            bool targetValid = false;
            if (threatResponseBehavior.currentTarget)
            {
                const entt::entity t = *threatResponseBehavior.currentTarget;
                if (admin.isValid(t) && admin.hasAll<PositionComponent>(t))
                {
                    const float dist =
                        (admin.get<PositionComponent>(t).value - pos.value).length();
                    if (dist <= range)
                    {
                        targetValid = true;
                    }
                }
            }

            if (!targetValid)
            {
                threatResponseBehavior.currentTarget = std::nullopt;
                float bestDist = range;

                for (const CombatantInfo& c : combatants)
                {
                    if (c.entity == e) { continue; }

                    bool isValidTarget = false;
                    if (!faction.isEnemy)
                    {
                        isValidTarget = c.isEnemy;
                    }
                    else
                    {
                        isValidTarget = !c.isEnemy;
                    }
                    if (!isValidTarget) { continue; }

                    const float dist = (c.position - pos.value).length();
                    if (dist < bestDist)
                    {
                        bestDist = dist;
                        threatResponseBehavior.currentTarget = c.entity;
                    }
                }
            }

            if (threatResponseBehavior.currentTarget)
            {
                const entt::entity t = *threatResponseBehavior.currentTarget;
                QVector2D dest = pos.value;
                if (admin.isValid(t) && admin.hasAll<PositionComponent>(t))
                {
                    dest = admin.get<PositionComponent>(t).value;
                }
                if (3 > intent.priority)
                {
                    intent = MovementIntentComponent{3, dest};
                }
            }
            else
            {
                if (3 > intent.priority)
                {
                    if (admin.hasAll<RallyBehaviorComponent>(e))
                    {
                        intent = MovementIntentComponent{
                            3, admin.get<RallyBehaviorComponent>(e).rallyPoint};
                    }
                    else if (!faction.isEnemy)
                    {
                        intent = MovementIntentComponent{
                            3, QVector2D(pos.value.x() + 1000.0f, pos.value.y())};
                    }
                    else
                    {
                        intent = MovementIntentComponent{
                            3, QVector2D(-10000.0f, pos.value.y())};
                    }
                }
            }
        });
}

// ---------------------------------------------------------------------------
// tickRepairBehavior  (priority 2)
// ---------------------------------------------------------------------------

void AiSystem::tickRepairBehavior(EntityAdmin& admin, BuildingSystem& buildings)
{
    // Snapshot all entities with health for repair targeting.
    struct RepairableInfo
    {
        entt::entity entity;
        QVector2D position;
        bool isEnemy;
        bool isShip;
        float hp;
        float maxHp;
    };
    std::vector<RepairableInfo> repairables;

    admin.forEach<ShipIdentityComponent, PositionComponent, FactionComponent, HealthComponent>(
        [&repairables](entt::entity e, const ShipIdentityComponent& /*si*/,
                       const PositionComponent& pos, const FactionComponent& f,
                       const HealthComponent& h)
        {
            repairables.push_back({e, pos.value, f.isEnemy, true, h.hp, h.maxHp});
        });

    admin.forEach<StationBodyComponent, PositionComponent, FactionComponent, HealthComponent>(
        [&repairables](entt::entity e, const StationBodyComponent& /*sb*/,
                       const PositionComponent& pos, const FactionComponent& f,
                       const HealthComponent& h)
        {
            repairables.push_back({e, pos.value, f.isEnemy, false, h.hp, h.maxHp});
        });

    // Snapshot enemy ships for threat detection.
    struct EnemyInfo
    {
        QVector2D position;
    };
    std::vector<EnemyInfo> enemies;
    admin.forEach<ShipIdentityComponent, PositionComponent, FactionComponent>(
        [&enemies](entt::entity /*e*/, const ShipIdentityComponent& /*si*/,
                   const PositionComponent& pos, const FactionComponent& f)
        {
            if (f.isEnemy)
            {
                enemies.push_back({pos.value});
            }
        });

    admin.forEach<RepairBehaviorComponent, PositionComponent,
                  FactionComponent, SensorRangeComponent, MovementIntentComponent>(
        [&](entt::entity e, RepairBehaviorComponent& rb,
            PositionComponent& pos, FactionComponent& /*faction*/,
            SensorRangeComponent& sensor, MovementIntentComponent& intent)
        {
            const float repairRange = rb.maxRepairRange;

            // Flee if enemy nearby.
            bool enemyNearby = false;
            for (const EnemyInfo& enemy : enemies)
            {
                if ((enemy.position - pos.value).length() <= sensor.value)
                {
                    enemyNearby = true;
                    break;
                }
            }
            if (enemyNearby)
            {
                if (2 > intent.priority)
                {
                    intent = MovementIntentComponent{
                        2, QVector2D(-10000.0f, pos.value.y())};
                }
                return;
            }

            // Validate current target.
            bool targetValid = false;
            if (rb.currentTarget)
            {
                const entt::entity t = *rb.currentTarget;
                if (admin.isValid(t) && admin.hasAll<HealthComponent>(t))
                {
                    const HealthComponent& th = admin.get<HealthComponent>(t);
                    if (th.hp > 0.0f && th.hp < th.maxHp)
                    {
                        targetValid = true;
                    }
                }
            }

            if (!targetValid)
            {
                rb.currentTarget = std::nullopt;
                float bestDist = sensor.value;

                for (const RepairableInfo& r : repairables)
                {
                    if (r.entity == e) { continue; }
                    if (r.isEnemy) { continue; }
                    if (r.hp >= r.maxHp) { continue; }
                    const float dist = (r.position - pos.value).length();
                    if (dist < bestDist)
                    {
                        bestDist = dist;
                        rb.currentTarget = r.entity;
                    }
                }
            }

            if (!rb.currentTarget)
            {
                if (2 > intent.priority)
                {
                    intent = MovementIntentComponent{
                        2, QVector2D(pos.value.x() + 1000.0f, pos.value.y())};
                }
                return;
            }

            const entt::entity target = *rb.currentTarget;
            QVector2D targetPos = pos.value;
            if (admin.isValid(target) && admin.hasAll<PositionComponent>(target))
            {
                targetPos = admin.get<PositionComponent>(target).value;
            }

            if (2 > intent.priority)
            {
                intent = MovementIntentComponent{2, targetPos};
            }
        });
}

// ---------------------------------------------------------------------------
// tickRepairTools
// ---------------------------------------------------------------------------

void AiSystem::tickRepairTools(EntityAdmin& admin)
{
    admin.forEach<RepairToolComponent, ModuleOwnerComponent>(
        [&](entt::entity /*e*/, RepairToolComponent& rt, const ModuleOwnerComponent& owner)
        {
            if (!admin.hasAll<RepairBehaviorComponent>(owner.owner)) { return; }
            const RepairBehaviorComponent& rb =
                admin.get<RepairBehaviorComponent>(owner.owner);
            if (!rb.currentTarget) { return; }

            const entt::entity target = *rb.currentTarget;
            if (!admin.isValid(target) || !admin.hasAll<HealthComponent>(target)) { return; }

            const PositionComponent& ownerPos = admin.get<PositionComponent>(owner.owner);
            const PositionComponent& targetPos = admin.get<PositionComponent>(target);
            const float dist = (targetPos.value - ownerPos.value).length();
            if (dist > rt.range) { return; }

            HealthComponent& targetHealth = admin.get<HealthComponent>(target);
            targetHealth.hp = std::min(targetHealth.hp + rt.ratePerTick, targetHealth.maxHp);
        });
}

// ---------------------------------------------------------------------------
// tickSalvageBehavior  (priority 1)
// ---------------------------------------------------------------------------

void AiSystem::tickSalvageBehavior(EntityAdmin& admin, ScrapSystem& scraps,
                                    BuildingSystem& buildings)
{
    // Snapshot enemy ships for threat detection.
    struct EnemyShipPos
    {
        QVector2D position;
    };
    std::vector<EnemyShipPos> enemyShips;
    admin.forEach<ShipIdentityComponent, PositionComponent, FactionComponent>(
        [&enemyShips](entt::entity /*e*/, const ShipIdentityComponent& /*si*/,
                      const PositionComponent& pos, const FactionComponent& f)
        {
            if (f.isEnemy)
            {
                enemyShips.push_back({pos.value});
            }
        });

    // Aggregate cargo across all salvage-module children per owning ship.
    struct AggregatedCargo
    {
        int totalCurrent  = 0;
        int totalCapacity = 0;
    };
    std::unordered_map<entt::entity, AggregatedCargo> cargoByShip;
    admin.forEach<SalvageCargoComponent, ModuleOwnerComponent>(
        [&](entt::entity /*ce*/, const SalvageCargoComponent& c, const ModuleOwnerComponent& o)
        {
            AggregatedCargo& agg = cargoByShip[o.owner];
            agg.totalCurrent  += c.current;
            agg.totalCapacity += c.capacity;
        });

    const std::vector<ScrapInfo> allScrap = scraps.allScrapInfo();

    admin.forEach<SalvageBehaviorComponent, PositionComponent,
                  SensorRangeComponent, MovementIntentComponent>(
        [&](entt::entity e, SalvageBehaviorComponent& salvageBehavior,
            PositionComponent& pos,
            SensorRangeComponent& sensor, MovementIntentComponent& intent)
        {
            const float collectRange = salvageBehavior.maxCollectionRange;
            const AggregatedCargo& cargoState = cargoByShip[e];

            // Assign nearest SalvageBay if needed.
            if (salvageBehavior.deliveryBay == kInvalidBuildingId)
            {
                const Building* bay = buildings.findNearestBuilding(pos.value,
                                                                     BuildingType::SalvageBay);
                if (bay)
                {
                    salvageBehavior.deliveryBay = bay->id;
                }
            }

            const BuildingId bayId = salvageBehavior.deliveryBay;

            QVector2D bayPos = pos.value;
            if (bayId != kInvalidBuildingId)
            {
                const Building* bay = buildings.findBuilding(bayId);
                if (bay)
                {
                    bayPos = QVector2D(bay->anchor.x() + bay->footprint.width()  / 2.0f,
                                       bay->anchor.y() + bay->footprint.height() / 2.0f);
                }
            }

            const bool cargoFull = (cargoState.totalCurrent >= cargoState.totalCapacity
                                    && cargoState.totalCapacity > 0);

            if (cargoFull)
            {
                if (1 > intent.priority)
                {
                    intent = MovementIntentComponent{1, bayPos};
                }
                if (bayId != kInvalidBuildingId
                    && (pos.value - bayPos).length() <= 1.0f)
                {
                    // Decrement first non-empty salvage child.
                    bool delivered = false;
                    admin.forEach<SalvageCargoComponent, ModuleOwnerComponent>(
                        [&](entt::entity /*ce*/, SalvageCargoComponent& c,
                            const ModuleOwnerComponent& o)
                        {
                            if (delivered || o.owner != e || c.current <= 0) { return; }
                            if (buildings.deliverScrapToSalvageBay(bayId))
                            {
                                --c.current;
                                delivered = true;
                            }
                        });
                }
                return;
            }

            // Retreat if enemy near and cargo empty.
            bool retreating = false;
            if (cargoState.totalCurrent == 0)
            {
                for (const EnemyShipPos& enemy : enemyShips)
                {
                    if ((enemy.position - pos.value).length() <= collectRange)
                    {
                        if (1 > intent.priority)
                        {
                            intent = MovementIntentComponent{
                                1, QVector2D(-10000.0f, pos.value.y())};
                        }
                        retreating = true;
                        break;
                    }
                }
            }
            if (retreating) { return; }

            // Collect nearby scrap — increment first non-full salvage child.
            for (const ScrapInfo& si : allScrap)
            {
                if ((si.position - pos.value).length() <= collectRange)
                {
                    bool collected = false;
                    admin.forEach<SalvageCargoComponent, ModuleOwnerComponent>(
                        [&](entt::entity /*ce*/, SalvageCargoComponent& c,
                            const ModuleOwnerComponent& o)
                        {
                            if (collected || o.owner != e || c.current >= c.capacity) { return; }
                            if (scraps.consume(si.entity))
                            {
                                ++c.current;
                                salvageBehavior.scrapTarget = std::nullopt;
                                collected = true;
                            }
                        });
                    break;
                }
            }

            // Move toward scrap target or find a new one.
            if (salvageBehavior.scrapTarget)
            {
                if (1 > intent.priority)
                {
                    intent = MovementIntentComponent{1, *salvageBehavior.scrapTarget};
                }
            }
            else
            {
                float bestDist = sensor.value;
                std::optional<QVector2D> bestPos;
                for (const ScrapInfo& si : allScrap)
                {
                    const float dist = (si.position - pos.value).length();
                    if (dist < bestDist)
                    {
                        bestDist = dist;
                        bestPos  = si.position;
                    }
                }
                if (bestPos)
                {
                    salvageBehavior.scrapTarget = bestPos;
                    if (1 > intent.priority)
                    {
                        intent = MovementIntentComponent{1, *bestPos};
                    }
                }
                else
                {
                    if (1 > intent.priority)
                    {
                        intent = MovementIntentComponent{
                            1, QVector2D(pos.value.x() + 1000.0f, pos.value.y())};
                    }
                }
            }
        });
}