import base64
import json
from io import BytesIO
from typing import Any

from PIL import Image
from pydantic import BaseModel, field_validator

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


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 ActionStep(BaseModel):
    """Represents a single episode in the trace."""

    initial_screenshot: Image.Image
    task: str | None = None
    action: Action

    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 action to a human-readable string."""
        return f"\nAction: {self.action.to_string()}\n"


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

    steps: list[ActionStep]
    task: str | None = None

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


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

    step_actions: list[Action]
    step_screenshots: list[bytes]
    step_tasks: list[str]

    @field_validator("step_screenshots", mode="before")
    def validate_step_screenshots(cls, v):
        step_screenshots: list[bytes] = []
        for step_screenshot in v:
            if step_screenshot is None:
                step_screenshots.append(b"\x00")
            elif isinstance(step_screenshot, bytes):
                step_screenshots.append(step_screenshot)
            else:
                step_screenshots.append(step_screenshot["bytes"])
        return step_screenshots

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

        for step_index in range(len(self.step_actions)):
            operations = self.step_actions[step_index].operations
            if self.step_screenshots[step_index] == b"\x00":
                # The screenshot is missing, skip the rest of the episode
                break

            screenshot = Image.open(BytesIO(self.step_screenshots[step_index]))
            step = ActionStep(
                initial_screenshot=screenshot,
                action=Action(operations=operations),
                task=self.step_tasks[step_index],
            )

            steps.append(step)

        return Trace(steps=steps)


class CompleteTrace(BaseModel):
    """Represents a complete trace from the dataset."""

    actions: list[Action]  # Each action should be a list of operations
    screenshots_b64: list[str]
    pos_candidates: list[list[str]] | None = None
    neg_candidates: list[list[str]] | None = None
    target_action_indexes: list[int] | None = None  # There is no target action index for android control
    step_tasks: list[str]  # Rename confirmed_tasks to step_tasks
    wider_goal: str | None = None  # There is no wider goal for MMind2web

    def to_trace(self) -> Trace:
        """Convert the sample to a Trace object."""
        steps: list[ActionStep] = []
        for i, action in enumerate(self.actions):
            screenshot = Image.open(BytesIO(base64.b64decode(self.screenshots_b64[i])))

            step = ActionStep(
                initial_screenshot=screenshot,
                action=action,
                task=self.step_tasks[i],
            )

            # For any operation, we click on the element before. For now, HOVER, ENTER, SELECT are bound to CLICK. TYPE required a CLICK first.
            # step.add_action(Action(op="Click", value=step.bounding_box.to_click_center_position_string()))
            # if operation.op == "TYPE":
            #     step.add_action(Action(op="Type", value=f'"{operation.value}"'))

            steps.append(step)

        return Trace(steps=steps, task=self.wider_goal if self.wider_goal else None)
