#!/usr/bin/env python3
"""Lay the zones of world/zones.json onto the island and draw them.

Reads the surface-biome grid of the recipe's seed (written by
`/cobblesurvival biomepreview`, which tools/seed_search.py runs for every
candidate) and assigns every cell to the first zone whose rule matches, in the
order of world/zones.json. Writes:

  build/zone-map.png     the island coloured by zone, with a legend
  build/zone-layout.json the zone of every grid cell plus area per zone
  datapacks/core/data/cobblesurvival/zones/island.json
                         the same grid for the companion mod (zone lookup,
                         the "zone" spawn condition)

Usage: tools/zone_layout.py [--preview PATH] [--scale PX_PER_CELL]
"""
import argparse
import json
import math
import sys
from collections import deque
from pathlib import Path

import numpy as np

ROOT = Path(__file__).resolve().parent.parent
SAMPLE_RADIUS = 1550  # survey points stay this close to 0,0, well inside the world border
DATAPACK_LAYOUT = ROOT / "datapacks/core/data/cobblesurvival/zones/island.json"
PREVIEWS = ROOT / "build/server/seed-search/cobblesurvival/biome_preview"


def components(mask):
    """Label 8-connected True cells; returns (labels, sizes) with labels -1 outside the mask."""
    labels = np.full(mask.shape, -1, dtype=np.int32)
    sizes = []
    rows, cols = mask.shape
    for start in zip(*np.nonzero(mask)):
        if labels[start] != -1:
            continue
        label, queue = len(sizes), deque([start])
        labels[start] = label
        size = 0
        while queue:
            r, c = queue.popleft()
            size += 1
            for dr in (-1, 0, 1):
                for dc in (-1, 0, 1):
                    rr, cc = r + dr, c + dc
                    if 0 <= rr < rows and 0 <= cc < cols and mask[rr, cc] and labels[rr, cc] == -1:
                        labels[rr, cc] = label
                        queue.append((rr, cc))
        sizes.append(size)
    return labels, sizes


def drop_small(mask, min_cells):
    labels, sizes = components(mask)
    keep = np.array([s >= min_cells for s in sizes] + [False])
    return keep[labels]


def dilate(mask, cells):
    out = mask.copy()
    for _ in range(cells):
        grown = out.copy()
        grown[1:, :] |= out[:-1, :]
        grown[:-1, :] |= out[1:, :]
        grown[:, 1:] |= out[:, :-1]
        grown[:, :-1] |= out[:, 1:]
        grown[1:, 1:] |= out[:-1, :-1]
        grown[:-1, :-1] |= out[1:, 1:]
        grown[1:, :-1] |= out[:-1, 1:]
        grown[:-1, 1:] |= out[1:, :-1]
        out = grown
    return out


def inherit_from_land(zone, mask):
    """Cells in mask take the most common zone among their assigned neighbours, spreading inwards."""
    todo = mask & (zone == -1)
    while todo.any():
        changed = False
        for r, c in zip(*np.nonzero(todo)):
            around = zone[max(r - 1, 0):r + 2, max(c - 1, 0):c + 2]
            known = around[around >= 0]
            if known.size:
                zone[r, c] = np.bincount(known).argmax()
                changed = True
        todo = mask & (zone == -1)
        if not changed:
            break


