import json
import logging
import os
from typing import Any, Generator

import numpy as np
from datasets import Dataset, load_dataset
from smolagents import Model
from torch.utils.data import Dataset as TorchDataset

from screensuite.basebenchmark import EvaluationConfig
from screensuite.benchmark_result import BenchmarkResult
from screensuite.benchmarks.hubbasebenchmark import HubBaseBenchmark
from screensuite.benchmarks.model import NormalizedClickCoordinates
from screensuite.benchmarks.singlestep.androidcontrol.config import AndroidControlConfig
from screensuite.benchmarks.singlestep.androidcontrol.models import (
    Action,
    ActionStep,
    CompleteTrace,
    Operation,
)
from screensuite.benchmarks.singlestep.androidcontrol.prompt import AndroidControlPrompt
from screensuite.benchmarks.utils import (
    bootstrap_confidence_interval,
    evaluate_click_distance,
    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 (
    AnnotatedInput,
    AnnotatedOutput,
    get_model_responses,
)

logger = logging.getLogger(__name__)

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


class AndroidControlDataset(TorchDataset):
    def __init__(
        self,
        data_path: str = "data/aguvis/train/android_control.json",
        images_folder: str = "android_control/images/",
        base_image_folder: str = "",  # Base path to prepend to images_folder
    ):
        super().__init__()
        self.images_folder = os.path.join(base_image_folder, images_folder)

        # Load the dataset
        print(f"Loading {data_path}")
        with open(data_path) as file:
            self.data = json.load(file)
        print(f"Loaded {len(self.data)} samples from {data_path}")

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data[idx]

        # Handle image paths
        if "image" in sample:
            image_file = sample["image"]
            if isinstance(image_file, list):
                image_paths = [os.path.join(self.images_folder, img) for img in image_file]
            else:
                image_paths = [os.path.join(self.images_folder, image_file)]

            # Here you would process the images with your processor
            # For demonstration, just returning paths
            return {"id": sample.get("id", idx), "image_paths": image_paths, "conversations": sample["conversations"]}
        else:
            return {"id": sample.get("id", idx), "conversations": sample["conversations"]}


def convert_to_hf_dataset(dataset: AndroidControlDataset) -> Dataset:
    """Convert the custom dataset to a Hugging Face dataset."""
    data_dict: dict[str, list[Any]] = {"id": [], "image_paths": [], "conversations": []}

    for i in range(len(dataset)):
        item = dataset[i]
        data_dict["id"].append(item.get("id"))
        data_dict["image_paths"].append(item.get("image_paths", []))
        data_dict["conversations"].append(item.get("conversations", []))

    return Dataset.from_dict(data_dict)


def convert_android_control_operation(operation_dict: dict[str, Any]) -> Operation:
    """Convert an Android Control dictionary into an Operation object."""
    action_type = operation_dict["action_type"].lower()
    # Create an operation based on the action type
    if action_type == "click" or action_type.replace("_", "") == "longpress":
        args = {"x": operation_dict["x"], "y": operation_dict["y"]}
        action_type = "click"
    elif action_type == "input_text":
        action_type = "type"
        args = {"text": operation_dict["text"]}
    elif action_type == "scroll":
        args = {"direction": operation_dict["direction"]}
    elif action_type == "wait":
        print(operation_dict)
        args = {"duration": 2.0}
    elif action_type == "open_app":
        args = {"app_name": operation_dict["app_name"]}
    elif action_type == "navigate_back" or action_type == "navigate_home":
        args = {}
    else:
        raise ValueError(f"Unknown action type: {operation_dict['action_type']}")

    return Operation(type=action_type, args=args)


class AndroidControlBenchmark(HubBaseBenchmark[AndroidControlConfig, None]):
    def load(self, streaming: bool = False) -> None:
        self.dataset: Dataset = load_dataset(self.config.hf_repo, split=self.config.split, streaming=streaming)  # type: ignore

    def _get_annotated_input_from_sample(
        self, sample_dict: dict[str, list[Any] | str], evaluation_config: EvaluationConfig
    ) -> Generator[AnnotatedInput[ActionGroundTruth], None, None]:
        """
        Transform the sample_dict into an AnnotatedInput with a list of messages and a ground truth.
        """
        sample = CompleteTrace(
            actions=[
                Action(operations=[convert_android_control_operation(operation)])  # type: ignore
                for operation in sample_dict["actions"]
            ],
            screenshots_b64=sample_dict["screenshots_b64"],
            step_tasks=sample_dict["step_instructions"],
            wider_goal=sample_dict["goal"],
        )
        trace = sample.to_trace()

        # Yield the query and the target action for each step
        for current_step_id in range(trace.num_steps):  # Skip the first step (task)
            current_step = trace.steps[current_step_id]
            assert isinstance(current_step, ActionStep)
            role_message: ChatMessage = ChatMessage(role="user")

            role_message.content.append(
                TextContent(
                    type="text",
                    text=AndroidControlPrompt.INSTRUCTION_PROMPT.value.format(task=current_step.task),
                )
            )

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

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

            resolution = current_step.initial_screenshot.width, current_step.initial_screenshot.height

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

            prompt = prompt_template.format(resolution=resolution)
            # 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=current_step.initial_screenshot))
                role_message.content.append(TextContent(text=prompt))
            else:
                role_message.content.append(TextContent(text=prompt))
                role_message.content.append(ImageContent(image=current_step.initial_screenshot))
            messages = dump_chat_message_list([role_message])

            # bbox = trace.get_step(current_step).bounding_box
            # if bbox is None:
            #     raise ValueError("Bounding box is None, cannot evaluate the action")
            operations = current_step.action.operations

            screenshot_dimensions = current_step.initial_screenshot.size

            optional_system_prompt = (
                [ChatMessage(role="system", content=[TextContent(text=evaluation_config.custom_system_prompt)])]
                if evaluation_config.custom_system_prompt
                else []
            )

            # Below if for debugging
            # check_messages_below_max_context_length(messages, 32000)
            yield AnnotatedInput(
                messages=dump_chat_message_list(optional_system_prompt + [role_message]),
                ground_truth=(operations, screenshot_dimensions),
            )

    def score_single_output(
        self,
        annotated_output: AnnotatedOutput[ActionGroundTruth],
    ) -> float:
        prediction, (target_actions, target_dimensions) = (
            annotated_output.output,
            annotated_output.ground_truth,
        )
        prediction = prediction.lower()
        assert len(target_actions) == 1, "Only one target action is supported for now"
        target_action = target_actions[0]
        if prediction.split("(")[0].replace("_", "").replace("longpress", "click").lower() != target_action.type:
            return 0.0

        if target_action.type == "click":
            click_coordinates = NormalizedClickCoordinates.from_xy(
                xy=(target_action.args["x"], target_action.args["y"]),
                image_width=target_dimensions[0],
                image_height=target_dimensions[1],
            )
            click_acc = evaluate_click_distance(
                prediction, click_coordinates, image_width=target_dimensions[0], image_height=target_dimensions[1]
            )
            return click_acc
        elif target_action.type == "type":
            type_acc = evaluate_type_action(prediction, target_action.args["text"])
            return type_acc
        elif target_action.type == "open_app":
            return target_action.args["app_name"].lower() in prediction
        elif target_action.type == "scroll":
            return target_action.args["direction"] in prediction.split("(")[1]
        elif target_action.type == "wait":
            return "wait" in prediction
        elif target_action.type == "navigate_back":
            return "navigate_back" in prediction
        elif target_action.type == "navigate_home":
            return "navigate_home" in prediction
        else:
            raise NotImplementedError(f"Unknown action type: {target_action.type}")

    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=False)

        accuracy_scores: list[float] = []
        responses: list[AnnotatedOutput[ActionGroundTruth] | None] = []

        # self._aggregate_traces()
        responses = get_model_responses(
            self.dataset,
            model,
            self._get_annotated_input_from_sample,
            f"{self.name}_{self.config.split}",
            evaluation_config,
            max_tokens=self.config.max_tokens,
            temperature=self.config.temperature,
        )
        for annotated_output in responses:
            if annotated_output is not None:
                score = self.score_single_output(annotated_output)
                accuracy_scores.append(score)
        metrics = {"action_acc": float(np.mean(accuracy_scores)) if accuracy_scores else np.nan}
        if accuracy_scores:
            _, (lower, upper) = bootstrap_confidence_interval(accuracy_scores)
            metrics["action_acc_confidence_interval_lower"] = float(lower)
            metrics["action_acc_confidence_interval_upper"] = float(upper)
        reference_field = "action_acc"
        metrics["proportion_missing"] = self._calculate_proportion_missing(responses)
        metrics["count_samples"] = len(responses)
        return BenchmarkResult(metrics=metrics, reference_field=reference_field)


android_control = AndroidControlBenchmark(
    name="android_control",
    config=AndroidControlConfig.create(),
    tags=["android_control", "hf_dataset", "offline", "agent", "web", "trace", "to_evaluate"],
)


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

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