import logging
import re
import string
from collections import defaultdict
from typing import Any, Tuple, cast

from inspect_ai.scorer import (
    Metric,
    SampleScore,
    Score,
    Scorer,
    Target,
    ValueToFloat,
    metric,
    scorer,
    value_to_float,
)
from inspect_ai.solver import TaskState

from inspect_evals.mind2web.dataset import ActionType, Answer, Mind2WebMetadata

logger = logging.getLogger(__name__)


def extract_answer(completion: str) -> Answer:
    """Extract the answer from the model's completion.

    This function parses the model's completion to extract the letter choice,
    action type, and value (if applicable). It handles various formats and
    patterns that might appear in the completion.

    Args:
        completion: The raw text completion from the model.

    Returns:
        An Answer object containing the extracted letter, action, and value.
    """
    # Clean up the completion by normalizing whitespace
    # First normalize all newlines and remove leading/trailing whitespace
    completion = completion.strip()
    lines = [line.strip() for line in completion.split("\n")]
    # Remove empty lines and join
    completion = "\n".join(line for line in lines if line)

    # Try to extract letter with flexible pattern matching
    # This pattern matches:
    # - The letter after "Answer:" with optional whitespace
    # - The letter can be surrounded by quotes, brackets, parentheses
    # - The letter can be preceded by "Option" or other text
    # - The letter can be followed by a period or other punctuation
    letter_pattern = (
        r"Answer:\s*(?:[\"'\[\(]|\bOption\b\s*)*([A-F])(?:\.|\)|\]|\"|'|\s|$)"
    )
    letter_match = re.search(letter_pattern, completion, re.IGNORECASE)

    # If no letter found, try a more aggressive search
    if not letter_match:
        # Look for any standalone A-F letter that might be an answer
        letter_match = re.search(
            r"(?:^|\s)([A-F])(?:\.|\)|\]|\"|'|\s|$)(?!\S)", completion, re.IGNORECASE
        )

    # If still no letter found, return empty Answer to indicate invalid answer
    if not letter_match:
        logger.warning(f"No letter match found in completion: {completion}")
        return Answer(letter="", action=None, value=None)

    letter = letter_match.group(1).upper()
    logger.debug(f"Extracted letter: {letter}")

    # For "A" (None of the above), we don't need action/value
    if letter == "A":
        return Answer(letter=letter, action=None, value=None)

    action_match = re.search(r"Action:\s*([^\n]+)", completion, re.IGNORECASE)
    value_match = re.search(r"Value:\s*([^\n]+)", completion, re.IGNORECASE)

    action = None
    value = None

    # Add action if found
    if action_match:
        action_text = action_match.group(1).strip().upper()
        # Convert to ActionType if valid
        if action_text in ["CLICK", "SELECT", "TYPE"]:
            action = cast(ActionType, action_text)
            logger.debug(f"Extracted action: {action}")

            # Add value if found and action is TYPE or SELECT
            if value_match and (action in ["TYPE", "SELECT"]):
                value = value_match.group(1).strip()
                logger.debug(f"Extracted value: {value}")
            # For CLICK, we don't include value
            elif action == "CLICK":
                logger.debug("CLICK action - skipping value")
            # For other actions, include value if present
            elif value_match:
                value = value_match.group(1).strip()
                logger.debug(f"Extracted value: {value}")

    result = Answer(letter=letter, action=action, value=value)
    logger.debug(f"Final extracted answer: {result}")
    return result


@metric
def step_success_rate(to_float: ValueToFloat = value_to_float()) -> Metric:
    """Compute proportion of steps that are fully successful.

    A step is considered successful if both element selection (element_acc) and
    operation prediction (action_f1) are completely correct (score of 1.0).
    This metric measures the model's ability to correctly predict both which
    element to interact with and what action to perform on it.

    Args:
        to_float: Function for mapping values to floats for computing metrics.
            Defaults to the standard value_to_float function.

    Returns:
        A Metric function that calculates the step success rate.
    """

    def metric(scores: list[SampleScore]) -> float:
        all_values = []
        for item in scores:
            # A step is successful if both element and operation are correct
            metadata: dict[str, Any] = item.score.metadata or {}
            element_acc = metadata["element_acc"]
            action_f1 = metadata["action_f1"]
            step_success = float(element_acc == 1.0 and action_f1 == 1.0)
            all_values.append(step_success)
        return sum(all_values) / len(all_values) if all_values else 0.0

    return metric


@metric
def task_success_rate(to_float: ValueToFloat = value_to_float()) -> Metric:
    """Compute proportion of tasks that are fully successful.

    A task is considered successful only if all its steps are successful, meaning
    every step has both element accuracy and operation F1 scores of 1.0. This is
    a stringent metric that measures the model's ability to complete entire tasks
    without any errors.

    Args:
        to_float: Function for mapping values to floats for computing metrics.
            Defaults to the standard value_to_float function.

    Returns:
        A Metric function that calculates the task success rate.
    """

    def metric(scores: list[SampleScore]) -> float:
        # Group scores by task_id using defaultdict
        tasks = defaultdict(list)
        for item in scores:
            metadata: dict[str, Any] = item.score.metadata or {}
            task_id = metadata.get("task_id")
            if not task_id:
                continue
            # A step is successful if both element and operation are correct
            element_acc = metadata["element_acc"]
            action_f1 = metadata["action_f1"]
            step_success = element_acc == 1.0 and action_f1 == 1.0
            tasks[task_id].append(step_success)

        # Calculate success rate - a task is successful if all its steps are successful
        successful_tasks = sum(all(steps) for steps in tasks.values())
        total_tasks = len(tasks)
        return successful_tasks / total_tasks if total_tasks > 0 else 0.0

    return metric


