#!/usr/bin/env python3
"""
Arrows Escape v4 - Chain-Lock Level Generator
================================================
NEW MECHANIC: strict chain-lock. At any point in the game, EXACTLY ONE
piece on the board is extractable (its exit lane is clear). Extracting
it frees exactly the next piece in a fixed sequence, and so on, like a
line of dominoes / a single unlock chain. The green target piece is
always LAST in the chain -- it cannot become free until every other
piece has been removed in the one and only order the puzzle allows.

Construction approach (place in REVERSE chain order -- target first,
then each predecessor, each one required to cross the exit-lane of
the piece placed immediately before it in this loop, which is exactly
the piece that comes right after it in the unlock chain):

  chain = [p0, p1, ..., p_{N-1}=target]   (p0 unlocks first, chain
                                            proceeds in index order)

  Place p_{N-1} (target) with a long lane -- nothing needs to block
  IT from being free early since it's placed first and by definition
  is last in the unlock order; every other piece will be required to
  cross its lane.

  For k = N-2 down to 0: place p_k such that its body crosses the
  lane of p_{k+1} (already placed), so p_{k+1} cannot free up until
  p_k is gone. p_k's own lane becomes the thing p_{k-1} must cross
  next.

  After construction, we run a full simulation (simulate_chain_is_
  strict) that walks the exact unlock order and confirms EXACTLY ONE
  piece is free at every step -- never zero (deadlock), never more
  than one (would break the single-choice guarantee) -- before
  accepting the level.

Each level entry in levels.json:
{
  id, difficulty, rows, cols,
  pieces: [ { id, cells: [[r,c],...], headDir } ],
  targetId: <id of the last piece in the chain (green)>,
  solutionOrder: [id, id, ...],   // the ONLY valid order; at every
                                   // step exactly one piece is free
  coinReward: int
}
"""

import json
import random

DIRS = {"N": (-1, 0), "S": (1, 0), "E": (0, 1), "W": (0, -1)}
ALL_DIRS = list(DIRS.keys())
OPPOSITE_DIR = {"N": "S", "S": "N", "E": "W", "W": "E"}


def in_bounds(r, c, rows, cols):
    return 0 <= r < rows and 0 <= c < cols


def exit_lane_cells(head_r, head_c, direction, rows, cols):
    dr, dc = DIRS[direction]
    cells = []
    r, c = head_r + dr, head_c + dc
    while in_bounds(r, c, rows, cols):
        cells.append((r, c))
        r += dr
        c += dc
    return cells


def grow_snake_complex(start, occupied, rows, cols, target_len, rng):
    """Random self-avoiding walk with a bias toward turns, producing a
    maze-like bent path. Returns the path (list of (r,c)) starting at
    `start`; may be shorter than target_len if it runs out of room."""
    path = [start]
    visited = set(occupied)
    visited.add(start)
    last_dir = None
    while len(path) < target_len:
        r, c = path[-1]
        cands = []
        for d in ALL_DIRS:
            if last_dir and d == OPPOSITE_DIR[last_dir] and rng.random() < 0.85:
                continue
            dr, dc = DIRS[d]
            nr, nc = r + dr, c + dc
            if in_bounds(nr, nc, rows, cols) and (nr, nc) not in visited:
                cands.append((d, nr, nc))
        if not cands:
            break
        if last_dir and rng.random() < 0.5:
            same = [cd for cd in cands if cd[0] == last_dir]
            d, nr, nc = same[0] if same else rng.choice(cands)
        else:
            d, nr, nc = rng.choice(cands)
        path.append((nr, nc))
        visited.add((nr, nc))
        last_dir = d
    return path


