import numpy as np
from PIL import Image, ImageDraw, ImageFilter, ImageFont

from .fonts import get_cached_font

# (font_path, word, size, orient) -> (ink[h,w] uint8, bbox_left, bbox_top).
# The HD passes -- clearance refinement and the independent overlap audit --
# rasterize the same words at the same sizes, and rasterizing is the single
# most expensive thing either of them does, so they share one cache.
_HD_INK_CACHE = {}


def _word_ink(word, size, orient, font_path):
    """Return (ink array, bbox_left, bbox_top) for a word, or None if empty."""
    key = (font_path, word, int(size), orient)
    cached = _HD_INK_CACHE.get(key)
    if cached is not None or key in _HD_INK_CACHE:
        return cached
    font = get_cached_font(font_path, size)
    if orient:
        font = ImageFont.TransposedFont(font, orientation=orient)
    bbox = ImageDraw.Draw(Image.new("L", (1, 1))).textbbox((0, 0), word, font=font)
    glyph = font.getmask(word, mode="L")
    glyph_width, glyph_height = glyph.size
    if glyph_width <= 0 or glyph_height <= 0:
        result = None
    else:
        ink = np.asarray(glyph, dtype=np.uint8).reshape(glyph_height, glyph_width)
        result = (ink, int(bbox[0]), int(bbox[1]))
    _HD_INK_CACHE[key] = result
    if len(_HD_INK_CACHE) > 20000:
        for i, k in enumerate(list(_HD_INK_CACHE.keys())):
            if i % 2 == 0:
                _HD_INK_CACHE.pop(k, None)
    return result


def scale_layout_for_hd(layout, work_scale):
    if work_scale <= 0:
        raise ValueError("WORK_SCALE must be greater than zero")
    return [
        (
            text,
            max(1, int(size / work_scale)),
            (int(y / work_scale), int(x / work_scale)),
            orient,
            color,
        )
        for text, size, (y, x), orient, color in layout
    ]


def count_layout_overlap_pixels(layout, mask_shape, font_path, alpha_threshold=0):
    """Count final rendered pixels occupied by more than one layout entry."""
    height, width = mask_shape
    occupied = np.zeros((height, width), dtype=bool)
    overlaps = np.zeros((height, width), dtype=bool)

    for word, size, (y, x), orient, _color in layout:
        measured = _word_ink(word, size, orient, font_path)
        if measured is None:
            continue
        ink_arr, bbox_left, bbox_top = measured
        glyph_height, glyph_width = ink_arr.shape
        ink = ink_arr > alpha_threshold

        ink_x = int(x) + bbox_left
        ink_y = int(y) + bbox_top
        x0 = max(0, ink_x)
        y0 = max(0, ink_y)
        x1 = min(width, ink_x + glyph_width)
        y1 = min(height, ink_y + glyph_height)
        if x0 >= x1 or y0 >= y1:
            continue

        visible_ink = ink[y0 - ink_y:y1 - ink_y, x0 - ink_x:x1 - ink_x]
        occupied_region = occupied[y0:y1, x0:x1]
        overlaps[y0:y1, x0:x1] |= occupied_region & visible_ink
        occupied_region |= visible_ink

    return int(overlaps.sum())


def _find_free_placement(blocked, occupied, collision, stamp, base_y, base_x):
    """Find any canvas position where `collision` hits nothing already taken.

    Returns the (dy, dx) offset from (base_y, base_x), or None. Candidates are
    ranked by distance from the original spot so a relocated word stays as close
    to its intended position as possible.

    A position whose whole footprint is empty is guaranteed to fit, so the
    search first looks for those using an integral image, which rejects the vast
    majority of positions with two additions instead of a per-pixel test.
    """
    height, width = blocked.shape
    gh, gw = collision.shape
    if height - gh < 0 or width - gw < 0:
        return None

    # Search expanding windows around the intended spot instead of the whole
    # canvas: a relocated word almost always finds room nearby, and the integral
    # image costs time proportional to the area examined. The last radius covers
    # the full canvas, so nothing is missed if the neighbourhood really is full.
    for radius in (256, 1024, max(height, width)):
        y_lo = max(0, base_y - radius)
        x_lo = max(0, base_x - radius)
        y_hi = min(height, base_y + radius + gh)
        x_hi = min(width, base_x + radius + gw)
        if y_hi - y_lo < gh or x_hi - x_lo < gw:
            continue

        taken = blocked[y_lo:y_hi, x_lo:x_hi] | occupied[y_lo:y_hi, x_lo:x_hi]
        integral = np.pad(taken.astype(np.int32), ((1, 0), (1, 0))).cumsum(0).cumsum(1)
        # Footprint sum for every candidate top-left corner in the window. A
        # position whose whole footprint is empty is guaranteed to fit, so no
        # per-pixel mask test is needed.
        counts = (
            integral[gh:, gw:]
            - integral[:-gh, gw:]
            - integral[gh:, :-gw]
            + integral[:-gh, :-gw]
        )
        ys, xs = np.nonzero(counts == 0)
        if ys.size == 0:
            continue
        dy = ys.astype(np.int64) + y_lo - base_y
        dx = xs.astype(np.int64) + x_lo - base_x
        best = int(np.argmin(dy * dy + dx * dx))
        return int(dy[best]), int(dx[best])

    return None


