using Theriapolis.Core.Util;
using Theriapolis.Core.World.Polylines;

namespace Theriapolis.Core.World.Generation.Stages;

/// <summary>
/// Stage 10 — HydrologyGen
/// Drainage simulation → rivers as polylines in world-pixel space + lakes.
/// Per-tile HasRiver and RiverAdjacent flags derived from polylines.
/// </summary>
public sealed class HydrologyGenStage : IWorldGenStage
{
    public string Name => "HydrologyGen";

    private const int W = C.WORLD_WIDTH_TILES;
    private const int H = C.WORLD_HEIGHT_TILES;

    // 8 neighbor deltas: N, NE, E, SE, S, SW, W, NW
    private static readonly (int dx, int dy)[] Dirs8 =
    {
        ( 0,-1), ( 1,-1), ( 1, 0), ( 1, 1),
        ( 0, 1), (-1, 1), (-1, 0), (-1,-1),
    };

    public void Run(WorldGenContext ctx)
    {
        var world = ctx.World;
        var rng   = ctx.Rngs["hydro"];

        ctx.ReportProgress(Name, 0f);

        // ── Step 1: Flow direction ────────────────────────────────────────────
        var flowDir = ComputeFlowDirections(world);

        ctx.ReportProgress(Name, 0.20f);

        // ── Step 2: Flow accumulation ─────────────────────────────────────────
        var accumulation = ComputeFlowAccumulation(world, flowDir);

        ctx.ReportProgress(Name, 0.40f);

        // ── Step 3: Lake detection ────────────────────────────────────────────
        DetectAndFillLakes(world, flowDir, accumulation, rng);

        ctx.ReportProgress(Name, 0.55f);

        // ── Step 4: Extract river paths ───────────────────────────────────────
        var riverPaths = ExtractRiverPaths(world, accumulation, flowDir);

        ctx.ReportProgress(Name, 0.65f);

        // ── Step 5: Guarantee rivers in required regions ──────────────────────
        EnsureRequiredRivers(world, riverPaths, accumulation, rng);

        ctx.ReportProgress(Name, 0.70f);

        // ── Step 6: Carve elevation along rivers ──────────────────────────────
        foreach (var path in riverPaths)
            CarveRiverPath(world, path.tiles);

        ctx.ReportProgress(Name, 0.80f);

        // ── Step 7: Convert to polylines ──────────────────────────────────────
        var meanderRng = SeededRng.ForSubsystem(world.WorldSeed, C.RNG_MEANDER);
        int polyId = 0;
        foreach (var path in riverPaths)
        {
            var poly = BuildRiverPolyline(world, path, polyId++, meanderRng);
            world.Rivers.Add(poly);
        }

        ctx.ReportProgress(Name, 0.92f);

        // ── Step 8: Rasterize tile flags ──────────────────────────────────────
        foreach (var poly in world.Rivers)
            PolylineBuilder.RasterizeToTileFlags(world, poly, setAdjacency: true);

        world.StageHashes["HydrologyGen"] = world.HashPolylines();
        ctx.LogMessage($"[HydrologyGen] Generated {world.Rivers.Count} rivers.");
        ctx.ReportProgress(Name, 1f);
    }

    // ── Flow direction map ───────────────────────────────────────────────────

    private byte[,] ComputeFlowDirections(WorldState world)
    {
        var flow = new byte[W, H];

        // Initialize to Dir.None
        for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++)
            flow[x, y] = Dir.None;