def walk_to_lane_then_grow(start, occupied, rows, cols, target_len, must_cross, rng, avoid_cells=None):
    """Grow a path starting at `start`: first greedily move toward the
    nearest cell of `must_cross` until touching it, then continue a
    normal complex random walk for the remaining length. Returns None
    if it can never reach must_cross from this start.

    `avoid_cells`, if given, is a set of cells the path should treat
    as blocked (in addition to `occupied`) EXCEPT where they overlap
    `must_cross` -- used to stop a piece from accidentally wandering
    through some OTHER piece's exit lane (which would let that other
    piece free up out of chain order)."""
    avoid = (avoid_cells or set()) - must_cross
    path = [start]
    visited = set(occupied) | avoid
    visited.add(start)
    cur = start
    reached = cur in must_cross
    steps = 0
    while not reached and steps < max(rows, cols) + 5:
        r, c = cur
        cands = []
        for d, (dr, dc) in DIRS.items():
            nr, nc = r + dr, c + dc
            if in_bounds(nr, nc, rows, cols) and (nr, nc) not in visited:
                cands.append((nr, nc))
        if not cands:
            break
        cands.sort(key=lambda cell: min(abs(cell[0]-lr)+abs(cell[1]-lc) for lr, lc in must_cross))
        nxt = cands[0] if (len(cands) == 1 or rng.random() < 0.8) else cands[1]
        path.append(nxt)
        visited.add(nxt)
        cur = nxt
        steps += 1
        if cur in must_cross:
            reached = True
    if not reached:
        return None

    piece_len = target_len
    last_dir = None
    while len(path) < piece_len:
        r, c = path[-1]
        cands = []
        for d, (dr, dc) in DIRS.items():
            if last_dir and d == OPPOSITE_DIR[last_dir] and rng.random() < 0.85:
                continue
            nr, nc = r + dr, c + dc
            if in_bounds(nr, nc, rows, cols) and (nr, nc) not in visited:
                cands.append((d, nr, nc))
        if not cands:
            break

        is_final_step = (len(path) == piece_len - 1)
        if is_final_step and len(cands) > 1:
            # bias the LAST move toward whichever direction leaves the
            # longest resulting exit lane, so the head doesn't end up
            # facing a wall 1 cell away (which would make this piece
            # nearly impossible for the next chain link to reach)
            def resulting_lane_len(cand):
                _, nr, nc = cand
                best = 0
                for dname, (ddr, ddc) in DIRS.items():
                    lane_cells = exit_lane_cells(nr, nc, dname, rows, cols)
                    free_len = 0
                    for cell in lane_cells:
                        if cell in visited:
                            break
                        free_len += 1
                    best = max(best, free_len)
                return best
            cands.sort(key=resulting_lane_len, reverse=True)
            # mostly take the best, sometimes 2nd-best for variety
            d, nr, nc = cands[0] if rng.random() < 0.7 else cands[min(1, len(cands)-1)]
        elif last_dir and rng.random() < 0.5:
            same = [cd for cd in cands if cd[0] == last_dir]
            d, nr, nc = same[0] if same else rng.choice(cands)
        else:
            d, nr, nc = rng.choice(cands)
        path.append((nr, nc))
        visited.add((nr, nc))
        last_dir = d
    return path


def bfs_reachable_to(goal_cells, occupied, rows, cols):
    """Return the set of all empty cells connected (via 4-directional
    free-space movement) to any cell in `goal_cells`. Cheap way to
    know in advance which starting positions can possibly reach the
    lane, instead of discovering it via repeated failed random walks."""
    from collections import deque
    visited = set()
    q = deque()
    for cell in goal_cells:
        if cell not in occupied:
            visited.add(cell)
            q.append(cell)
    while q:
        r, c = q.popleft()
        for dr, dc in DIRS.values():
            nr, nc = r + dr, c + dc
            if in_bounds(nr, nc, rows, cols) and (nr, nc) not in occupied and (nr, nc) not in visited:
                visited.add((nr, nc))
                q.append((nr, nc))
    return visited


def head_dir_of(path):
    if len(path) < 2:
        return None
    hr, hc = path[-1]
    phr, phc = path[-2]
    dr, dc = hr - phr, hc - phc
    for dname, (ddr, ddc) in DIRS.items():
        if (ddr, ddc) == (dr, dc):
            return dname
    return None


