from __future__ import annotations

from typing import List, Optional

from PIL import Image

from surya.common.blank import is_blank_region
from surya.inference import SuryaInferenceManager, get_default_manager
from surya.inference.parsers import denorm_bbox, parse_layout
from surya.inference.prompts import LAYOUT_JSON_SCHEMA, PROMPT_TYPE_LAYOUT
from surya.inference.schema import BatchInputItem
from surya.layout.label import LAYOUT_PRED_RELABEL, TEXT_LABELS
from surya.layout.schema import LayoutBox, LayoutResult
from surya.logging import get_logger
from surya.settings import settings
from surya.timing import timing_span

logger = get_logger()


class LayoutPredictor:
    """Run LAYOUT_PROMPT on full pages, parse JSON, return LayoutResult per image."""

    def __init__(self, manager: Optional[SuryaInferenceManager] = None):
        self.manager = manager  # If None, get_default_manager() is used at call time
        self._disable_tqdm = settings.DISABLE_TQDM

    @property
    def disable_tqdm(self) -> bool:
        return self._disable_tqdm

    @disable_tqdm.setter
    def disable_tqdm(self, value: bool) -> None:
        self._disable_tqdm = bool(value)

    def to(self, *args, **kwargs):
        # Manager-backed; .to() is a no-op for compatibility with BasePredictor callers.
        return

    def __call__(
        self,
        images: List[Image.Image],
        target_image_sizes: Optional[List[tuple]] = None,
        max_tokens: Optional[int] = None,
    ) -> List[LayoutResult]:
        """Run layout on a batch of images.

        target_image_sizes: optional list of (width, height) tuples — if
        provided, bboxes are denormalized to these sizes instead of each
        input image's size. Useful when layout runs on a low-DPI render but
        you want bboxes in the OCR image's coordinate space.
        """
        if not images:
            return []
        manager = self.manager or get_default_manager()

        max_tokens = max_tokens or settings.SURYA_MAX_TOKENS_LAYOUT
        guided = LAYOUT_JSON_SCHEMA if settings.SURYA_GUIDED_LAYOUT else None
        with timing_span("layout_build_batch", image_count=len(images), max_tokens=max_tokens):
            batch = [
                BatchInputItem(
                    image=img,
                    prompt_type=PROMPT_TYPE_LAYOUT,
                    max_tokens=max_tokens,
                    guided_json=guided,
                )
                for img in images
            ]
        with timing_span("layout_manager_generate", item_count=len(batch)):
            outputs = manager.generate(batch)

        if target_image_sizes is not None and len(target_image_sizes) != len(images):
            raise ValueError("target_image_sizes must match images length")

        with timing_span("layout_parse_outputs", item_count=len(outputs)):
            results: List[LayoutResult] = []
            for idx, (img, out) in enumerate(zip(images, outputs)):
                if target_image_sizes is not None:
                    w, h = target_image_sizes[idx]
                else:
                    w, h = img.size
                page_bbox = [0, 0, float(w), float(h)]
                if out.error or not out.raw:
                    results.append(
                        LayoutResult(
                            bboxes=[], image_bbox=page_bbox, raw=out.raw, error=True
                        )
                    )
                    continue
                try:
                    parsed = parse_layout(out.raw)
                except Exception as e:
                    logger.warning(f"Layout parse failed: {e}; raw[:300]={out.raw[:300]!r}")
                    results.append(
                        LayoutResult(
                            bboxes=[], image_bbox=page_bbox, raw=out.raw, error=True
                        )
                    )
                    continue

                confidence = out.mean_token_prob if out.mean_token_prob is not None else 1.0
                img_w, img_h = img.size
                boxes: List[LayoutBox] = []
                dropped_blank = 0
                for blk in parsed:
                    canon = LAYOUT_PRED_RELABEL.get(blk.label, blk.label)
                    # Drop text-labeled blocks the model hallucinated over an
                    # essentially-blank region (mostly white OR near-uniform
                    # color). Visual blocks (Picture / Figure / Table / etc.)
                    # are allowed to be uniform — that's normal content.
                    if canon in TEXT_LABELS:
                        img_bbox = denorm_bbox(
                            blk.bbox, img_w, img_h, scale=settings.BBOX_SCALE
                        )
                        x0, y0, x1, y1 = (max(0, int(v)) for v in img_bbox)
                        if x1 > x0 and y1 > y0:
                            if is_blank_region(img.crop((x0, y0, x1, y1))):
                                dropped_blank += 1
                                continue
                    pixel_bbox = denorm_bbox(blk.bbox, w, h, scale=settings.BBOX_SCALE)
                    boxes.append(
                        LayoutBox(
                            polygon=list(pixel_bbox),
                            label=canon,
                            raw_label=blk.label,
                            position=len(boxes),
                            count=blk.count,
                            confidence=confidence,
                        )
                    )
                if dropped_blank:
                    logger.info(
                        f"dropped {dropped_blank} text-labeled layout block(s) over "
                        f"blank/uniform regions"
                    )
                results.append(
                    LayoutResult(
                        bboxes=boxes, image_bbox=page_bbox, raw=out.raw, error=False
                    )
                )
        return results
