import collections
import concurrent.futures
import logging
from typing import Any, Generator, Literal, Sequence

import numpy as np
from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
from smolagents import Model
from tqdm import tqdm

from screensuite.basebenchmark import AnnotatedInput, EvaluationConfig
from screensuite.benchmark_result import BenchmarkResult
from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmark
from screensuite.benchmarks.singlestep.mmind2web.config import MMind2WebConfig
from screensuite.benchmarks.singlestep.mmind2web.models import (
    AggregatedTraceSample,
    BoundingBox,
    Operation,
)
from screensuite.benchmarks.singlestep.mmind2web.prompt import MMind2WebPrompt
from screensuite.benchmarks.utils import (
    evaluate_click_in_bounding_box,
    evaluate_type_action,
)
from screensuite.chat_message import (
    ChatMessage,
    ImageContent,
    TextContent,
    dump_chat_message_list,
)
from screensuite.registry import registry
from screensuite.response_generation import AnnotatedOutput, get_model_responses

logger = logging.getLogger(__name__)

ActionGroundTruth = tuple[list[Operation], BoundingBox, tuple[int, int]]


class MMind2WebBenchmark(HubBaseBenchmark[MMind2WebConfig, None]):
    def _get_annotated_input_from_sample(
        self, sample: dict[str, list[Any]], evaluation_config: EvaluationConfig
    ) -> Generator[AnnotatedInput[ActionGroundTruth], None, None]:
        """
        Transforms the sample into a list of input messages and a ground truth.
        """

        sample_model = AggregatedTraceSample(
            operations=sample["operation"],
            initial_screenshots=sample["screenshot"],
            step_tasks=sample["confirmed_task"],
            pos_candidates=sample["pos_candidates"],
            neg_candidates=sample["neg_candidates"],
            target_action_indexes=sample["target_action_index"],
        )

        trace = sample_model.to_trace()

        # Yield the query and the target action for each step
        for current_step in range(trace.num_steps):
            role_message: ChatMessage = ChatMessage(role="user")

            role_message.content.append(
                TextContent(
                    type="text",
                    text=MMind2WebPrompt.INSTRUCTION_PROMPT.value.format(task=trace.steps[current_step].task),
                )
            )

            # Add previous steps' screenshot and actions
            action_index = 0
            for i in range(current_step):
                if i < current_step - self.config.max_step_memory:
                    continue
                role_message.content.append(TextContent(text=f"Screenshot {action_index}:\n"))
                role_message.content.append(ImageContent(image=trace.steps[i].initial_screenshot))
                # For testing
                # role_message.content.append(DummyImageContent(type="image", image=f"screenshot_{i}.png"))
                role_message.content.append(TextContent(text=trace.steps[i].action_to_string()))
                action_index += 1

            # Add first step prompt if it's the first episode
            if current_step == 0:
                role_message.content.append(TextContent(text="\nNo previous actions.\n"))

            # For testing
            # role_message.content.append(DummyImageContent(type="image", image=f"screenshot_{current_step}.png"))
            resolution = (
                trace.steps[current_step].initial_screenshot.width,
                trace.steps[current_step].initial_screenshot.height,
            )

            if evaluation_config.use_uitars_action_space:
                prompt_template = MMind2WebPrompt.ASSISTANT_PROMPT_ABSOLUTE_CLICK_UITARS.value
            elif evaluation_config.use_h_action_space:
                prompt_template = MMind2WebPrompt.ASSISTANT_PROMPT_ABSOLUTE_CLICK.value.replace("click(", "Click(")
            else:
                prompt_template = MMind2WebPrompt.ASSISTANT_PROMPT_ABSOLUTE_CLICK.value

            # Add assistant prompt
            if evaluation_config.image_before_text:
                role_message.content.append(TextContent(text=f"Screenshot {action_index}:\n"))
                role_message.content.append(ImageContent(image=trace.steps[current_step].initial_screenshot))
                role_message.content.append(TextContent(text=prompt_template.format(resolution=resolution)))
            else:
                role_message.content.append(TextContent(text=prompt_template.format(resolution=resolution)))
                role_message.content.append(ImageContent(image=trace.steps[current_step].initial_screenshot))
            messages = dump_chat_message_list([role_message])

            bbox = trace.steps[current_step].bounding_box
            if bbox is None:
                raise ValueError("Bounding box is None, cannot evaluate the action")
            operations = trace.steps[current_step].action.operations
            optional_system_prompt = (
                [ChatMessage(role="system", content=[TextContent(text=evaluation_config.custom_system_prompt)])]
                if evaluation_config.custom_system_prompt
                else []
            )
            yield AnnotatedInput(
                messages=dump_chat_message_list(optional_system_prompt + [role_message]),
                ground_truth=(operations, bbox, resolution),
            )

    # This is a function to aggregate the traces. TODO: Push the dataset to the Hub with the traces aggregated.
    def _aggregate_traces(self):
        """
        Aggregate the traces in format for the model, using multithreading.
        """
        new_datasets: dict[str, DatasetDict | Dataset | IterableDatasetDict | IterableDataset] = {}
        features = [
            "confirmed_task",
            "operation",
            "screenshot",
            "pos_candidates",
            "neg_candidates",
            "target_action_index",
        ]
        new_row_counter = 0
        old_row_counter = 0

        def process_split(
            split: str, dataset: DatasetDict | Dataset | IterableDatasetDict | IterableDataset
        ) -> tuple[str, Dataset, int, int]:
            if isinstance(dataset, (DatasetDict, IterableDatasetDict)):
                raise ValueError("DatasetDict or IterableDatasetDict is not supported")

            # First pass: group samples by trace
            traces_samples: list[list[dict]] = []
            current_trace: list[dict] = []
            current_idx: int = -1
            total_samples = len(dataset) if not isinstance(dataset, IterableDataset) else None
            split_old_counter = 0
            split_new_counter = 0
            for sample in tqdm(dataset, total=total_samples, desc=f"Grouping {self.name} traces - split: '{split}'"):
                if not isinstance(sample, dict):
                    raise ValueError("Dataset should yield dicts")

                # reset the current index if it is greater than the target action index -> new trace
                if current_idx >= int(sample["target_action_index"]):
                    if current_trace:
                        traces_samples.append(current_trace)
                        current_trace = []
                    current_idx = -1

                # ensure the trace is in the correct order
                if current_idx != int(sample["target_action_index"]) - 1:
                    raise ValueError("Current index is not equal to the target action index")

                current_idx = int(sample["target_action_index"])
                current_trace.append(sample)
                split_old_counter += 1

            if current_trace:
                traces_samples.append(current_trace)

            # Second pass: process traces in parallel
            def process_trace(trace_samples: list[dict]) -> tuple[dict[str, list[Any]], int]:
                accumulated_samples: dict[str, list[Any]] = collections.defaultdict(list)
                for sample in trace_samples:
                    for feature in features:
                        accumulated_samples[feature].append(sample[feature])
                return accumulated_samples, len(trace_samples)

            traces: list[dict[str, list[Any]]] = []
            with concurrent.futures.ThreadPoolExecutor() as executor:
                futures = [executor.submit(process_trace, trace) for trace in traces_samples]
                for future in tqdm(
                    concurrent.futures.as_completed(futures),
                    total=len(futures),
                    desc=f"Processing {self.name} traces - split: '{split}'",
                ):
                    accumulated_samples, num_samples = future.result()
                    traces.append(accumulated_samples)
                    split_new_counter += num_samples

            new_dataset = Dataset.from_list(traces)
            return split, new_dataset, split_old_counter, split_new_counter

        # Process splits in parallel
        with concurrent.futures.ThreadPoolExecutor() as executor:
            futures = []
            for split, dataset in self.datasets.items():
                if split != "test_website":
                    continue
                futures.append(executor.submit(process_split, split, dataset))

            for future in concurrent.futures.as_completed(futures):
                split, new_dataset, split_old_counter, split_new_counter = future.result()
                new_datasets[split] = new_dataset
                old_row_counter += split_old_counter
                new_row_counter += split_new_counter

        self.datasets = new_datasets

    def _calculate_weighted_metrics(
        self, accuracy_scores: dict[str, list[float]]
    ) -> tuple[dict[str, float], Literal["weighted_average_acc"]]:
        """
        Calculate weighted metrics from accuracy scores.

        Args:
            accuracy_scores: Dictionary containing accuracy scores for each split and metric type

        Returns:
            - Dictionary of metrics including weighted accuracies
            - Reference field name for the benchmark result
        """

        if not self.datasets:
            raise ValueError("Datasets are not yet evaluated")

        def get_mean(scores: list[float]) -> float:
            return float(np.mean(scores)) if scores else 0.0

        metrics: dict[str, float] = {}
        total_samples = 0
        type_samples = 0
        weighted_click_acc = 0.0
        weighted_type_acc = 0.0
        weighted_avg_acc = 0.0

        for split in self.datasets.keys():
            split_samples = len(accuracy_scores[f"{split}_click_acc"])
            total_samples += split_samples

            # Calculate weighted accuracies
            click_mean = get_mean(accuracy_scores[f"{split}_click_acc"])
            not_nan_type_acc = [s for s in accuracy_scores[f"{split}_type_acc"] if not np.isnan(s)]
            type_mean = get_mean(not_nan_type_acc) if not_nan_type_acc else np.nan
            avg_mean = get_mean(accuracy_scores[f"{split}_average_acc"])

            weighted_click_acc += click_mean * split_samples
            weighted_avg_acc += avg_mean * split_samples
            if not np.isnan(type_mean):
                weighted_type_acc += type_mean * split_samples
                type_samples += split_samples

            metrics.update(
                {
                    f"{split}_click_acc": click_mean,
                    f"{split}_type_acc": type_mean,
                    f"{split}_average_acc": avg_mean,
                }
            )

        # Add weighted metrics
        if total_samples > 0:
            metrics["weighted_click_acc"] = float(weighted_click_acc / total_samples)
            metrics["weighted_type_acc"] = float(weighted_type_acc / type_samples) if type_samples > 0 else np.nan
            metrics["weighted_average_acc"] = float(weighted_avg_acc / total_samples)

        return metrics, "weighted_average_acc"

    def _calculate_proportion_missing(self, responses: Sequence[AnnotatedOutput[ActionGroundTruth] | None]) -> float:
        """
        Calculate the proportion of missing responses.

        Args:
            responses: Dictionary containing responses for each split

        Returns:
            Proportion of missing responses
        """
        if not self.datasets:
            raise ValueError("Datasets are not yet evaluated")

        total_responses = 0
        missing_responses = 0
        total_responses += len(responses)
        for annotated_output in responses:
            if annotated_output is None or annotated_output.output is None:
                missing_responses += 1

        if total_responses == 0:
            return 0.0
        return float(missing_responses) / total_responses

    def evaluate(
        self,
        model: Model,
        evaluation_config: EvaluationConfig,
        env_config: None = None,
    ) -> BenchmarkResult:
        """
        Evaluate the model on the benchmark

        Args:
            model: The model to evaluate
            evaluation_config: Configuration for inference

        Returns:
            Evaluation results
        """
        if not self.datasets:
            self.load(streaming=True)

        accuracy_scores: dict[str, list[float]] = collections.defaultdict(list)
        responses: dict[str, list[AnnotatedOutput[ActionGroundTruth] | None]] = {}

        self._aggregate_traces()
        for split, dataset in self.datasets.items():
            assert isinstance(dataset, Dataset)
            responses[split] = get_model_responses(
                dataset,
                model,
                self._get_annotated_input_from_sample,
                f"{self.name}_{split}",
                evaluation_config,
                max_tokens=self.config.max_tokens,
                temperature=self.config.temperature,
            )
            for annotated_output in responses[split]:
                if annotated_output is not None and annotated_output.output is not None:
                    prediction, (target_actions, target_bbox, image_resolution) = (
                        annotated_output.output,
                        annotated_output.ground_truth,
                    )
                    prediction = prediction.lower()
                    for target_action in target_actions:
                        if prediction is not None:
                            logger.debug(
                                "Prediction: ", prediction, " | Target: ", target_action, target_bbox.to_xyxy()
                            )
                            if target_action.type == "click":
                                click_acc = evaluate_click_in_bounding_box(
                                    prediction,
                                    target_bbox,
                                    image_width=image_resolution[0],
                                    image_height=image_resolution[1],
                                )
                                logger.debug(
                                    "Predicted: ",
                                    prediction,
                                    " | Target: ",
                                    target_action,
                                    target_bbox.to_xyxy(),
                                    "| Click acc:",
                                    click_acc,
                                )
                                accuracy_scores[f"{split}_click_acc"].append(click_acc)
                            if target_action.type == "type":
                                type_acc = evaluate_type_action(prediction, target_action.args["text"])
                                accuracy_scores[f"{split}_type_acc"].append(type_acc)
                                accuracy_scores[f"{split}_average_acc"].append((click_acc + type_acc) / 2)
                            else:
                                accuracy_scores[f"{split}_type_acc"].append(np.nan)
                                accuracy_scores[f"{split}_average_acc"].append(click_acc)

        metrics, reference_field = self._calculate_weighted_metrics(accuracy_scores)
        metrics["proportion_missing"] = self._calculate_proportion_missing(responses[split])
        metrics["count_samples"] = len(responses[split])
        return BenchmarkResult(metrics=metrics, reference_field=reference_field)


mmind2web = MMind2WebBenchmark(
    name="mmind2web",
    config=MMind2WebConfig.create(),
    tags=["mmind2web", "hf_dataset", "offline", "agent", "web", "trace", "to_evaluate"],
)


if __name__ == "__main__":
    # test
    from smolagents import OpenAIServerModel

    model = OpenAIServerModel(model_id="gpt-4o")
    mmind2web.evaluate(
        model,
        EvaluationConfig(timeout=10, parallel_workers=1, test_mode=True, run_name="test_mmind2web"),
    )
else:
    registry.register(mmind2web)