def build_level(level_num, cfg, rng):
    rows, cols = cfg["rows"], cfg["cols"]
    num_pieces = cfg["num_pieces"]
    min_len, max_len = cfg["piece_len"]
    target_id = num_pieces - 1  # last in the chain (green)

    for _attempt in range(50):
        # Each attempt gets a fresh, independent RNG derived from the
        # outer seed. Without this, a failed attempt leaves the shared
        # rng in a "used up" state that makes subsequent attempts
        # within the same call much less likely to succeed than an
        # equivalent fresh call would be.
        attempt_rng = random.Random(rng.randint(0, 10**9))
        occupied = set()
        piece_data = {}

        # ---- place target with a long lane (through the longer grid
        # dimension) ----
        long_dim_is_rows = rows >= cols
        if long_dim_is_rows:
            t_dir = attempt_rng.choice(["N", "S"])
            t_c = attempt_rng.randrange(cols)
            t_r = (rows - 2) if t_dir == "N" else 1
        else:
            t_dir = attempt_rng.choice(["E", "W"])
            t_r = attempt_rng.randrange(rows)
            t_c = (cols - 2) if t_dir == "W" else 1

        lane = exit_lane_cells(t_r, t_c, t_dir, rows, cols)
        min_lane_len = max(6, int(0.55 * max(rows, cols)))
        if len(lane) < min_lane_len:
            continue

        t_len = attempt_rng.randint(min_len, max_len)
        t_body = grow_snake_complex((t_r, t_c), occupied, rows, cols, t_len, attempt_rng)
        if not t_body or len(t_body) < 2:
            continue
        for cell in t_body:
            occupied.add(cell)
        piece_data[target_id] = {
            "cells": [[r, c] for r, c in reversed(t_body)],
            "headDir": t_dir,
        }
        prev_lane = set(lane)  # the lane the NEXT-placed piece must cross
        all_lanes_so_far = set(lane)  # union of every placed piece's own
                                       # exit lane, used to keep future
                                       # pieces from wandering through a
                                       # lane they aren't assigned to

        # ---- place the rest of the chain, each one crossing the
        # previous piece's lane, walking backward from target_id-1
        # down to 0 ----
        ok = True
        for pid in range(target_id - 1, -1, -1):
            placed = False

            # BFS from the lane cells backward through free space to
            # find every empty cell that can actually reach the lane
            # (much cheaper than blindly retrying random starts, which
            # tends to pick unreachable pockets repeatedly at high
            # board density).
            reachable = bfs_reachable_to(prev_lane, occupied, rows, cols)
            if not reachable:
                ok = False
                break
            candidates = list(reachable)
            attempt_rng.shuffle(candidates)

            for start in candidates[:80]:
                plen = attempt_rng.randint(min_len, max_len)
                path = walk_to_lane_then_grow(
                    start, occupied, rows, cols, plen, prev_lane, attempt_rng,
                    avoid_cells=all_lanes_so_far,
                )
                if not path or len(path) < 2:
                    continue
                hd = head_dir_of(path)
                if hd is None:
                    continue
                own_lane = exit_lane_cells(path[-1][0], path[-1][1], hd, rows, cols)
                if pid > 0 and len(own_lane) < 3:
                    continue
                # this new piece's body must not cross any lane other
                # than the one it was explicitly walked toward (that
                # would let some other piece free up out of order)
                path_set = set(path)
                if (path_set & all_lanes_so_far) - prev_lane:
                    continue

                for cell in path:
                    occupied.add(cell)
                piece_data[pid] = {
                    "cells": [[r, c] for r, c in path],
                    "headDir": hd,
                }
                placed = True
                prev_lane = set(own_lane)
                all_lanes_so_far |= set(own_lane)
                break
            if not placed:
                ok = False
                break

        if not ok or len(piece_data) != num_pieces:
            continue

        pieces_out = [
            {"id": pid, "cells": piece_data[pid]["cells"], "headDir": piece_data[pid]["headDir"]}
            for pid in range(num_pieces)
        ]
        chain_order = list(range(num_pieces))

        level = {
            "id": level_num,
            "difficulty": cfg["name"],
            "rows": rows,
            "cols": cols,
            "pieces": pieces_out,
            "targetId": target_id,
            "solutionOrder": chain_order,
            "coinReward": cfg["reward"],
        }

        if simulate_chain_is_strict(level):
            return level

    raise RuntimeError(f"Failed to build level {level_num}")