@metric
def element_accuracy() -> Metric:
    """Calculate the average element selection accuracy across all samples.

    This metric measures how often the model correctly identifies which
    element to interact with on the webpage.

    Returns:
        A Metric function that calculates the average element accuracy.
    """
    return avg_score(field="element_acc")


@metric
def action_f1() -> Metric:
    """Calculate the average action F1 score across all samples.

    This metric measures how accurately the model predicts the correct
    operation (CLICK, TYPE, SELECT) and any associated values.

    Returns:
        A Metric function that calculates the average action F1 score.
    """
    return avg_score(field="action_f1")


@metric
def avg_score(field: str, to_float: ValueToFloat = value_to_float()) -> Metric:
    """Compute the average score for a specific metadata field.

    Args:
        field: The metadata field to average (e.g., "element_acc", "action_f1").
        to_float: Function for mapping values to floats for computing metrics.
            Defaults to the standard value_to_float function.

    Returns:
        A Metric function that calculates the average score for the specified field.
    """

    def metric(scores: list[SampleScore]) -> float:
        total = 0.0
        for item in scores:
            if not isinstance(item.score.metadata, dict):
                continue
            if item.score.metadata[field] is None:
                continue
            total += to_float(item.score.metadata[field])
        return total / float(len(scores))

    return metric


def calculate_text_similarity(pred: str, label: str) -> float:
    """Calculate F1 score between predicted and target action strings.

    Args:
        pred: Predicted action string
        label: Target action string

    Returns:
        float: F1 score between 0 and 1
    """
    pred_words = set(pred.strip().split())
    label_words = set(label.strip().split())
    # Remove punctuation
    pred_words = set([x for x in pred_words if x not in string.punctuation])
    label_words = set([x for x in label_words if x not in string.punctuation])

    if len(pred_words) == 0 and len(label_words) == 0:
        return 1.0
    if len(pred_words) == 0 or len(label_words) == 0:
        return 0.0

    tp = len(pred_words & label_words)
    fp = len(pred_words - label_words)
    fn = len(label_words - pred_words)
    precision = tp / (tp + fp)
    recall = tp / (tp + fn)
    if precision == 0 or recall == 0:
        return 0.0
    f1 = 2 * precision * recall / (precision + recall)
    return f1


def score_answer(pred_answer: Answer, target_answer: Answer) -> Tuple[float, float]:
    """Score the predicted answer against the target answer.

    Args:
        pred_answer: The extracted answer from the completion
        target_answer: The target answer from metadata

    Returns:
        Tuple[float, float]: (element_accuracy, action_f1_score)
    """
    # Handle empty or None cases
    if not pred_answer.letter or not target_answer.letter:
        return (0.0, 0.0)

    # Get letters and convert to uppercase for comparison
    pred_letter = pred_answer.letter.upper()
    target_letter = target_answer.letter.upper()
    element_acc = float(pred_letter == target_letter)

    # For answer A (None of the above), only check the letter match
    if pred_letter == "A" or target_letter == "A":
        return (element_acc, element_acc)

    # For TYPE and SELECT actions, compare both action and value
    if target_answer.action in ["TYPE", "SELECT"]:
        if not pred_answer.action or not pred_answer.value:
            return (element_acc, 0.0)  # Missing required action/value
        pred_action = f"{pred_answer.action} {pred_answer.value}".strip()
        target_action = f"{target_answer.action} {target_answer.value}".strip()
        action_f1 = calculate_text_similarity(pred_action, target_action)
    else:
        # For CLICK actions, only compare the action type
        action_f1 = float(pred_answer.action == target_answer.action)

    return (element_acc, action_f1)


def is_answer_provided(p: str | None) -> int:
    """Check if a valid answer is provided.

    Args:
        p: The answer string to check

    Returns:
        1 if any text is provided, 0 if empty/None
    """
    if p is None or not p.strip():
        return 0
    return 1


@scorer(
    metrics=[
        element_accuracy(),
        action_f1(),
        step_success_rate(),
        task_success_rate(),
    ]
)
def mind2web_scorer() -> Scorer:
    """Scoring function for the Mind2Web evaluation.

    This function extracts the predicted answer from the model's completion,
    compares it with the target answer, and calculates element accuracy and
    action F1 scores. It handles different action types (CLICK, TYPE, SELECT)
    appropriately and supports the "None of the above" option.

    The scorer uses the following metrics:
    - element_accuracy: Accuracy of selecting the correct element
    - action_f1: F1 score for predicting the correct operation and value
    - step_success_rate: Proportion of steps where both element and operation are correct
    - task_success_rate: Proportion of tasks where all steps are successful

    Returns:
        A Scorer function that evaluates Mind2Web predictions.
    """

    async def score(state: TaskState, target: Target) -> Score:
        completion = state.output.completion
        if not completion:
            return Score(
                value="I",
                answer="",
                metadata={
                    "element_acc": 0.0,
                    "is_answered": 0,
                    "action_f1": 0.0,
                    "task_id": state.metadata.get("task_id"),
                },
            )

        # Extract answer from completion
        pred_answer = extract_answer(completion)

        # Convert state metadata to Mind2WebMetadata
        metadata = state.metadata_as(Mind2WebMetadata)

        # Get target answer from metadata
        target_answer = metadata.answer

        element_acc, action_f1 = score_answer(pred_answer, target_answer)

        # Get task_id from metadata
        task_id = metadata.task_id
        return Score(
            value="C" if element_acc > 0 else "I",
            answer=completion,
            metadata={
                "element_acc": element_acc,
                "is_answered": is_answer_provided(completion),
                "action_f1": action_f1,
                "task_id": task_id,
            },
        )

    return score