def refine_layout_with_hd_clearance(
    layout,
    mask,
    font_path,
    clearance=1,
    max_shift=24,
):
    """Validate the whole HD batch and minimally move rasterization collisions."""
    height, width = mask.shape
    blocked = np.asarray(mask) != 0
    occupied = np.zeros((height, width), dtype=bool)
    refined = []
    shifted_words = 0
    max_applied_shift = 0

    offset_rings = [[(0, 0)]]
    for radius in range(1, max_shift + 1):
        ring = [
            (dy, dx)
            for dy in range(-radius, radius + 1)
            for dx in range(-radius, radius + 1)
            if max(abs(dy), abs(dx)) == radius
        ]
        ring.sort(key=lambda item: (item[0] * item[0] + item[1] * item[1], item[0], item[1]))
        offset_rings.append(ring)

    for word, size, (draw_y, draw_x), orient, color in layout:
        measured = _word_ink(word, size, orient, font_path)
        if measured is None:
            refined.append((word, size, (draw_y, draw_x), orient, color))
            continue

        glyph_ink, bbox_left, bbox_top = measured
        glyph_height, glyph_width = glyph_ink.shape
        pad = max(0, int(clearance))
        padded = np.zeros((glyph_height + 2 * pad, glyph_width + 2 * pad), dtype=np.uint8)
        padded[pad:pad + glyph_height, pad:pad + glyph_width] = glyph_ink
        if pad:
            collision = np.asarray(
                Image.fromarray(padded).filter(ImageFilter.MaxFilter(2 * pad + 1)),
                dtype=np.uint8,
            ) > 0
        else:
            collision = padded > 0
        stamp = padded > 0

        base_y = int(draw_y) + bbox_top - pad
        base_x = int(draw_x) + bbox_left - pad

        def fits(y0, x0):
            y1 = y0 + collision.shape[0]
            x1 = x0 + collision.shape[1]
            if y0 < 0 or x0 < 0 or y1 > height or x1 > width:
                return False
            if np.any(blocked[y0:y1, x0:x1] & stamp):
                return False
            return not np.any(occupied[y0:y1, x0:x1] & collision)

        placed_offset = None
        for ring in offset_rings:
            for dy, dx in ring:
                y0 = base_y + dy
                x0 = base_x + dx
                if not fits(y0, x0):
                    continue
                placed_offset = (dy, dx)
                break
            if placed_offset is not None:
                break

        if placed_offset is None:
            # Nothing within max_shift. Rather than fail the batch -- which
            # makes the caller rebuild the entire cloud on a larger canvas, by
            # far the most expensive thing that can happen -- look for any free
            # spot on the whole canvas. This only runs for the occasional word
            # whose work-grid position does not survive the scale-up to HD.
            placed_offset = _find_free_placement(
                blocked, occupied, collision, stamp, base_y, base_x
            )

        if placed_offset is None:
            return None, {
                "shifted_words": shifted_words,
                "max_shift": max_applied_shift,
                "failed_word": word,
            }

        dy, dx = placed_offset
        y0 = base_y + dy
        x0 = base_x + dx
        occupied[y0:y0 + stamp.shape[0], x0:x0 + stamp.shape[1]] |= stamp
        if dy or dx:
            shifted_words += 1
            max_applied_shift = max(max_applied_shift, abs(dy), abs(dx))
        refined.append((word, size, (int(draw_y) + dy, int(draw_x) + dx), orient, color))

    return refined, {
        "shifted_words": shifted_words,
        "max_shift": max_applied_shift,
        "failed_word": None,
    }


