from __future__ import annotations

from collections.abc import Iterable
from functools import lru_cache
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont

from .protocol import Primitive

PURPLE = (120, 78, 188)
GREEN = (145, 195, 49)
PINK = (214, 120, 172)
CYAN = (55, 171, 183)
PALETTE = (PURPLE, GREEN, PINK, CYAN)


@lru_cache(maxsize=32)
def _load_font(size: int, multilingual: bool) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    assets = Path(__file__).resolve().parents[1] / "assets"
    if multilingual:
        # Pillow does not perform the browser's font fallback. Bundle CJK glyphs
        # so Chinese OCR labels render identically on macOS and the Linux Space.
        return ImageFont.truetype(str(assets / "NotoSansCJKsc-Regular.otf"), size=size)
    candidates = (
        str(assets / "Manrope-variable.ttf"),
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
        "/usr/share/fonts/truetype/liberation2/LiberationSans-Regular.ttf",
    )
    for candidate in candidates:
        if Path(candidate).is_file():
            font = ImageFont.truetype(candidate, size=size)
            if Path(candidate).name == "Manrope-variable.ttf":
                font.set_variation_by_axes([600])
            return font
    return ImageFont.load_default()


def _font(size: int, text: str = "") -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
    return _load_font(size, not text.isascii())


def _xy(values: tuple[int, ...], width: int, height: int) -> tuple[int, ...]:
    if len(values) == 2:
        return round(values[0] * (width - 1) / 999), round(values[1] * (height - 1) / 999)
    return (
        round(values[0] * (width - 1) / 999),
        round(values[1] * (height - 1) / 999),
        round(values[2] * (width - 1) / 999),
        round(values[3] * (height - 1) / 999),
    )


def _intersects(a, b, gap=2):
    return a[0] < b[2]+gap and b[0] < a[2]+gap and a[1] < b[3]+gap and b[1] < a[3]+gap