        for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++)
        {
            ref var tile = ref world.TileAt(x, y);
            if (tile.Biome == BiomeId.Ocean) continue;

            float lowestElev = tile.Elevation;
            int bestDx = 0, bestDy = 0;

            foreach (var (ddx, ddy) in Dirs8)
            {
                int nx = x + ddx;
                int ny = y + ddy;
                if ((uint)nx >= W || (uint)ny >= H) continue;
                float ne = world.Tiles[nx, ny].Elevation;
                if (ne < lowestElev)
                {
                    lowestElev = ne;
                    bestDx = ddx; bestDy = ddy;
                }
            }

            if (bestDx != 0 || bestDy != 0)
                flow[x, y] = Dir.FromDelta(bestDx, bestDy);
        }

        // Handle flat areas: BFS from edges outward
        ResolveFlatAreas(world, flow);
        return flow;
    }

    private void ResolveFlatAreas(WorldState world, byte[,] flow)
    {
        // BFS-compute ocean distance for every cell.
        // Flow always moves towards lower ocean distance → guaranteed cycle-free.
        var oceanDist = new int[W, H];
        const int INF = int.MaxValue / 2;
        for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++)
            oceanDist[x, y] = world.Tiles[x, y].Biome == BiomeId.Ocean ? 0 : INF;

        var bfsQ = new Queue<(int x, int y)>();
        for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++)
            if (oceanDist[x, y] == 0) bfsQ.Enqueue((x, y));

        while (bfsQ.Count > 0)
        {
            var (cx, cy) = bfsQ.Dequeue();
            foreach (var (ddx, ddy) in Dirs8)
            {
                int nx = cx + ddx, ny = cy + ddy;
                if ((uint)nx >= W || (uint)ny >= H) continue;
                if (oceanDist[nx, ny] != INF) continue;
                oceanDist[nx, ny] = oceanDist[cx, cy] + 1;
                bfsQ.Enqueue((nx, ny));
            }
        }

        // Assign direction to every Dir.None land tile:
        //  - Prefer the neighbor with strictly lower ocean distance.
        //  - Tiebreak (same distance): prefer lower elevation, then earliest in Dirs8.
        for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++)
        {
            if (flow[x, y] != Dir.None) continue;
            if (world.Tiles[x, y].Biome == BiomeId.Ocean) continue;

            int myDist    = oceanDist[x, y];
            int bestDist  = myDist;
            float bestElev = float.MaxValue;
            int bdx = 0, bdy = 0;

            for (int di = 0; di < Dirs8.Length; di++)
            {
                var (ddx, ddy) = Dirs8[di];
                int nx = x + ddx, ny = y + ddy;
                if ((uint)nx >= W || (uint)ny >= H) continue;

                int   nd = oceanDist[nx, ny];
                float ne = world.Tiles[nx, ny].Elevation;

                if (nd < bestDist || (nd == bestDist && ne < bestElev))
                {
                    bestDist = nd; bestElev = ne;
                    bdx = ddx; bdy = ddy;
                }
            }

            if (bdx != 0 || bdy != 0)
                flow[x, y] = Dir.FromDelta(bdx, bdy);
        }
    }

    // ── Flow accumulation ────────────────────────────────────────────────────

    /// <summary>
    /// Topological-sort based accumulation.  Processes each tile exactly once in
    /// upstream-to-downstream order, so flat areas with cycles are handled correctly.
    /// Cycle tiles (in-degree never reaches 0) receive a self-contribution of 1.
    /// </summary>
    private int[,] ComputeFlowAccumulation(WorldState world, byte[,] flowDir)
    {
        var accum    = new int[W, H];
        var inDeg    = new int[W, H];
        var processed = new bool[W, H];

        // Count in-degrees: how many non-ocean land tiles flow into each tile
        for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++)
        {
            if (world.Tiles[x, y].Biome == BiomeId.Ocean) continue;
            byte d = flowDir[x, y];
            if (d == Dir.None) continue;
            var (ddx, ddy) = Dir.ToDelta(d);
            int nx = x + ddx, ny = y + ddy;
            if ((uint)nx >= W || (uint)ny >= H) continue;
            if (world.Tiles[nx, ny].Biome == BiomeId.Ocean) continue;
            inDeg[nx, ny]++;
        }

        // Seed: all source tiles (no upstream contributors)
        var queue = new Queue<(int x, int y)>();
        for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++)
        {
            if (world.Tiles[x, y].Biome == BiomeId.Ocean) continue;
            if (inDeg[x, y] == 0) queue.Enqueue((x, y));
        }

        while (queue.Count > 0)
        {
            var (cx, cy) = queue.Dequeue();
            if (processed[cx, cy]) continue;
            processed[cx, cy] = true;

            accum[cx, cy]++; // self-contribution

            byte d = flowDir[cx, cy];
            if (d == Dir.None) continue;
            var (ddx, ddy) = Dir.ToDelta(d);
            int nx = cx + ddx, ny = cy + ddy;
            if ((uint)nx >= W || (uint)ny >= H) continue;
            if (world.Tiles[nx, ny].Biome == BiomeId.Ocean) continue;

            accum[nx, ny] += accum[cx, cy];
            if (--inDeg[nx, ny] == 0)
                queue.Enqueue((nx, ny));
        }

        // Tiles still unprocessed are in cycles — assign self-contribution
        for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++)
        {
            if (world.Tiles[x, y].Biome == BiomeId.Ocean) continue;
            if (!processed[x, y]) accum[x, y] = Math.Max(1, accum[x, y]);
        }

        return accum;
    }

    // ── Lake detection and filling ───────────────────────────────────────────

    private void DetectAndFillLakes(WorldState world, byte[,] flowDir, int[,] accum, SeededRng rng)
    {
        var visited = new bool[W, H];

        for (int y = 1; y < H - 1; y++)
        for (int x = 1; x < W - 1; x++)
        {
            if (visited[x, y]) continue;
            if (world.Tiles[x, y].Biome == BiomeId.Ocean) continue;
            if (flowDir[x, y] != Dir.None) continue; // has a direction — not a sink

            // BFS flood fill from this sink point
            var basin = new List<(int, int)>();
            var q = new Queue<(int, int)>();
            q.Enqueue((x, y));
            visited[x, y] = true;
            float outletElev = world.Tiles[x, y].Elevation;

            while (q.Count > 0)
            {
                var (cx, cy) = q.Dequeue();
                basin.Add((cx, cy));

                foreach (var (ddx, ddy) in Dirs8)
                {
                    int nx2 = cx + ddx;
                    int ny2 = cy + ddy;
                    if ((uint)nx2 >= W || (uint)ny2 >= H) continue;
                    if (visited[nx2, ny2]) continue;
                    if (world.Tiles[nx2, ny2].Biome == BiomeId.Ocean) continue;
                    if (world.Tiles[nx2, ny2].Elevation > outletElev + 0.05f) continue;
                    visited[nx2, ny2] = true;
                    q.Enqueue((nx2, ny2));
                }
            }

            if (basin.Count < C.LAKE_MIN_AREA) continue;

            // Fill the basin as a lake (set elevation to outlet and biome to Ocean for rendering)
            foreach (var (lx, ly) in basin)
            {
                ref var tile = ref world.TileAt(lx, ly);
                tile.Elevation = outletElev;
                tile.Biome     = BiomeId.Ocean; // reuse ocean for inland lakes (renders as water)
                tile.Features |= FeatureFlags.HasRiver; // mark as water body
                flowDir[lx, ly] = Dir.None;
            }
        }
    }

    // ── River path extraction ────────────────────────────────────────────────

    private record RiverPath(List<(int x, int y)> tiles, int maxAccum);

    private List<RiverPath> ExtractRiverPaths(WorldState world, int[,] accum, byte[,] flowDir)
    {
        var paths = new List<RiverPath>();
        var onRiver = new bool[W, H];

        // Find all tiles with flow >= threshold, sorted by accumulation (descending)
        var highAccum = new List<(int accum, int x, int y)>();
        for (int y = 0; y < H; y++)
        for (int x = 0; x < W; x++)
        {
            if (accum[x, y] >= C.RIVER_MIN_FLOW_ACCUM && world.Tiles[x, y].Biome != BiomeId.Ocean)
                highAccum.Add((accum[x, y], x, y));
        }
        // Sort ASCENDING so we start at headwaters (low accum = most upstream).
        // Each headwater traces one long path to coast; downstream tiles are marked onRiver
        // and skipped by subsequent iterations — producing one river per tributary.
        highAccum.Sort((a, b) => a.accum.CompareTo(b.accum));

        foreach (var (_, sx, sy) in highAccum)
        {
            if (paths.Count >= C.RIVER_MAX_COUNT) break; // hard cap
            if (onRiver[sx, sy]) continue;

            var tilePath = new List<(int x, int y)>();
            int curX = sx, curY = sy;
            int maxAccumVal = 0;
            int safety = W + H;

            while (safety-- > 0)
            {
                if (onRiver[curX, curY]) break;
                if (world.Tiles[curX, curY].Biome == BiomeId.Ocean) break;

                tilePath.Add((curX, curY));
                maxAccumVal = Math.Max(maxAccumVal, accum[curX, curY]);

                // Navigate downstream via precomputed flow direction
                byte dir = flowDir[curX, curY];
                if (dir == Dir.None) break; // true sink
                var (ddx, ddy) = Dir.ToDelta(dir);
                int nextX = curX + ddx, nextY = curY + ddy;
                if ((uint)nextX >= W || (uint)nextY >= H) break;
                curX = nextX; curY = nextY;
            }

            // Only mark tiles as onRiver AFTER confirming path length.
            // Premature marking (before the check) would block subsequent upstream traces.
            if (tilePath.Count >= 3)
            {
                foreach (var (tx, ty) in tilePath)
                    onRiver[tx, ty] = true;
                paths.Add(new RiverPath(tilePath, maxAccumVal));
            }
        }

        return paths;
    }

    // ── Ensure required macro-region rivers ──────────────────────────────────

    private void EnsureRequiredRivers(WorldState world, List<RiverPath> paths, int[,] accum, SeededRng rng)
    {
        // Check which required regions have rivers
        var required = new[]
        {
            "temperate_deciduous", "temperate_forest",
            "temperate_grassland",
            "subtropical_forest",
        };

        var covered = new HashSet<string>();
        foreach (var path in paths)
        {
            if (path.maxAccum < C.RIVER_MODERATE_THRESHOLD) continue;
            foreach (var (tx, ty) in path.tiles)
            {
                var macro = world.MacroCellForTile(world.TileAt(tx, ty));
                covered.Add(macro.BiomeType.ToLowerInvariant());
            }
        }

        // For each uncovered required region, find its highest tile and add a short river
        foreach (var req in required)
        {
            if (covered.Contains(req)) continue;

            // Find the highest-elevation tile in this macro region
            int bestX = -1, bestY = -1;
            float bestElev = 0f;
            for (int y = 0; y < H; y += 4)
            for (int x = 0; x < W; x += 4)
            {
                ref var tile = ref world.TileAt(x, y);
                if (tile.Biome == BiomeId.Ocean) continue;
                var macro = world.MacroCellForTile(in tile);
                if (!macro.BiomeType.ToLowerInvariant().Contains(req.Split('_')[0])) continue;
                if (tile.Elevation > bestElev)
                {
                    bestElev = tile.Elevation; bestX = x; bestY = y;
                }
            }
            if (bestX < 0) continue;

            // Trace a path from that tile downhill
            var tilePath = new List<(int, int)>();
            int cx = bestX, cy = bestY;
            int safety = 200;
            while (safety-- > 0 && world.Tiles[cx, cy].Biome != BiomeId.Ocean)
            {
                tilePath.Add((cx, cy));
                int bx = cx, by = cy;
                float be = world.Tiles[cx, cy].Elevation;
                foreach (var (ddx, ddy) in Dirs8)
                {
                    int nx = cx + ddx, ny = cy + ddy;
                    if ((uint)nx >= W || (uint)ny >= H) continue;
                    if (world.Tiles[nx, ny].Elevation < be)
                    {
                        be = world.Tiles[nx, ny].Elevation; bx = nx; by = ny;
                    }
                }
                if (bx == cx && by == cy) break;
                cx = bx; cy = by;
            }

            if (tilePath.Count >= 5)
            {
                // Override accumulation to moderate threshold so it qualifies
                foreach (var (tx, ty) in tilePath)
                    accum[tx, ty] = Math.Max(accum[tx, ty], C.RIVER_MODERATE_THRESHOLD);
                paths.Add(new RiverPath(tilePath, C.RIVER_MODERATE_THRESHOLD));
            }
        }
    }

    // ── Elevation carving ────────────────────────────────────────────────────

    private static void CarveRiverPath(WorldState world, List<(int x, int y)> tiles)
    {
        foreach (var (x, y) in tiles)
        {
            ref var tile = ref world.TileAt(x, y);
            if (tile.Biome == BiomeId.Ocean) continue;
            // Don't carve mountain tiles — river carving would push them below the macro
            // ElevationFloor constraint that BiomeAssignStage enforced.
            if (tile.Biome is BiomeId.MountainAlpine or BiomeId.MountainForested) continue;
            tile.Elevation = Math.Max(WorldState.SeaLevel, tile.Elevation - C.RIVER_CARVE_DEPTH);
        }
    }

    // ── Polyline conversion ──────────────────────────────────────────────────

    private static Polyline BuildRiverPolyline(WorldState world, RiverPath path, int id, SeededRng rng)
    {
        var poly = new Polyline { Type = PolylineType.River, Id = id };
        poly.FlowAccumulation = path.maxAccum;

        // Classify
        if (path.maxAccum >= C.RIVER_MAJOR_THRESHOLD)
        {
            poly.RiverClassification = RiverClass.MajorRiver;
            poly.Width = 3f;
        }
        else if (path.maxAccum >= C.RIVER_MODERATE_THRESHOLD)
        {
            poly.RiverClassification = RiverClass.River;
            poly.Width = 2f;
        }
        else
        {
            poly.RiverClassification = RiverClass.Stream;
            poly.Width = 1f;
        }

        // Control points from tile centers
        var controlPts = path.tiles
            .Select(t => PolylineBuilder.TileToWorldPixel(t.x, t.y))
            .ToList();

        // Smooth
        var smoothed = PolylineBuilder.CatmullRomSmooth(controlPts);

        // Meander amplitude based on average elevation
        float avgElev = path.tiles.Average(t => world.Tiles[t.x, t.y].Elevation);
        float amp = avgElev > 0.55f ? C.MEANDER_AMP_MOUNTAIN : C.MEANDER_AMP_FLAT;

        ulong noiseSeed = world.WorldSeed ^ C.RNG_HYDRO ^ (ulong)id;
        PolylineBuilder.ApplyMeanderNoise(smoothed, amp, C.MEANDER_FREQ, noiseSeed);

        poly.Points.AddRange(smoothed);
        poly.SimplifiedPoints = PolylineBuilder.RDPSimplify(poly.Points);

        return poly;
    }
}