def append_layout_with_hd_clearance(
    base_layout,
    additions,
    mask,
    font_path,
    clearance=1,
    max_shift=24,
    allow_global_search=False,
):
    """Place only *additions* against an already validated HD layout.

    The normal refinement pass must rebuild occupancy for every word because it
    is allowed to move the whole batch. Auto-repeat words are appended after the
    base batch has already passed refinement, so rescanning that batch is wasted
    work. This helper seeds occupancy from the base once, then processes only
    the new words and returns the largest collision-free prefix.
    """
    height, width = mask.shape
    blocked = np.asarray(mask) != 0
    occupied = np.zeros((height, width), dtype=bool)

    def stamp_existing(item):
        word, size, (draw_y, draw_x), orient, _color = item
        measured = _word_ink(word, size, orient, font_path)
        if measured is None:
            return
        ink, bbox_left, bbox_top = measured
        ink_y = int(draw_y) + bbox_top
        ink_x = int(draw_x) + bbox_left
        y0 = max(0, ink_y)
        x0 = max(0, ink_x)
        y1 = min(height, ink_y + ink.shape[0])
        x1 = min(width, ink_x + ink.shape[1])
        if y0 < y1 and x0 < x1:
            occupied[y0:y1, x0:x1] |= ink[y0 - ink_y:y1 - ink_y, x0 - ink_x:x1 - ink_x] > 0

    for item in base_layout:
        stamp_existing(item)

    offset_rings = [[(0, 0)]]
    for radius in range(1, max_shift + 1):
        ring = [
            (dy, dx)
            for dy in range(-radius, radius + 1)
            for dx in range(-radius, radius + 1)
            if max(abs(dy), abs(dx)) == radius
        ]
        ring.sort(key=lambda item: (item[0] * item[0] + item[1] * item[1], item[0], item[1]))
        offset_rings.append(ring)

    accepted = []
    shifted_words = 0
    max_applied_shift = 0
    failed_word = None
    for word, size, (draw_y, draw_x), orient, color in additions:
        measured = _word_ink(word, size, orient, font_path)
        if measured is None:
            accepted.append((word, size, (draw_y, draw_x), orient, color))
            continue

        glyph_ink, bbox_left, bbox_top = measured
        glyph_height, glyph_width = glyph_ink.shape
        pad = max(0, int(clearance))
        padded = np.zeros((glyph_height + 2 * pad, glyph_width + 2 * pad), dtype=np.uint8)
        padded[pad:pad + glyph_height, pad:pad + glyph_width] = glyph_ink
        if pad:
            collision = np.asarray(
                Image.fromarray(padded).filter(ImageFilter.MaxFilter(2 * pad + 1)),
                dtype=np.uint8,
            ) > 0
        else:
            collision = padded > 0
        stamp = padded > 0
        base_y = int(draw_y) + bbox_top - pad
        base_x = int(draw_x) + bbox_left - pad

        def fits(y0, x0):
            y1 = y0 + collision.shape[0]
            x1 = x0 + collision.shape[1]
            if y0 < 0 or x0 < 0 or y1 > height or x1 > width:
                return False
            if np.any(blocked[y0:y1, x0:x1] & stamp):
                return False
            return not np.any(occupied[y0:y1, x0:x1] & collision)

        placed_offset = None
        for ring in offset_rings:
            for dy, dx in ring:
                if fits(base_y + dy, base_x + dx):
                    placed_offset = (dy, dx)
                    break
            if placed_offset is not None:
                break
        if placed_offset is None and allow_global_search:
            placed_offset = _find_free_placement(
                blocked, occupied, collision, stamp, base_y, base_x
            )
        if placed_offset is None:
            # A local-only append is deliberately best-effort: skipping one
            # extra word is much cheaper than scanning the full HD canvas and
            # keeps the latency predictable for large auto-repeat batches.
            failed_word = word
            continue

        dy, dx = placed_offset
        y0 = base_y + dy
        x0 = base_x + dx
        occupied[y0:y0 + stamp.shape[0], x0:x0 + stamp.shape[1]] |= stamp
        if dy or dx:
            shifted_words += 1
            max_applied_shift = max(max_applied_shift, abs(dy), abs(dx))
        accepted.append((word, size, (int(draw_y) + dy, int(draw_x) + dx), orient, color))

    return accepted, {
        "shifted_words": shifted_words,
        "max_shift": max_applied_shift,
        "failed_word": failed_word,
    }