def layout_box_labels(primitives, width, height, font, padding=4):
    """One collision-free badge per exact category; returned boxes stay untouched."""
    groups = {}
    for primitive in primitives:
        if len(primitive.coordinates) == 4:
            groups.setdefault(primitive.label, []).append(_xy(primitive.coordinates, width, height))
    area = lambda b: max(0,b[2]-b[0])*max(0,b[3]-b[1])
    small = [b for boxes in groups.values() for b in boxes if area(b)<width*height*.01]
    measure = ImageDraw.Draw(Image.new("RGB", (1, 1)))
    occupied, placed = [], []
    # Give small targets first choice of nearby space. Large regions have more
    # candidate anchors. Ordering affects label placement only, never predictions.
    for label, boxes in sorted(groups.items(), key=lambda item: max(map(area,item[1]))):
        label_font = _font(getattr(font, "size", 13), label) if not label.isascii() else font
        text = label if len(label)<=38 else label[:37]+"…"
        while text:
            bounds = measure.textbbox((0,0),text,font=label_font)
            tw, th = bounds[2]-bounds[0]+2*padding, bounds[3]-bounds[1]+2*padding
            if tw<=width-2 and th<=height-2: break
            text = text[:-2]+"…" if len(text)>2 else ""
        if not text: continue
        candidates=[]
        for target in sorted(boxes,key=area,reverse=True):
            x1,y1,x2,y2=target
            for x,y in [(x1,y1-th-2),(x1,y1+2),(x1,y2+2),(x2-tw,y1-th-2),
                        (x2-tw,y2-th-2),(x2+2,y1),(x1-tw-2,y1)]:
                x=max(1,min(x,width-tw-1));y=max(1,min(y,height-th-1))
                candidates.append(((x,y,x+tw,y+th),target))
        chosen=next(((b,t) for b,t in candidates if not any(_intersects(b,o) for o in occupied+small)),None)
        if chosen is None:
            # Rare crowded anchors use the nearest free area with a leader line.
            # Never put a badge on another badge or cover a small prediction.
            free=[]
            for y in range(1,max(2,height-th),max(4,th//2)):
                for x in range(1,max(2,width-tw),max(6,tw//4)):
                    b=(x,y,x+tw,y+th)
                    if any(_intersects(b,o) for o in occupied+small): continue
                    for t in boxes:
                        dx=max(t[0]-b[2],b[0]-t[2],0);dy=max(t[1]-b[3],b[1]-t[3],0)
                        # Prefer top-left anchors among equally near locations.
                        score=dx*dx+dy*dy+.01*((x-t[0])**2+(y-t[1])**2)
                        free.append((score,b,t))
            if free:
                _,b,t=min(free,key=lambda item:item[0]);chosen=(b,t)
        if chosen:
            b,target=chosen;occupied.append(b)
            placed.append(dict(label=label,text=text,box=b,target=target,text_bounds=bounds,font=label_font))
    return placed


def draw_primitives(
    image: Image.Image,
    primitives: Iterable[Primitive],
    references: Iterable[Primitive] = (),
    title: str | None = None,
) -> Image.Image:
    canvas = image.convert("RGBA")
    overlay = Image.new("RGBA", canvas.size, (0, 0, 0, 0))
    draw = ImageDraw.Draw(overlay)
    width, height = canvas.size
    primitives = list(primitives)
    dense = len(primitives) > 12
    # Fine boundaries keep a dense scene readable; short corners make small
    # boxes discoverable without filling, moving, or removing any prediction.
    line = max(1 if dense else 2, round(min(width, height) / (420 if dense else 260)))
    radius = max(6, line * 3)
    font = _font(max(12 if dense else 13, round(min(width, height) / 48)))
    label_colors: dict[str, tuple[int, int, int]] = {}

    for ref in references:
        if len(ref.coordinates) != 4:
            continue
        box = _xy(ref.coordinates, width, height)
        draw.rectangle(box, outline=(*GREEN, 255), width=line * 2)
        x1, y1, x2, y2 = box
        dash = max(5, line * 3)
        for x in range(x1, x2, dash * 2):
            draw.line((x, y1, min(x + dash, x2), y1), fill=(*GREEN, 255), width=line)
            draw.line((x, y2, min(x + dash, x2), y2), fill=(*GREEN, 255), width=line)

    # Opaque RGBA writes replace earlier pixels in a Pillow layer. In particular,
    # a later lake/cloud fill used to erase every small object drawn inside it.
    # Keep region interiors untouched; geometry and labels are separate passes.
    for primitive in primitives:
        color = label_colors.setdefault(primitive.label, PALETTE[len(label_colors) % len(PALETTE)])
        coords = _xy(primitive.coordinates, width, height)
        if len(coords) == 2:
            x, y = coords
            draw.ellipse(
                (x - radius, y - radius, x + radius, y + radius),
                outline=(*color, 255),
                width=line,
            )
            draw.line((x - radius * 2, y, x + radius * 2, y), fill=(*color, 255), width=line)
            draw.line((x, y - radius * 2, x, y + radius * 2), fill=(*color, 255), width=line)
        else:
            x1, y1, x2, y2 = coords
            draw.rectangle(coords, outline=(*color, 255), width=line)
            if dense:
                corner = min(max(4, round(min(width, height) / 80)), (x2-x1)//3, (y2-y1)//3)
                if corner > 0:
                    for x, y, dx, dy in ((x1,y1,1,1),(x2,y1,-1,1),(x1,y2,1,-1),(x2,y2,-1,-1)):
                        draw.line((x+dx*corner,y,x,y,x,y+dy*corner), fill=(*color,255), width=max(2,line))

    # Labels remain visible in dense scenes, once per category across all boxes.
    padding=max(4,line*2)
    for badge in layout_box_labels(primitives,width,height,font,padding):
        text,box,target=badge["text"],badge["box"],badge["target"]
        color=label_colors[badge["label"]]
        tx,ty=box[:2];tb=badge["text_bounds"]
        # Attach displaced badges to their actual prediction with a fine leader.
        if not _intersects(box,target,gap=4):
            cx,cy=(box[0]+box[2])/2,(box[1]+box[3])/2
            ex=max(target[0],min(cx,target[2]));ey=max(target[1],min(cy,target[3]))
            sx=max(box[0],min(ex,box[2]));sy=max(box[1],min(ey,box[3]))
            draw.line((sx,sy,ex,ey),fill=(*color,230),width=1)
        light=tuple(round(c*.10+255*.90) for c in color)
        draw.rounded_rectangle(box,radius=3,fill=(*light,248),outline=(*color,255),width=1)
        draw.text((tx+padding-tb[0],ty+padding-tb[1]),text,font=badge["font"],fill=(38,30,49,255))

    if title:
        title_font = _font(getattr(font, "size", 13), title)
        tb = draw.textbbox((0, 0), title, font=title_font)
        draw.rounded_rectangle(
            (10, 10, min(width - 10, tb[2] + 28), tb[3] + 24),
            radius=10,
            fill=(255, 255, 255, 225),
            outline=(*PURPLE, 255),
            width=line,
        )
        draw.text((18, 16), title, font=title_font, fill=(*PURPLE, 255))
    return Image.alpha_composite(canvas, overlay).convert("RGB")