"""Utils for the ScreenQA benchmark. Referenced from https://github.com/X-LANCE/WebSRC-Baseline/blob/master/src/utils_evaluate.py#L79-L131"""

import collections
import logging
import re
import string
from math import sqrt
from typing import Any

import numpy as np
from transformers.models.auto.processing_auto import AutoProcessor

from screensuite.benchmarks.model import BoundingBox, NormalizedClickCoordinates

logger = logging.getLogger(__name__)


def normalize_answer(s: str) -> str:
    """Lower text and remove punctuation, articles and extra whitespace."""

    def remove_articles(text):
        regex = re.compile(r"\b(a|an|the)\b", re.UNICODE)
        return re.sub(regex, " ", text)

    def white_space_fix(text):
        return " ".join(text.split())

    def remove_punc(text):
        exclude = set(string.punctuation)
        return "".join(ch for ch in text if ch not in exclude)

    def lower(text):
        return text.lower()

    return white_space_fix(remove_articles(remove_punc(lower(s))))


def get_tokens(s: str) -> list[str]:
    r"""
    Get the word list in the input.
    """
    if not s:
        return []
    return normalize_answer(s).split()


def compute_exact(a_gold: str, a_pred: str) -> int:
    r"""
    Calculate the exact match.
    """
    return int(normalize_answer(a_gold) == normalize_answer(a_pred))


def compute_f1(a_pred: str, a_gold: str) -> float:
    r"""
    Calculate the f1 score.
    """
    gold_toks = get_tokens(a_gold)
    pred_toks = get_tokens(a_pred)
    common = collections.Counter(gold_toks) & collections.Counter(pred_toks)
    num_same = sum(common.values())
    if len(gold_toks) == 0 or len(pred_toks) == 0:
        # If either is no-answer, then F1 is 1 if they agree, 0 otherwise
        return int(gold_toks == pred_toks)
    if num_same == 0:
        return 0
    precision = 1.0 * num_same / len(pred_toks)
    recall = 1.0 * num_same / len(gold_toks)
    f1 = (2 * precision * recall) / (precision + recall)
    return f1


def parse_click_action(pred: str, image_width: int = 1, image_height: int = 1) -> NormalizedClickCoordinates | None:
    pred = pred.lower()

    # Extract everything within click() parentheses
    click_match = re.search(r"click\((.*?)\)", pred)
    if not click_match:
        logger.warning(f"No click action found in: {pred}")
        return None

    click_content = click_match.group(1)

    # Find all numbers (integers or floats) in the click content
    numbers = re.findall(r"([\d.]+)", click_content)

    if len(numbers) >= 2:
        try:
            x, y = float(numbers[0]), float(numbers[1])
            return NormalizedClickCoordinates.from_xy(xy=(x, y), image_width=image_width, image_height=image_height)
        except Exception:
            logger.warning(
                f"Error parsing click coordinates from: {pred} with image width {image_width} and image height {image_height}"
            )
            return None

    logger.warning(f"Could not find two numeric coordinates in: {pred}")
    return None


def parse_bounding_box_action(pred: str, image_width: int = 1, image_height: int = 1) -> BoundingBox | None:
    pred = pred.lower()
    # format: boundingbox(x1,y1,x2,y2)
    matches = re.findall(r"boundingbox\(\s*([\d.]+)\s*,\s*([\d.]+)\s*,\s*([\d.]+)\s*,\s*([\d.]+)\s*\)", pred)

    try:
        if matches and len(matches) == 1:
            x1, y1, x2, y2 = (float(m) for m in matches[0])
            return BoundingBox.from_xyxy(xyxy_bbox=(x1, y1, x2, y2), image_width=image_width, image_height=image_height)
    except Exception as e:
        logger.warning(f"Error parsing bounding box: {pred}")
        return None
    logger.warning(f"No bounding box or too many bounding boxes found: {pred}")
    return None


def parse_type_action(pred: str) -> str | None:
    pred = pred.lower()

    # First, extract everything inside type(...)
    type_match = re.search(r"type\(([^)]+)\)", pred)
    if not type_match:
        return None

    content = type_match.group(1)

    # Look for quoted content within that
    quote_match = re.search(r"[\"'](.+?)[\"']", content)
    if quote_match:
        return quote_match.group(1).strip()

    # If no quotes found, return the content as-is (trimmed)
    return content.strip()