def render_layout_occupancy(layout, mask_shape, font_path):
    h, w = mask_shape
    canvas = Image.new("L", (w, h), 0)
    draw = ImageDraw.Draw(canvas)
    for word, size, (y, x), orient, _color in layout:
        font = get_cached_font(font_path, size)
        if orient:
            font = ImageFont.TransposedFont(font, orientation=orient)
        draw.text((x, y), word, font=font, fill=255)
    return (np.array(canvas) > 0).astype(np.uint8)


def compute_fill_ratio_fast(layout, mask, font_path):
    if not layout:
        return 0.0, None
    occ = render_layout_occupancy(layout, mask.shape, font_path)
    free_area = np.sum(mask == 0)
    if free_area == 0:
        return 0.0, occ
    filled_area = np.sum((mask == 0) & (occ == 1))
    return filled_area / free_area, occ


def compute_coverage_score(occupancy, mask, block_size=8):
    """Shape-aware fill quality: how broadly the ink reaches across the mask.

    Unlike :func:`compute_fill_ratio_fast` (filled pixels / free pixels), this
    rewards reaching *every part* of the mask silhouette -- protrusions, limb
    tips, heart-shaped cusps -- that a centre-out Fermat spiral abandons first
    when words are scarce or font sizes are large.

    The fillable region is tiled into ``block_size`` cells. A cell counts as a
    "region block" when at least 30% of its pixels are free; it is "covered"
    when at least one of those free pixels is inked. The score is the share of
    region blocks that are covered:

        coverage = covered_blocks / region_blocks

    A compact disc packed around the centroid touches only the central blocks,
    so it scores low even at a high pixel fill ratio; a layout that spreads
    into every arm of the mask touches blocks in each arm and scores high.
    Block granularity (not per-pixel weighting) is what makes this robust to
    the mask's geometry: a thin tip is one block whether it is 3px or 30px
    wide, so reaching it is rewarded consistently. ``occupancy`` may be None
    (treated as empty).
    """
    mask_arr = np.asarray(mask)
    if mask_arr.ndim != 2 or mask_arr.size == 0:
        return 0.0
    free = mask_arr == 0
    if not free.any():
        return 0.0

    h, w = mask_arr.shape
    occ = np.asarray(occupancy) if occupancy is not None else None
    if occ is None or occ.shape != mask_arr.shape:
        return 0.0
    inked = (occ == 1) & free

    # Block-aligned tile counts. Trailing partial blocks are merged into the
    # last full block by clamping the end index, so no free pixels are dropped.
    region_blocks = 0
    covered_blocks = 0
    for by in range(0, h, block_size):
        y1 = min(h, by + block_size)
        for bx in range(0, w, block_size):
            x1 = min(w, bx + block_size)
            block_free = free[by:y1, bx:x1]
            free_count = int(block_free.sum())
            if free_count == 0:
                continue
            # A block is a region if a meaningful share of it is fillable;
            # this ignores blocks that only clip a mask corner.
            if free_count < 0.30 * block_free.size:
                continue
            region_blocks += 1
            if np.any(inked[by:y1, bx:x1] & block_free):
                covered_blocks += 1

    if region_blocks == 0:
        return 0.0
    return covered_blocks / region_blocks


def largest_empty_square_size(occupancy, mask):
    """Return the largest fully empty square inside the fillable mask."""
    blocked = (np.asarray(mask) != 0) | (np.asarray(occupancy) != 0)
    if blocked.ndim != 2 or blocked.size == 0:
        return 0

    integral = np.pad(blocked.astype(np.uint32), ((1, 0), (1, 0)))
    integral = integral.cumsum(axis=0).cumsum(axis=1)

    low = 0
    high = min(blocked.shape)
    while low < high:
        size = (low + high + 1) // 2
        window_sums = (
            integral[size:, size:]
            - integral[:-size, size:]
            - integral[size:, :-size]
            + integral[:-size, :-size]
        )
        if np.any(window_sums == 0):
            low = size
        else:
            high = size - 1
    return int(low)