def assign(preview, plan):
    n, step, radius = preview["size"], preview["step"], preview["radius"]
    names = preview["biomes"]
    surface = np.array(preview["surface"]).reshape(n, n)
    water_biomes = [b for spec in plan["zones"] for b in spec["rule"].get("water", [])] + plan["inherit"]
    # Zones follow biomes, not terrain height: a lake in the forest belongs to the forest's zone.
    land = ~np.isin(surface, [names.index(b) for b in water_biomes if b in names])
    coords = -radius + np.arange(n) * step
    x = np.broadcast_to(coords[None, :], (n, n))
    z = np.broadcast_to(coords[:, None], (n, n))
    dist = np.hypot(x, z)

    def biome_mask(biomes):
        ids = [names.index(b) for b in biomes if b in names]
        return np.isin(surface, ids)

    zone = np.full((n, n), -1, dtype=np.int32)
    for index, spec in enumerate(plan["zones"]):
        rule, free = spec["rule"], zone == -1
        if "within" in rule:
            mask = land & free & (dist <= rule["within"])
        elif "biomes" in rule:
            mask = drop_small(land & free & biome_mask(rule["biomes"]), plan["min_cells"])
            for biome, blocks in rule.get("absorb", {}).items():
                near = dilate(mask, math.ceil(blocks / step))
                mask |= land & free & near & biome_mask([biome])
        elif "region" in rule:
            r = rule["region"]
            mask = land & free
            if "x_min" in r: mask &= x >= r["x_min"]
            if "x_max" in r: mask &= x < r["x_max"]
            if "z_min" in r: mask &= z >= r["z_min"]
            if "z_max" in r: mask &= z < r["z_max"]
        elif "water" in rule:
            inherit_from_land(zone, biome_mask(plan["inherit"]))
            mask = ~land & (zone == -1) & biome_mask(rule["water"])
        else:
            sys.exit(f"zone {spec['id']}: unknown rule {rule}")
        zone[mask] = index
    return zone, land, np.array(preview["height"]).reshape(n, n)


def summary(zone, land, plan, step):
    total_land = int(land.sum()) or 1
    out = {}
    for index, spec in enumerate(plan["zones"]):
        mask = zone == index
        _, sizes = components(mask)
        out[spec["id"]] = {
            "cells": int(mask.sum()),
            "km2": round(float(mask.sum()) * step * step / 1e6, 2),
            "land_percent": round(100 * float((mask & land).sum()) / total_land, 1),
            "pieces": len(sizes),
        }
    out["_unassigned_land_cells"] = int((land & (zone == -1)).sum())
    return out


def hex_rgb(h):
    return tuple(int(h[i:i + 2], 16) for i in (1, 3, 5))


def render(zone, land, height, plan, stats, preview, out, scale):
    from PIL import Image, ImageDraw, ImageFont
    n = zone.shape[0]
    colours = np.array([hex_rgb(z["color"]) for z in plan["zones"]] + [(40, 40, 40)], dtype=np.float32)
    rgb = colours[zone]
    # Hill shading on land, so mountains read as mountains.
    shade = np.zeros_like(height, dtype=np.float32)
    shade[1:, 1:] = (height[1:, 1:] - height[:-1, :-1]) / 12.0
    above_sea = height >= preview["sea_level"]
    rgb[above_sea] *= np.clip(1.0 + shade[above_sea], 0.7, 1.25)[:, None]
    # Zone borders darker.
    edge = np.zeros_like(land)
    edge[1:, :] |= zone[1:, :] != zone[:-1, :]
    edge[:, 1:] |= zone[:, 1:] != zone[:, :-1]
    rgb[edge] *= 0.55
    img = Image.fromarray(np.clip(rgb, 0, 255).astype(np.uint8)).resize((n * scale, n * scale), Image.NEAREST)

    try:
        font = ImageFont.truetype("DejaVuSans.ttf", 15)
        bold = ImageFont.truetype("DejaVuSans-Bold.ttf", 13)
    except OSError:
        font = bold = ImageFont.load_default()
    legend_w = 470
    canvas = Image.new("RGB", (img.width + legend_w, img.height), (24, 24, 24))
    canvas.paste(img, (0, 0))
    draw = ImageDraw.Draw(canvas)

    # Name each land zone on its biggest piece.
    for index, spec in enumerate(plan["zones"]):
        if "water" in spec["rule"]:
            continue
        labels, sizes = components(zone == index)
        if not sizes:
            continue
        rows, cols = np.nonzero(labels == int(np.argmax(sizes)))
        cx, cy = (cols.mean() + 0.5) * scale, (rows.mean() + 0.5) * scale
        text = spec["name"]
        w = draw.textlength(text, font=bold)
        draw.rectangle((cx - w / 2 - 3, cy - 9, cx + w / 2 + 3, cy + 9), fill=(0, 0, 0))
        draw.text((cx - w / 2, cy - 8), text, fill=(255, 255, 255), font=bold)

    # Spawn.
    r0 = preview["radius"]
    sx, sz = (0 + r0) / preview["step"] * scale, (32 + r0) / preview["step"] * scale
    draw.ellipse((sx - 5, sz - 5, sx + 5, sz + 5), fill=(255, 220, 0), outline=(0, 0, 0))

    x0, y = img.width + 16, 16
    draw.text((x0, y), "Zones (levels are placeholders)", fill=(255, 255, 255), font=font)
    y += 30
    for spec in plan["zones"]:
        s = stats[spec["id"]]
        draw.rectangle((x0, y + 2, x0 + 16, y + 18), fill=hex_rgb(spec["color"]), outline=(0, 0, 0))
        lv = f"Lv {spec['levels'][0]}-{spec['levels'][1]}"
        where = f"{s['land_percent']}% of land" if "water" not in spec["rule"] else f"{s['km2']} sq km"
        draw.text((x0 + 24, y), f"{spec['name']}  {lv}  ({where})", fill=(235, 235, 235), font=font)
        draw.text((x0 + 24, y + 18), f"hazard: {spec['hazard']}"[:58], fill=(170, 170, 170), font=font)
        y += 44
    y += 6
    draw.text((x0, y), "Underground (by cave biome, island-wide):", fill=(255, 255, 255), font=font)
    y += 24
    for cave in plan["caves"]:
        draw.text((x0 + 8, y), f"{cave['name']}  Lv {cave['levels'][0]}-{cave['levels'][1]}  {cave['hazard']}"[:62],
                  fill=(235, 235, 235), font=font)
        y += 22
    y += 10
    draw.text((x0, y), f"1 cell = {preview['step']} blocks; north is up; yellow dot: spawn", fill=(170, 170, 170), font=font)
    canvas.save(out)


