# SPDX-FileCopyrightText: The Docling Contributors
# SPDX-License-Identifier: MIT

import itertools
import logging
import re
import sys
from abc import abstractmethod
from copy import deepcopy
from typing import TYPE_CHECKING, Any, List

if TYPE_CHECKING:
    from transformers import GenerationConfig

_log = logging.getLogger(__name__)


class GenerationStopper:
    """
    Base interface for stopping logic.
    - should_stop(s): True to stop given the current decoded text window.
    - lookback_tokens(): how many tokens should be considered (default: sys.maxsize).
    """

    @abstractmethod
    def should_stop(self, s: str) -> bool:
        pass

    def lookback_tokens(self) -> int:
        return sys.maxsize


def build_generation_config(
    base_config: "GenerationConfig | None",
    *,
    overrides: dict[str, Any] | None = None,
    max_new_tokens: int | None = None,
    use_cache: bool | None = None,
    do_sample: bool | None = None,
    temperature: float | None = None,
    pad_token_id: int | list[int] | None = None,
    eos_token_id: int | list[int] | None = None,
) -> "GenerationConfig":
    from transformers import GenerationConfig

    config = deepcopy(base_config) if base_config is not None else GenerationConfig()

    # Model-level defaults, kept at the lowest precedence so caller overrides
    # still win over them -- matching the previous behavior where these were
    # spread as loose generate() kwargs *before* extra_generation_config.
    if max_new_tokens is not None:
        config.max_new_tokens = max_new_tokens
    if use_cache is not None:
        config.use_cache = use_cache

    # Default pad_token_id to the tokenizer's, falling back to eos_token_id, to
    # silence the transformers pad_token_id warning. This only fills a gap, so a
    # caller override below can still replace it.
    if pad_token_id is not None:
        config.pad_token_id = pad_token_id
    elif config.pad_token_id is None and eos_token_id is not None:
        config.pad_token_id = eos_token_id

    # Caller overrides (e.g. vlm_options.extra_generation_config) win over the
    # model defaults above. update() validates and preserves custom entries such
    # as num_logits_to_keep, unlike a raw setattr loop.
    if overrides:
        config.update(allow_custom_entries=True, **overrides)

    # The explicit sampling decision has the final say, matching the old ordering
    # where do_sample/temperature were set after the override spread.
    if do_sample is not None:
        config.do_sample = do_sample
    if temperature is not None:
        config.temperature = temperature

    return config


class DocTagsRepetitionStopper(GenerationStopper):
    """
    Detects repetitive <tag>...<loc_x><loc_y><loc_w><loc_h>text</tag> blocks,
    but only when repeats are **consecutive** and both tag & inner text are identical.

    Performance:
    - Heavy check runs every N calls (default 32).
    - Only decodes the last LOOKBACK_TOKENS tokens per sequence (default 200).
    """

    def __init__(self, *, N: int = 32, lookback_tokens: int = 200):
        self.N = max(1, int(N))
        self._lookback_tokens = max(1, int(lookback_tokens))
        self._call_count = 0

        # <tag> ... <loc_x><loc_y><loc_w><loc_h> text ... </tag>
        self._PATTERN = re.compile(
            r"""
            <(?P<tag>[a-zA-Z0-9_]+)>\s*
            (?P<prefix>.*?)?
            <loc_(?P<x>\d+)><loc_(?P<y>\d+)><loc_(?P<w>\d+)><loc_(?P<h>\d+)>
            (?P<text>.*?)
            </(?P=tag)>
            """,
            re.DOTALL | re.VERBOSE,
        )

    # --- small helper ---
    def _regular(self, vals: List[int]) -> bool:
        """3+ strictly increasing values with ~regular spacing (±20%)."""
        if len(vals) < 3:
            return False
        diffs = [b - a for a, b in itertools.pairwise(vals)]
        if any(d <= 0 for d in diffs):
            return False
        mean = sum(diffs) / len(diffs)
        tol = 0.2 * mean
        return all(abs(d - mean) <= tol for d in diffs)

    def should_stop(self, s: str) -> bool:
        """
        Trip only on **consecutive** runs (no other matched blocks between) of ≥3 items
        with the same <tag> and identical inner text, where within that run we see:
          - any exact duplicate (x,y,w,h), or
          - stable X/W with regular Y progression, or
          - stable Y/H with regular X progression.
        """
        # Stream matches and evaluate runs on-the-fly to stay compact and fast.
        prev_tag = prev_text = None
        run = []  # list of (x,y,w,h)

        def run_repetitive(boxes: List[tuple]) -> bool:
            if len(boxes) < 3:
                return False
            # duplicates?
            if len(set(boxes)) < len(boxes):
                return True
            xs, ys, ws, hs = zip(*boxes)
            x_stable = all(x == xs[0] for x in xs)
            y_stable = all(y == ys[0] for y in ys)
            w_stable = all(w == ws[0] for w in ws)
            h_stable = all(h == hs[0] for h in hs)
            # horizontal (down the page): X/W stable, Y regular
            if (x_stable or w_stable) and self._regular(list(ys)):
                return True
            # vertical (across): Y/H stable, X regular
            if (y_stable or h_stable) and self._regular(list(xs)):
                return True
            return False

        for m in self._PATTERN.finditer(s):
            tag, text = m.group("tag"), m.group("text")
            box = (
                int(m.group("x")),
                int(m.group("y")),
                int(m.group("w")),
                int(m.group("h")),
            )

            if prev_tag == tag and prev_text == text:
                run.append(box)  # consecutive same-tag+text
            else:
                # evaluate previous run before starting a new one
                if run_repetitive(run):
                    return True
                prev_tag, prev_text = tag, text
                run = [box]

        # check the last run
        return run_repetitive(run)