import json
from io import BytesIO
from typing import Any

from PIL import Image
from pydantic import BaseModel, field_validator

from screensuite.benchmarks.model import BoundingBox
from screensuite.benchmarks.singlestep.common_models import Action, Operation


def convert_mind2web_operations(operation_json: str) -> list[Operation]:
    """Convert a Mind2Web operation to an Operation."""
    operation = json.loads(operation_json)
    operations: list[Operation] = []

    if operation["op"].lower() not in ["click", "hover", "enter", "select", "type"]:
        raise ValueError(f"Unknown operation: {operation['op']}")

    # For any operation, we click on the element before. For now, HOVER, ENTER, SELECT are bound to CLICK. TYPE required a CLICK first.
    operations.append(Operation(type="click", args={}))
    if operation["op"].lower() == "type":
        operations.append(Operation(type="type", args={"text": operation["value"]}))

    return operations


class Candidate(BaseModel):
    """Represents a candidate element with its attributes."""

    attributes: dict[str, Any]

    class Config:
        extra = "allow"

    @classmethod
    def from_json(cls, json_str: str) -> "Candidate":
        """Create a Candidate from a JSON string."""
        return cls(attributes=json.loads(json.loads(json_str)["attributes"]))


class Mind2WebActionStep(BaseModel):
    """Represents a single episode in the trace."""

    initial_screenshot: Image.Image
    action: Action
    pos_candidates: list[Candidate]
    neg_candidates: list[Candidate]
    target_action_index: int
    task: str
    bounding_box: BoundingBox | None = None

    class Config:
        arbitrary_types_allowed = True

    def add_operation(self, operation: Operation) -> None:
        """Add an operation to the episode."""
        self.action.operations.append(operation)

    def action_to_string(self) -> str:
        """Convert the actions to a human-readable string."""
        return f"\nAction: {'; '.join([operation.to_string() for operation in self.action.operations])}\n"


class Trace(BaseModel):
    """Represents a running trace of episodes."""

    steps: list[Mind2WebActionStep]

    @property
    def num_steps(self) -> int:
        """Get the number of steps in the trace."""
        return len(self.steps)


class AggregatedTraceSample(BaseModel):
    """Represents a complete sample from the dataset."""

    operations: list[str]
    initial_screenshots: list[bytes]
    pos_candidates: list[list[str]]
    neg_candidates: list[list[str]]
    target_action_indexes: list[int]
    step_tasks: list[str]

    @field_validator("target_action_indexes", mode="before")
    def validate_target_action_indexes(cls, v):
        v = [int(i) for i in v]
        if v != sorted(v):
            raise ValueError("Target action indexes must be sorted")
        return v

    @field_validator("initial_screenshots", mode="before")
    def validate_screenshots(cls, v):
        initial_screenshots: list[bytes] = []
        for screenshot in v:
            if screenshot is None:
                initial_screenshots.append(b"\x00")
            elif isinstance(screenshot, bytes):
                initial_screenshots.append(screenshot)
            elif isinstance(screenshot, dict) and "bytes" in screenshot:
                initial_screenshots.append(screenshot["bytes"])
            elif isinstance(screenshot, Image.Image):
                # Handle PIL Image objects
                buffer = BytesIO()
                screenshot.save(buffer, format="JPEG")
                initial_screenshots.append(buffer.getvalue())
            else:
                print(f"Invalid screenshot format: {type(screenshot)}")
                print("Skipping one screenshot")
                # raise ValueError(f"Invalid screenshot format: {type(screenshot)}")
        return initial_screenshots

    def to_trace(self) -> Trace:
        """Convert the sample to a Trace object."""
        episode: list[Mind2WebActionStep] = []

        for i in range(len(self.target_action_indexes)):
            if self.initial_screenshots[i] == b"\x00":
                # The screenshot is missing, skip the rest of the episode
                break
            try:
                # Use validated image loading
                screenshot = Image.open(BytesIO(self.initial_screenshots[i]))
            except Exception as e:
                print(f"Error opening screenshot: {e}")
                break

            operations = convert_mind2web_operations(self.operations[i])
            pos_candidates = [Candidate.from_json(c) for c in self.pos_candidates[i]]
            neg_candidates = [Candidate.from_json(c) for c in self.neg_candidates[i]]

            step = Mind2WebActionStep(
                initial_screenshot=screenshot,
                action=Action(operations=operations),
                pos_candidates=pos_candidates,
                neg_candidates=neg_candidates,
                target_action_index=int(self.target_action_indexes[i]),
                task=self.step_tasks[i],
            )

            # Add bounding box
            if pos_candidates:
                bbox_str = pos_candidates[0].attributes.get("bounding_box_rect")
                if bbox_str:
                    try:
                        step.bounding_box = BoundingBox.from_xywh_string(bbox_str, screenshot.width, screenshot.height)
                    except ValueError:
                        # The bounding box is out of the image dimensions or the format is invalid, skip the rest of the episode
                        break
                else:
                    # No bounding box, skip the rest of the episode
                    break
            else:
                # No positive candidates, skip the rest of the episode
                break

            episode.append(step)

        return Trace(steps=episode)
