import time
import paramiko
from vncdotool import api as vnc_api

from vagary.schemas import Machine, VagrantMachine, QemuMachine
from vagary.providers.base import BaseProvider
from vagary.providers.vagrant import VagrantProvider
from vagary.providers.qemu import QemuProvider


def _make_provider(config: Machine) -> BaseProvider:
    if isinstance(config, VagrantMachine):
        return VagrantProvider(
            box=config.box,
            memory=config.memory,
            cpus=config.cpus,
            vnc_port=config.vnc_port,
        )
    if isinstance(config, QemuMachine):
        return QemuProvider(
            disk_image=config.disk_image,
            disk_format=config.disk_format,
            memory=config.memory,
            cpus=config.cpus,
            ssh_user=config.ssh_user,
            ssh_password=config.ssh_password,
            ssh_key=config.ssh_key,
            vnc_port=config.vnc_port,
            ssh_port=config.ssh_port,
        )
    raise ValueError(f"Unknown machine provider: {type(config)}")


class MachineManager:
    def __init__(self, config: Machine):
        self.config = config
        self.provider: BaseProvider = _make_provider(config)

        self.ssh_client = None
        self.vnc_client = None

    @property
    def vnc_host(self) -> str:
        return self.provider.vnc_host

    @property
    def vnc_port(self) -> int:
        return self.provider.vnc_port

    def start(self):
        self.provider.start()

    def connect_ssh(self):
        cfg = self.provider.get_ssh_config()
        if cfg is None:
            print("> SSH not available for this provider, skipping")
            return

        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        client.connect(
            hostname=cfg["HostName"],
            port=int(cfg["Port"]),
            username=cfg["User"],
            key_filename=cfg.get("IdentityFile"),
            password=cfg.get("Password"),
        )

        self.ssh_client = client
        print("> SSH connected")

    def connect_vnc(self, password: str | None = None, timeout: float = 30.0):
        self.vnc_client = vnc_api.connect(self.vnc_host, password=password, timeout=timeout)
        time.sleep(0.5)
        print("> VNC connected")

    def exec(self, command: str):
        if self.provider.get_ssh_config() is None:
            raise RuntimeError(
                f"Provider '{type(self.provider).__name__}' does not support SSH. "
                "Cannot run shell step."
            )
        if not self.ssh_client:
            raise RuntimeError("SSH is not connected. Call connect_ssh() first.")
        stdin, stdout, stderr = self.ssh_client.exec_command(command)
        return stdout.read().decode(), stderr.read().decode()

    def vnc_key(self, key: str):
        self.vnc_client.keyPress(key)

    def vnc_type(self, text: str):
        for k in text:
            self.vnc_client.keyPress(k)

    def vnc_move(self, x: int, y: int):
        self.vnc_client.mouseMove(x, y)

    def vnc_click(self, x: int, y: int, button=1):
        self.vnc_client.mouseMove(x, y)
        self.vnc_client.mousePress(button)

    def vnc_screenshot(self, path: str):
        if not self.vnc_client:
            raise RuntimeError("VNC not connected")
        self.vnc_client.captureScreen(path)

    def vnc_expect_region(self, path: str, x: int, y: int, maxrms: float = 0.0):
        if not self.vnc_client:
            raise RuntimeError("VNC not connected")

        # length mismatch (e.g. RGBA vs RGB) silently skips the comparison and loops forever.
        # We need to convert to RGB to avoid this.

        from PIL import Image
        import tempfile, os

        expected = Image.open(path).convert("RGB")
        with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp:
            tmp_path = tmp.name
        try:
            expected.save(tmp_path)
            self.vnc_client.expectRegion(tmp_path, x, y, maxrms=maxrms)
        finally:
            os.unlink(tmp_path)

    def close(self):
        if self.ssh_client:
            self.ssh_client.close()
        if self.vnc_client:
            self.vnc_client.disconnect()
        vnc_api.shutdown()
        print("> Connections closed")
