from __future__ import annotations

import re
from dataclasses import dataclass

LOC_TOKEN_RE = re.compile(r"<LOC_(\d+)>")
SPECIAL_TOKEN_RE = re.compile(
    r"<\|(?:object_ref_start|object_ref_end|box_start|box_end|im_start|im_end)\|>"
    r"|</?ref>|</?box>|</?quad>"
)
BOX_PAIR_RE = re.compile(
    r"(.+?)<\|box_start\|>\s*\(?\s*(\d+(?:\.\d+)?)\s*,\s*(\d+(?:\.\d+)?)\s*\)?"
    r"\s*,?\s*\(?\s*(\d+(?:\.\d+)?)\s*,\s*(\d+(?:\.\d+)?)\s*\)?\s*<\|box_end\|>",
    flags=re.DOTALL,
)


@dataclass(frozen=True)
class TextSpan:
    text: str
    box: tuple[float, float, float, float] | None = None


def _normalize_box(
    values: list[float],
    image_width: int,
    image_height: int,
    normalized_max: float = 1000.0,
) -> tuple[float, float, float, float]:
    if len(values) >= 8:
        xs = values[0:8:2]
        ys = values[1:8:2]
    else:
        xs = values[0::2]
        ys = values[1::2]
    scale_x = image_width / normalized_max
    scale_y = image_height / normalized_max
    if max(xs) <= 1.5 and max(ys) <= 1.5:
        scale_x = image_width
        scale_y = image_height
    elif max(xs) > image_width * 1.2 or max(ys) > image_height * 1.2:
        scale_x = image_width / max(normalized_max, max(xs) or 1.0)
        scale_y = image_height / max(normalized_max, max(ys) or 1.0)
    x0 = max(0.0, min(xs) * scale_x)
    y0 = max(0.0, min(ys) * scale_y)
    x1 = min(float(image_width), max(xs) * scale_x)
    y1 = min(float(image_height), max(ys) * scale_y)
    if x1 <= x0:
        x1 = min(float(image_width), x0 + 8.0)
    if y1 <= y0:
        y1 = min(float(image_height), y0 + 8.0)
    return x0, y0, x1, y1


def parse_loc_spans(raw: str, image_width: int, image_height: int) -> list[TextSpan]:
    if not raw or not LOC_TOKEN_RE.search(raw):
        return []

    spans: list[TextSpan] = []
    buffer = ""
    coords: list[float] = []
    last = 0
    for match in LOC_TOKEN_RE.finditer(raw):
        buffer += raw[last:match.start()]
        coords.append(float(match.group(1)))
        last = match.end()
        if len(coords) == 8:
            text = SPECIAL_TOKEN_RE.sub("", buffer).strip(" \n\t:;,-")
            if text:
                spans.append(
                    TextSpan(text=text, box=_normalize_box(coords, image_width, image_height))
                )
            buffer = ""
            coords = []
    trailing = SPECIAL_TOKEN_RE.sub("", raw[last:]).strip()
    if trailing and not spans:
        return []
    return spans


def parse_box_tag_spans(raw: str, image_width: int, image_height: int) -> list[TextSpan]:
    spans: list[TextSpan] = []
    for match in BOX_PAIR_RE.finditer(raw or ""):
        text = SPECIAL_TOKEN_RE.sub("", match.group(1)).strip()
        if not text:
            continue
        coords = [float(match.group(i)) for i in range(2, 6)]
        spans.append(TextSpan(text=text, box=_normalize_box(coords, image_width, image_height)))
    return spans


def parse_spans(raw: str, image_width: int, image_height: int) -> list[TextSpan]:
    loc_spans = parse_loc_spans(raw, image_width, image_height)
    if loc_spans:
        return loc_spans
    box_spans = parse_box_tag_spans(raw, image_width, image_height)
    if box_spans:
        return box_spans
    return []


def strip_special_tokens(raw: str) -> str:
    text = SPECIAL_TOKEN_RE.sub("", raw or "")
    text = LOC_TOKEN_RE.sub("", text)
    return text


def boxes_iou(left: tuple[float, float, float, float], right: tuple[float, float, float, float]) -> float:
    lx0, ly0, lx1, ly1 = left
    rx0, ry0, rx1, ry1 = right
    ix0, iy0 = max(lx0, rx0), max(ly0, ry0)
    ix1, iy1 = min(lx1, rx1), min(ly1, ry1)
    inter = max(0.0, ix1 - ix0) * max(0.0, iy1 - iy0)
    if inter <= 0:
        return 0.0
    area_l = max(1.0, (lx1 - lx0) * (ly1 - ly0))
    area_r = max(1.0, (rx1 - rx0) * (ry1 - ry0))
    return inter / (area_l + area_r - inter)


def spans_to_text(spans: list[TextSpan], fallback: str) -> str:
    if not spans:
        return fallback
    ordered = sorted(
        spans,
        key=lambda span: (
            round((span.box[1] if span.box else 0.0) / 12.0),
            span.box[0] if span.box else 0.0,
        ),
    )
    return "\n".join(span.text for span in ordered if span.text)