from __future__ import annotations

import re
from dataclasses import dataclass

OBJECT_RE = re.compile(
    r"<\|object_ref_start\|>(.*?)<\|object_ref_end\|>"
    r"<\|box_start\|>(.*?)<\|box_end\|>",
    re.DOTALL,
)
COORD_RE = re.compile(r"<([0-9]{1,3})>")
POINT_ROUTES = {"dense_point", "refer_point", "gui", "referring_point_in_mask", "dense_point_in_mask"}


@dataclass(frozen=True)
class Primitive:
    label: str
    coordinates: tuple[int, ...]


@dataclass(frozen=True)
class ParseResult:
    primitives: tuple[Primitive, ...]
    errors: tuple[str, ...]

    @property
    def valid(self) -> bool:
        return bool(self.primitives) and not self.errors


def parse_gam(text: str, route: str) -> ParseResult:
    """Strictly parse canonical GAM special-token output.

    Reasoning before the first object block is ignored, but every detected
    object block and every coordinate must satisfy the released GAM contract.
    """

    arity = 2 if route in POINT_ROUTES else 4
    primitives: list[Primitive] = []
    errors: list[str] = []
    matches = list(OBJECT_RE.finditer(text or ""))
    if not matches:
        return ParseResult((), ("Canonical GAM object/box wrapper not found",))

    for block_index, match in enumerate(matches, start=1):
        label = match.group(1).strip()
        payload = match.group(2).strip()
        if not label:
            errors.append(f"Object {block_index} has an empty label")
        tuples = [item.strip() for item in payload.split(",") if item.strip()]
        if not tuples:
            errors.append(f"Object {block_index} has no coordinates")
            continue
        for tuple_index, item in enumerate(tuples, start=1):
            values = tuple(int(v) for v in COORD_RE.findall(item))
            rebuilt = "".join(f"<{v}>" for v in values)
            if rebuilt != item:
                errors.append(f"Coordinate {block_index}.{tuple_index} contains non-canonical content")
                continue
            if len(values) != arity:
                errors.append(
                    f"Coordinate {block_index}.{tuple_index} has arity={len(values)}, expected {arity}"
                )
                continue
            if any(v < 0 or v > 999 for v in values):
                errors.append(f"Coordinate {block_index}.{tuple_index} is outside 0..999")
                continue
            if arity == 4 and (values[0] >= values[2] or values[1] >= values[3]):
                errors.append(f"Box {block_index}.{tuple_index} is degenerate or inverted")
                continue
            primitives.append(Primitive(label=label, coordinates=values))

    return ParseResult(tuple(primitives), tuple(errors))


def parse_reference_boxes(prompt: str) -> tuple[Primitive, ...]:
    if "Given reference boxes" not in prompt:
        return ()
    match = re.search(r"<\|box_start\|>(.*?)<\|box_end\|>", prompt, re.DOTALL)
    if not match:
        return ()
    result: list[Primitive] = []
    for item in match.group(1).split(","):
        values = tuple(int(v) for v in COORD_RE.findall(item))
        if len(values) == 4 and values[0] < values[2] and values[1] < values[3]:
            result.append(Primitive("reference", values))
    return tuple(result)