from __future__ import annotations
from typing import Annotated, Union, List, Literal, Optional
from pydantic import BaseModel, Field
import yaml


class VagrantMachine(BaseModel):
    provider: Literal["vagrant"] = "vagrant"
    box: str = Field(description="Vagrant box name")
    memory: int = Field(default=8096, description="Memory (in MB)")
    cpus: int = Field(default=2, description="Number of vCPUs")
    vnc_port: int = Field(default=5900, description="VNC port on host")


class QemuMachine(BaseModel):
    provider: Literal["qemu"] = "qemu"
    disk_image: str = Field(description="Path to disk image (e.g. built via bootc install to-disk)")
    disk_format: str = Field(default="raw", description="Disk image format (raw, qcow2, ...)")
    memory: int = Field(default=8096, description="Memory (in MB)")
    cpus: int = Field(default=2, description="Number of vCPUs")
    vnc_port: int = Field(default=5900, description="VNC port on host")
    ssh_port: int = Field(default=2222, description="Host port forwarded to guest SSH (22)")
    ssh_user: Optional[str] = Field(default=None, description="SSH username; if not set, SSH is disabled")
    ssh_password: Optional[str] = Field(default=None, description="SSH password")
    ssh_key: Optional[str] = Field(default=None, description="Path to SSH private key")


Machine = Annotated[
    Union[VagrantMachine, QemuMachine],
    Field(discriminator="provider"),
]


class Position(BaseModel):
    x: int
    y: int


class BaseStep(BaseModel):
    type: str
    name: Optional[str] = None

    def get_display_name(self) -> str:
        return self.name if self.name else self.type

    def get_display_info(self) -> str:
        return ""

    def format_step(self, indent: int = 0) -> str:
        prefix = "  " * indent
        icon = "▶"
        name = self.get_display_name()
        info = self.get_display_info()
        suffix = f" → {info}" if info else ""

        return f"{prefix}{icon} [{self.type}] {name}{suffix}"

class ShellStep(BaseStep):
    type: Literal["shell"] = "shell"
    run: str

    def get_display_info(self) -> str:
        preview = self.run[:50] + "..." if len(self.run) > 50 else self.run
        return f"`{preview}`"


class KeyStep(BaseStep):
    type: Literal["key"] = "key"
    key: str

    def get_display_info(self) -> str:
        return f"'{self.key}'"


class TypeStep(BaseStep):
    type: Literal["type"] = "type"
    text: str

    def get_display_info(self) -> str:
        preview = self.text[:50] + "..." if len(self.text) > 50 else self.text
        return f'"{preview}"'


class SleepStep(BaseStep):
    type: Literal["sleep"] = "sleep"
    time: float

    def get_display_info(self) -> str:
        return f"{self.time}s"


class ScreenshotStep(BaseStep):
    type: Literal["screenshot"] = "screenshot"
    path: str

    def get_display_info(self) -> str:
        return self.path


class ClickStep(BaseStep):
    type: Literal["click"] = "click"
    pos: Position

    def get_display_info(self) -> str:
        return f"({self.pos.x}, {self.pos.y})"


class MoveStep(BaseStep):
    type: Literal["move"] = "move"
    pos: Position

    def get_display_info(self) -> str:
        return f"({self.pos.x}, {self.pos.y})"


class GroupStep(BaseStep):
    type: Literal["group"] = "group"
    steps: List[Step]

    def get_display_info(self) -> str:
        return f"({len(self.steps)} steps)"


class ExpectRegionStep(BaseStep):
    type: Literal["expect_region"] = "expect_region"
    path: str
    pos: Position
    maxrms: float = 0.0

    def get_display_info(self) -> str:
        return f"{self.path} at ({self.pos.x}, {self.pos.y}) (maxrms={self.maxrms})"


Step = Union[
    ShellStep,
    KeyStep,
    TypeStep,
    SleepStep,
    ScreenshotStep,
    ExpectRegionStep,
    ClickStep,
    MoveStep,
    GroupStep,
]


class Settings(BaseModel):
    sleep_between_steps: float = 0


class Scenario(BaseModel):
    machine: Machine
    steps: List[Step]
    settings: Settings = Field(default_factory=Settings)

    @classmethod
    def from_yaml(cls, yaml_string: str) -> 'Scenario':
        """Загрузить сценарий из YAML"""
        data = yaml.safe_load(yaml_string)
        return cls(**data)

    def to_yaml(self) -> str:
        """Сохранить сценарий в YAML"""
        return yaml.dump(
            self.model_dump(by_alias=True, exclude_none=True),
            sort_keys=False,
            allow_unicode=True
        )