def simulate_chain_is_strict(level):
    """Verify the CHAIN-LOCK property: at every step of extracting
    pieces in solutionOrder, EXACTLY ONE piece is free (matches the
    next expected piece), never zero (deadlock) and never more than
    one (breaks the 'only one option' guarantee)."""
    rows, cols = level["rows"], level["cols"]
    pieces = {p["id"]: p for p in level["pieces"]}
    order = level["solutionOrder"]
    present = set(pieces.keys())

    def free_pieces(present_set):
        free = []
        for pid in present_set:
            piece = pieces[pid]
            hr, hc = piece["cells"][-1]
            lane = exit_lane_cells(hr, hc, piece["headDir"], rows, cols)
            occ = set()
            for other in present_set:
                if other == pid:
                    continue
                for cell in pieces[other]["cells"]:
                    occ.add(tuple(cell))
            if not any(cell in occ for cell in lane):
                free.append(pid)
        return free

    for expected in order:
        free = free_pieces(present)
        if len(free) != 1:
            return False
        if free[0] != expected:
            return False
        present.discard(expected)

    return len(present) == 0


def difficulty_for_level(n, total_levels=500):
    tier_size = total_levels // 5
    tier = min(4, (n - 1) // tier_size)
    # Piece count is capped in the 6-8 range across all tiers: the
    # strict chain-lock guarantee (every piece's lane reserved for
    # exactly one predecessor, no accidental crossings allowed) has a
    # hard reliability ceiling around 6-8 pieces regardless of grid
    # size, verified empirically. Difficulty instead comes from larger
    # grids (more complex, longer paths to visually trace) and longer
    # individual pieces.
    tiers = [
        {"name": "Easy",   "rows": 12, "cols": 9,  "num_pieces": 6, "piece_len": (4, 8),  "reward": 10},
        {"name": "Medium", "rows": 14, "cols": 10, "num_pieces": 7, "piece_len": (5, 9),  "reward": 10},
        {"name": "Hard",   "rows": 16, "cols": 11, "num_pieces": 7, "piece_len": (5, 10), "reward": 20},
        {"name": "Expert", "rows": 18, "cols": 12, "num_pieces": 8, "piece_len": (6, 11), "reward": 20},
        {"name": "Master", "rows": 20, "cols": 13, "num_pieces": 8, "piece_len": (6, 12), "reward": 20},
    ]
    return tiers[tier]


def validate_level(level):
    return simulate_chain_is_strict(level)


def main():
    TOTAL_LEVELS = 500
    rng = random.Random(20260704)
    levels = []
    for n in range(1, TOTAL_LEVELS + 1):
        cfg = difficulty_for_level(n, TOTAL_LEVELS)
        built = None
        for _try in range(50):
            try:
                candidate = build_level(n, cfg, random.Random(rng.randint(0, 10**9)))
                if validate_level(candidate):
                    built = candidate
                    break
            except RuntimeError:
                continue
        if not built:
            raise RuntimeError(f"Could not generate level {n}")
        levels.append(built)
        if n % 50 == 0:
            print(f"  Generated level {n}/{TOTAL_LEVELS}...")

    with open("levels.json", "w") as f:
        json.dump({"levels": levels}, f, indent=2)

    print(f"Generated {len(levels)} levels -> levels.json")
    by_diff = {}
    for lvl in levels:
        by_diff.setdefault(lvl["difficulty"], 0)
        by_diff[lvl["difficulty"]] += 1
    print(by_diff)
    piece_counts = [len(lvl["pieces"]) for lvl in levels]
    print(f"Pieces per level: min={min(piece_counts)}, max={max(piece_counts)}")


if __name__ == "__main__":
    main()