def evaluate_click_in_bounding_box(
    pred: str, target_bbox: BoundingBox, image_width: int = 1, image_height: int = 1
) -> float:
    click_position = parse_click_action(pred, image_width, image_height)
    if click_position is None:
        return 0.0

    pred_x, pred_y = click_position.to_xy()
    gt_x1, gt_y1, gt_x2, gt_y2 = target_bbox.to_xyxy()

    is_inside = (gt_x1 <= pred_x <= gt_x2) and (gt_y1 <= pred_y <= gt_y2)

    return 1.0 if is_inside else 0.0


def evaluate_click_distance(
    pred: str, target_coordinates: NormalizedClickCoordinates, image_width: int = 1, image_height: int = 1
) -> float:
    click_position = parse_click_action(pred, image_width, image_height)
    if click_position is None:
        return 0.0

    pred_x, pred_y = click_position.to_xy()
    gt_x, gt_y = target_coordinates.to_xy()

    distance = sqrt(((pred_x - gt_x) ** 2 + (pred_y - gt_y) ** 2) / 2)

    return max(0, 1 - distance * 2)  # Should be 1 for perfect click, 0 for max distance


# IoU evaluation
def evaluate_bounding_box_action(
    pred: str, target_bbox: BoundingBox, image_width: int = 1, image_height: int = 1
) -> float:
    bounding_box = parse_bounding_box_action(pred, image_width, image_height)
    if bounding_box is None:
        return 0.0

    pred_x1, pred_y1, pred_x2, pred_y2 = bounding_box.to_xyxy()
    gt_x1, gt_y1, gt_x2, gt_y2 = target_bbox.to_xyxy()

    x_left = max(pred_x1, gt_x1)
    y_top = max(pred_y1, gt_y1)
    x_right = min(pred_x2, gt_x2)
    y_bottom = min(pred_y2, gt_y2)

    if x_right < x_left or y_bottom < y_top:
        return 0.0

    intersection_area = (x_right - x_left) * (y_bottom - y_top)

    # Calculate union area
    pred_area = (pred_x2 - pred_x1) * (pred_y2 - pred_y1)
    gt_area = (gt_x2 - gt_x1) * (gt_y2 - gt_y1)
    union_area = pred_area + gt_area - intersection_area

    # Calculate IoU
    accuracy = intersection_area / union_area if union_area > 0 else 0.0
    return accuracy


def evaluate_type_action(pred: str, gold: str) -> float:
    text = parse_type_action(pred)
    if text is None:
        return 0.0
    lower_text = text.lower()
    lower_gold = gold.lower()
    return float(lower_text == lower_gold)


def check_messages_below_max_context_length(messages: list[dict[str, Any]], max_length: int = 32000):
    processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-72B-Instruct")
    input_tokens = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
    )
    token_count = len(input_tokens["input_ids"][0])

    if token_count > max_length:
        raise ValueError(f"Token count ({token_count}) exceeds the {max_length} limit")
    else:
        print(f"All is right on the token count: {token_count} tokens")


def bootstrap_confidence_interval(values, *, n_boot: int = 1000, ci: int = 95, seed: int = 42):
    """
    Compute the bootstrap confidence interval of the mean of a numeric variable.

    Args:
        values (1-D array-like of numeric values)
        n_boot (int): number of bootstrap resamples
        ci (int): confidence level, e.g. 95 for 95 % CI
        seed (int): seed for reproducibility
    """
    assert 0 < ci < 100, "CI must be between 0 and 100"
    values = np.asarray(values, dtype=float)
    n = values.size
    # Handle edge case where there are no samples
    if n == 0:
        return np.nan, (np.nan, np.nan)

    rng = np.random.default_rng(seed)

    # Draw `n_boot` resamples of size `n`
    resample_idxs = rng.choice(n, size=(n_boot, n), replace=True)
    resample_means = values[resample_idxs].mean(axis=1)

    # Percentile-based CI
    alpha = (100 - ci) / 2  # e.g. 2.5 for a 95 % CI
    lower, upper = np.percentile(resample_means, [alpha, 100 - alpha])

    return values.mean(), (lower, upper)
