#include "AiSystem.h"

#include <optional>
#include <vector>

#include <QVector2D>

#include "Building.h"
#include "BuildingSystem.h"
#include "BuildingType.h"
#include "EcsComponents.h"
#include "EntityAdmin.h"
#include "EntityId.h"
#include "MovementIntent.h"
#include "Scrap.h"
#include "ScrapSystem.h"
#include "Ship.h"

// ---------------------------------------------------------------------------
// tickHomeReturn  (priority 4)
// ---------------------------------------------------------------------------

void AiSystem::tickHomeReturn(EntityAdmin& admin)
{
    admin.forEach<HomeReturn, Health, MovementIntent>(
        [](entt::entity /*e*/, const HomeReturn& hr, const Health& h, MovementIntent& intent)
        {
            if (h.hp / h.maxHp < hr.retreatHpFraction)
            {
                if (4 > intent.priority)
                {
                    intent = MovementIntent{4, hr.homePos};
                }
            }
        });
}

// ---------------------------------------------------------------------------
// tickThreatResponse  (priority 3)
// ---------------------------------------------------------------------------

void AiSystem::tickThreatResponse(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<Position, Faction, ShipIdentity>(
        [&combatants](entt::entity e, const Position& pos, const Faction& f, const ShipIdentity& /*si*/)
        {
            combatants.push_back({e, pos.value, f.isEnemy, false});
        });

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

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

    admin.forEach<ThreatResponse, Position, Faction, SensorRange, MovementIntent>(
        [&](entt::entity e, ThreatResponse& threat, Position& pos, Faction& faction,
            SensorRange& sensor, MovementIntent& intent)
        {
            const float range = sensor.value;

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

            if (!targetValid)
            {
                threat.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;
                        threat.currentTarget = c.entity;
                    }
                }
            }

            if (threat.currentTarget)
            {
                const entt::entity t = *threat.currentTarget;
                QVector2D dest = pos.value;
                if (admin.isValid(t) && admin.hasAll<Position>(t))
                {
                    dest = admin.get<Position>(t).value;
                }
                if (3 > intent.priority)
                {
                    intent = MovementIntent{3, dest};
                }
            }
            else
            {
                if (3 > intent.priority)
                {
                    if (admin.hasAll<RallyBehavior>(e))
                    {
                        intent = MovementIntent{3, admin.get<RallyBehavior>(e).rallyPoint};
                    }
                    else if (!faction.isEnemy)
                    {
                        intent = MovementIntent{3, QVector2D(pos.value.x() + 1000.0f,
                                                              pos.value.y())};
                    }
                    else
                    {
                        intent = MovementIntent{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<ShipIdentity, Position, Faction, Health>(
        [&repairables](entt::entity e, const ShipIdentity& /*si*/,
                       const Position& pos, const Faction& f, const Health& h)
        {
            repairables.push_back({e, pos.value, f.isEnemy, true, h.hp, h.maxHp});
        });

    admin.forEach<StationBody, Position, Faction, Health>(
        [&repairables](entt::entity e, const StationBody& /*sb*/,
                       const Position& pos, const Faction& f, const Health& 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<ShipIdentity, Position, Faction>(
        [&enemies](entt::entity /*e*/, const ShipIdentity& /*si*/,
                   const Position& pos, const Faction& f)
        {
            if (f.isEnemy)
            {
                enemies.push_back({pos.value});
            }
        });

    admin.forEach<RepairBehavior, RepairTool, Position, Faction, SensorRange, MovementIntent>(
        [&](entt::entity e, RepairBehavior& rb, RepairTool& rt, Position& pos,
            Faction& /*faction*/, SensorRange& sensor, MovementIntent& intent)
        {
            const float repairRange = rt.range;

            // 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 = MovementIntent{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<Health>(t))
                {
                    const Health& th = admin.get<Health>(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 = MovementIntent{2, QVector2D(pos.value.x() + 1000.0f,
                                                          pos.value.y())};
                }
                return;
            }

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

            const float distToTarget = (targetPos - pos.value).length();
            if (distToTarget <= repairRange)
            {
                if (admin.isValid(target) && admin.hasAll<Health>(target))
                {
                    Health& targetHealth = admin.get<Health>(target);
                    targetHealth.hp = std::min(targetHealth.hp + rt.ratePerTick,
                                               targetHealth.maxHp);
                }
            }

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

// ---------------------------------------------------------------------------
// tickScrapCollector  (priority 1)
// ---------------------------------------------------------------------------

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

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

    admin.forEach<ScrapCollector, SalvageCargo, Position, SensorRange, MovementIntent>(
        [&](entt::entity /*e*/, ScrapCollector& sc, SalvageCargo& cargo,
            Position& pos, SensorRange& sensor, MovementIntent& intent)
        {
            const float collectRange = cargo.collectionRange;

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

            const EntityId bayId = sc.deliveryBay;

            QVector2D bayPos = pos.value;
            if (bayId != kInvalidEntityId)
            {
                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 = (cargo.current >= cargo.capacity);

            if (cargoFull)
            {
                if (1 > intent.priority)
                {
                    intent = MovementIntent{1, bayPos};
                }
                if (bayId != kInvalidEntityId
                    && (pos.value - bayPos).length() <= 1.0f)
                {
                    if (buildings.deliverScrapToSalvageBay(bayId))
                    {
                        --cargo.current;
                    }
                }
                return;
            }

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

            // Collect nearby scrap.
            for (const ScrapInfo& si : allScrap)
            {
                if ((si.position - pos.value).length() <= collectRange)
                {
                    if (scraps.consume(si.entity))
                    {
                        ++cargo.current;
                        sc.scrapTarget = std::nullopt;
                    }
                    break;
                }
            }

            // Move toward scrap target or find a new one.
            if (sc.scrapTarget)
            {
                if (1 > intent.priority)
                {
                    intent = MovementIntent{1, *sc.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)
                {
                    sc.scrapTarget = bestPos;
                    if (1 > intent.priority)
                    {
                        intent = MovementIntent{1, *bestPos};
                    }
                }
                else
                {
                    if (1 > intent.priority)
                    {
                        intent = MovementIntent{1, QVector2D(pos.value.x() + 1000.0f,
                                                              pos.value.y())};
                    }
                }
            }
        });
}