from __future__ import annotations

from dataclasses import replace
from pathlib import Path
from struct import unpack

from core.models import Bounds, Scene, SceneElement


def build_scene(
    *,
    screen_width: int,
    screen_height: int,
    ui_elements: list[SceneElement] | None = None,
    ocr_elements: list[SceneElement] | None = None,
    iou_threshold: float = 0.5,
) -> Scene:
    ui_elements = ui_elements or []
    ocr_elements = ocr_elements or []
    merged: list[SceneElement] = []
    used_ocr: set[int] = set()

    for ui_index, ui_element in enumerate(ui_elements):
        best_index: int | None = None
        best_iou = 0.0
        for ocr_index, ocr_element in enumerate(ocr_elements):
            if ocr_index in used_ocr:
                continue
            score = bbox_iou(ui_element.bounds, ocr_element.bounds)
            if score > best_iou:
                best_iou = score
                best_index = ocr_index

        element = _with_id(ui_element, f"ui-{ui_index:03d}")
        if best_index is not None and best_iou >= iou_threshold:
            used_ocr.add(best_index)
            ocr_element = ocr_elements[best_index]
            element = replace(
                element,
                text=element.text or ocr_element.text,
                confidence=_best_confidence(element.confidence, ocr_element.confidence),
            )
        merged.append(element)

    for ocr_index, ocr_element in enumerate(ocr_elements):
        if ocr_index in used_ocr:
            continue
        merged.append(_with_id(ocr_element, f"ocr-{ocr_index:03d}"))

    return Scene(
        width=screen_width,
        height=screen_height,
        elements=merged,
        ocr_elements=list(ocr_elements),
    )


def bbox_iou(first: Bounds, second: Bounds) -> float:
    x_left = max(first.x, second.x)
    y_top = max(first.y, second.y)
    x_right = min(first.right, second.right)
    y_bottom = min(first.bottom, second.bottom)
    if x_right <= x_left or y_bottom <= y_top:
        return 0.0

    intersection = (x_right - x_left) * (y_bottom - y_top)
    first_area = first.width * first.height
    second_area = second.width * second.height
    union = first_area + second_area - intersection
    if union <= 0:
        return 0.0
    return intersection / union


def infer_png_size(image: bytes | str | Path | None) -> tuple[int, int]:
    if image is None:
        return (0, 0)
    data = Path(image).read_bytes() if not isinstance(image, bytes) else image
    if len(data) >= 24 and data[:8] == b"\x89PNG\r\n\x1a\n":
        width, height = unpack(">II", data[16:24])
        return (int(width), int(height))
    return (0, 0)


def _with_id(element: SceneElement, fallback_id: str) -> SceneElement:
    if element.id:
        return element
    return replace(element, id=fallback_id)


def _best_confidence(first: float | None, second: float | None) -> float | None:
    values = [value for value in (first, second) if value is not None]
    if not values:
        return None
    return max(values)