using System.Collections.Generic;
using HarmonyLib;
using Model;
using Model.Physics;
using UnityEngine;

namespace RailroaderPhysicsOverhaul;

public static class ConsistFreezer
{
    // Manual freeze — explicit /rpf freeze command
    static readonly HashSet<uint> _frozenSetIds = new();
    public static bool IsFrozen(uint setId) => _frozenSetIds.Contains(setId);
    public static bool Freeze(uint setId)   => _frozenSetIds.Add(setId);
    public static bool Unfreeze(uint setId) => _frozenSetIds.Remove(setId);
    public static void ClearAll()           => _frozenSetIds.Clear();

    // Auto-freeze: skip the Verlet tick for consists that are far from the camera
    // and moving slowly enough that physics quality there doesn't matter.
    // This is our integrated replacement for the stock optimizer (AllCarsAtRest only).
    public static bool  ExcludeLocomotives        = true;
    public static bool  AutoFreezeEnabled         = true;
    public static float AutoFreezeDistance       = 200f;   // meters
    public static float AutoFreezeSpeedThreshold = 0.3f;  // m/s (~0.7 mph)

    static float AutoFreezeDistanceSq => AutoFreezeDistance * AutoFreezeDistance;

    // Per-tick frozen-car counters — flushed to Last* by FlushFrameCounts().
    // internal so ShouldSkipTickPatch (same file, different class) can increment them.
    internal static int _frameAtRestCars;
    internal static int _frameDistanceCars;
    public static int LastAtRestCars    { get; private set; }
    public static int LastDistanceCars  { get; private set; }

    // Per-tick set of cars whose Verlet tick is being skipped — for debug visualization.
    public static readonly HashSet<Car> FrozenCars = new();

    internal static void FlushFrameCounts()
    {
        LastAtRestCars   = _frameAtRestCars;
        LastDistanceCars = _frameDistanceCars;
        _frameAtRestCars   = 0;
        _frameDistanceCars = 0;
        FrozenCars.Clear();
    }

    // Called from ShouldSkipTickPatch. Internal so the patch class (same file) can reach it.
    internal static bool ShouldAutoFreeze(IntegrationSet set)
    {
        if (!AutoFreezeEnabled) return false;
        // ShouldSkipTick is per-IntegrationSet — we can't freeze freight cars while keeping
        // a coupled loco ticking. If ExcludeLocomotives is on, the whole consist stays live
        // so AI waypoint queues on the loco don't get interrupted by a physics freeze.
        if (ExcludeLocomotives)
        {
            foreach (Car car in set.Cars)
                if (car is BaseLocomotive) return false;
        }
        Camera cam = Camera.main;
        if (cam == null) return false;
        Vector3 camPos       = cam.transform.position;
        float   distSqThresh = AutoFreezeDistanceSq;
        float   speedThresh  = AutoFreezeSpeedThreshold;

        foreach (Car car in set.Cars)
        {
            Transform body = car.BodyTransform;
            if (body == null) continue;

            // If ANY car is within distance → don't freeze the consist.
            if ((body.position - camPos).sqrMagnitude < distSqThresh) return false;

            // If ANY car is moving above threshold → don't freeze.
            if (Mathf.Abs(car.velocity) >= speedThresh) return false;
        }

        // Every car is both far and slow — safe to skip the Verlet tick.
        return true;
    }
}

// Completely replaces the stock ShouldSkipTick getter.
// Priority order:
//   1. Manual freeze  → always skip
//   2. ForceActive    → never skip (profiling bypass)
//   3. AllCarsAtRest  → stock optimizer behaviour, replicated explicitly
//   4. Auto-freeze    → our distance + speed tier
[HarmonyPatch(typeof(IntegrationSet), "get_ShouldSkipTick")]
static class ShouldSkipTickPatch
{
    static bool Prefix(IntegrationSet __instance, ref bool __result)
    {
        if (ConsistFreezer.IsFrozen(__instance.Id))
        {
            foreach (Car c in __instance.Cars) ConsistFreezer.FrozenCars.Add(c);
            __result = true;
            return false;
        }

        if (PhysicsTimer.ForceActive)
        {
            __result = false;
            return false;
        }

        // Replicate stock optimizer behaviour explicitly so it works with or without
        // the stock optimizer mod installed.
        if (__instance.AllCarsAtRest())
        {
            if (ConsistFreezer.ExcludeLocomotives)
            {
                foreach (Car c in __instance.Cars)
                    if (c is BaseLocomotive) { __result = false; return false; }
            }
            foreach (Car c in __instance.Cars) ConsistFreezer.FrozenCars.Add(c);
            ConsistFreezer._frameAtRestCars += __instance.NumberOfCars;
            __result = true;
            return false;
        }

        __result = ConsistFreezer.ShouldAutoFreeze(__instance);
        if (__result)
        {
            foreach (Car c in __instance.Cars) ConsistFreezer.FrozenCars.Add(c);
            ConsistFreezer._frameDistanceCars += __instance.NumberOfCars;
        }
        return false; // always override — stock getter never runs
    }
}