from __future__ import annotations

import asyncio
import contextlib
import socket
import tempfile
import time
from pathlib import Path

import asyncssh
import docker
import docker.errors
import pytest
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import ed25519

from mcp_vm_blackbox.targets import TargetConfig


def get_free_port() -> int:
    """Get a free port on localhost."""
    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
        s.bind(("", 0))
        return s.getsockname()[1]


@pytest.fixture(scope="session")
def docker_client():
    """Get Docker client from environment."""
    return docker.from_env()


@pytest.fixture(scope="session")
def ssh_keypair():
    """Generate a temporary ED25519 SSH keypair using cryptography."""
    with tempfile.TemporaryDirectory() as tmpdir:
        key_path = Path(tmpdir) / "id_ed25519"

        # Generate private key
        private_key = ed25519.Ed25519PrivateKey.generate()
        private_bytes = private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.OpenSSH,
            encryption_algorithm=serialization.NoEncryption(),
        )

        # Generate public key
        public_key = private_key.public_key()
        public_bytes = public_key.public_bytes(
            encoding=serialization.Encoding.OpenSSH, format=serialization.PublicFormat.OpenSSH
        )

        # Write to files
        key_path.write_bytes(private_bytes)
        key_path.chmod(0o600)

        pub_key_path = key_path.with_suffix(".pub")
        pub_key_path.write_bytes(public_bytes)
        pub_key_path.chmod(0o644)

        yield key_path, pub_key_path


@pytest.fixture(scope="session")
def docker_sshd_target(ssh_keypair, docker_client):
    """Start an Ubuntu sshd docker container and return a TargetConfig."""
    priv_key, pub_key = ssh_keypair
    port = get_free_port()

    dockerfile_dir = Path(__file__).parent
    image_tag = "mcp-vm-blackbox-test-sshd:latest"
    container_name = f"test-sshd-{port}"

    print(f"\nBuilding {image_tag} from {dockerfile_dir}/Dockerfile...")
    docker_client.images.build(path=str(dockerfile_dir), tag=image_tag)

    # We need to mount the public key into the container's authorized_keys.
    # The simplest way is to mount it, then start the container.
    pub_key_content = pub_key.read_text().strip()

    with tempfile.TemporaryDirectory() as auth_keys_dir:
        auth_keys_path = Path(auth_keys_dir) / "authorized_keys"
        auth_keys_path.write_text(pub_key_content + "\n")

        # Ensure it has right permissions for testuser (UID 1000 usually)
        Path(auth_keys_path).chmod(0o644)

        print(f"\nStarting container: {container_name}")
        container = docker_client.containers.run(
            image_tag,
            detach=True,
            remove=True,
            name=container_name,
            ports={"22/tcp": port},
            volumes={str(auth_keys_path): {"bind": "/home/testuser/.ssh/authorized_keys", "mode": "ro"}},
        )

        # Wait for SSH to be responsive
        cfg = TargetConfig(ssh_host="127.0.0.1", ssh_port=port, ssh_user="testuser", ssh_identity_file=str(priv_key))

        async def check_ssh() -> bool:
            try:
                async with asyncssh.connect(
                    "127.0.0.1", port=port, username="testuser", client_keys=[str(priv_key)], known_hosts=None
                ) as conn:
                    result = await conn.run("echo 'ready'")
                    stdout_str = result.stdout.decode() if isinstance(result.stdout, bytes) else (result.stdout or "")
                    return result.exit_status == 0 and "ready" in stdout_str
            except Exception as e:
                print(f"SSH check exception: {e}")
                return False

        max_retries = 30
        for _i in range(max_retries):
            if asyncio.run(check_ssh()):
                break
            time.sleep(0.5)
        else:
            with contextlib.suppress(docker.errors.NotFound):
                container.remove(force=True)
            raise RuntimeError(f"SSH container did not become ready on port {port}")

        yield cfg

        # Teardown
        with contextlib.suppress(docker.errors.NotFound):
            container.remove(force=True)