def samples(zone, plan, preview):
    """{zone id: [x, y or None, z]}: a point deep inside each zone, for the spawn survey."""
    n, step, radius = preview["size"], preview["step"], preview["radius"]
    coords = -radius + np.arange(n) * step
    near = np.hypot(coords[None, :], coords[:, None]) <= SAMPLE_RADIUS  # inside the world border
    dry = np.array(preview["height"]).reshape(n, n) > preview["sea_level"]
    out = {}
    for index, spec in enumerate(plan["zones"]):
        mask = (zone == index) & near
        if "water" not in spec["rule"]:
            mask &= dry  # land zones: stand on land, not on a lake inside the zone
        if not mask.any():
            continue
        # Peel the zone from its edges inwards; the last cells left are the deepest.
        inner = mask
        while True:
            shrunk = inner & ~dilate(~inner, 1)
            if not shrunk.any():
                break
            inner = shrunk
        rows, cols = np.nonzero(inner)
        r, c = rows[len(rows) // 2], cols[len(cols) // 2]
        out[spec["id"]] = [int(-radius + c * step), None, int(-radius + r * step)]
    names = preview["biomes"]
    for cave in plan["caves"]:
        if cave["biome"] not in names:
            continue
        biome = names.index(cave["biome"])
        best = None
        for y, cells in preview["caves"].items():
            grid = (np.array(cells).reshape(n, n) == biome) & near
            if not grid.any():
                continue
            inner = grid
            depth = 0
            while True:
                shrunk = inner & ~dilate(~inner, 1)
                if not shrunk.any():
                    break
                inner, depth = shrunk, depth + 1
            if best is None or depth > best[0]:
                rows, cols = np.nonzero(inner)
                r, c = rows[len(rows) // 2], cols[len(cols) // 2]
                best = (depth, [int(-radius + c * step), int(y), int(-radius + r * step)])
        if best:
            out[cave["id"]] = best[1]
    return out


def biome_samples(zone, plan, preview):
    """{zone id: [x, None, z]}: for zones made of biomes, a point deep inside their first biome (the hazard's own
    biome, e.g. the jagged peaks of Frostpeak), for the hazard tests."""
    n, step, radius = preview["size"], preview["step"], preview["radius"]
    names = preview["biomes"]
    surface = np.array(preview["surface"]).reshape(n, n)
    dry = np.array(preview["height"]).reshape(n, n) > preview["sea_level"]
    out = {}
    for index, spec in enumerate(plan["zones"]):
        biomes = spec["rule"].get("biomes")
        if not biomes or biomes[0] not in names:
            continue
        mask = (zone == index) & dry & (surface == names.index(biomes[0]))
        if not mask.any():
            continue
        inner = mask
        while True:
            shrunk = inner & ~dilate(~inner, 1)
            if not shrunk.any():
                break
            inner = shrunk
        rows, cols = np.nonzero(inner)
        r, c = rows[len(rows) // 2], cols[len(cols) // 2]
        out[spec["id"]] = [int(-radius + c * step), None, int(-radius + r * step)]
    return out


def write_datapack(zone, plan, preview):
    runs = []
    for value in zone.flatten().tolist():
        if runs and runs[-1][0] == value:
            runs[-1][1] += 1
        else:
            runs.append([value, 1])
    keep = ("id", "name", "levels")
    DATAPACK_LAYOUT.parent.mkdir(parents=True, exist_ok=True)
    DATAPACK_LAYOUT.write_text(json.dumps({
        "_comment": "Generated by tools/zone_layout.py from world/zones.json; do not edit by hand.",
        "seed": preview["seed"], "radius": preview["radius"], "step": preview["step"], "size": preview["size"],
        "zones": [{k: z[k] for k in keep} for z in plan["zones"]],
        "caves": [{**{k: c[k] for k in keep}, "biome": c["biome"]} for c in plan["caves"]],
        "level_pressure": plan.get("level_pressure", {}),
        "wanderers": plan.get("wanderers", {}),
        "samples": samples(zone, plan, preview),
        "biome_samples": biome_samples(zone, plan, preview),
        "grid_rle": runs,
    }, separators=(",", ":")) + "\n")


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--preview", type=Path, help="biome preview JSON (default: the recipe's seed)")
    parser.add_argument("--scale", type=int, default=6, help="pixels per grid cell in the map")
    args = parser.parse_args()
    recipe = json.loads((ROOT / "world/world.json").read_text())
    preview_path = args.preview or PREVIEWS / f"{recipe['seed']}.json"
    if not preview_path.exists():
        sys.exit(f"{preview_path} is missing: run tools/seed_search.py (or /cobblesurvival biomepreview "
                 f"{recipe['seed']} 1700 24 on the test server) first.")
    preview = json.loads(preview_path.read_text())
    plan = json.loads((ROOT / "world/zones.json").read_text())

    zone, land, height = assign(preview, plan)
    stats = summary(zone, land, plan, preview["step"])
    ids = [z["id"] for z in plan["zones"]]
    (ROOT / "build").mkdir(exist_ok=True)
    (ROOT / "build/zone-layout.json").write_text(json.dumps({
        "seed": preview["seed"], "radius": preview["radius"], "step": preview["step"], "size": preview["size"],
        "zones": ids, "grid": zone.flatten().tolist(), "stats": stats,
    }))
    render(zone, land, height, plan, stats, preview, ROOT / "build/zone-map.png", args.scale)
    write_datapack(zone, plan, preview)
    for spec in plan["zones"]:
        s = stats[spec["id"]]
        print(f"  {spec['name']:18s} Lv {spec['levels'][0]:2d}-{spec['levels'][1]:2d}  "
              f"{s['land_percent']:5.1f}% of land  {s['km2']:5.2f} km²  {s['pieces']} piece(s)")
    if stats["_unassigned_land_cells"]:
        print(f"  warning: {stats['_unassigned_land_cells']} land cells match no zone")
    print(f"wrote build/zone-map.png, build/zone-layout.json and {DATAPACK_LAYOUT.relative_to(ROOT)}")


if __name__ == "__main__":
    main()