"""
Install RDDS (Remote Desktop Driver Server) on a running OSWorld VM.

Reads the RDDS source from ./rdds/ (local copy), packages it as a tarball,
uploads it to the VM via the Flask API on port 5000, installs deps, and
starts the uvicorn server on the specified port.

Usage:
    from install_rdds import install_rdds
    install_rdds(env, rdds_port=8087)
"""

import base64
import logging
import tarfile
import time
from io import BytesIO
from pathlib import Path

import boto3
import requests

from desktop_env.desktop_env import DesktopEnv

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

_THIS_DIR = Path(__file__).resolve().parent
RDDS_PATH = _THIS_DIR / "rdds" / "hai_drivers" / "desktop" / "remote_desktop_driver_server"
INTERFACES_PATH = _THIS_DIR / "rdds" / "hai_drivers" / "interfaces"

# --- Scripts that run ON the VM via execute_python_command ---

_SCRIPT_INSTALL_UV = """
import subprocess
import os

home = os.path.expanduser("~")
subprocess.run(["wget", "-q", "-O", "/tmp/uv_install.sh", "https://astral.sh/uv/install.sh"], timeout=60)
subprocess.run(["bash", "/tmp/uv_install.sh"], capture_output=True, timeout=120, env={**os.environ, "HOME": home})

uv_path = f"{home}/.local/bin/uv"
if os.path.exists(uv_path):
    print("UV_INSTALLED")
else:
    print("UV_INSTALL_FAILED")
"""

_SCRIPT_CREATE_DIRS = """
import os
home = os.path.expanduser("~")
os.makedirs(f"{home}/hai/hai_drivers/desktop/remote_desktop_driver_server", exist_ok=True)
os.makedirs(f"{home}/hai/hai_drivers/interfaces", exist_ok=True)
print("DIRS_CREATED")
"""

_SCRIPT_UPLOAD_RDDS = """
import base64, tarfile, io, os
os.chdir(os.path.expanduser("~/hai/hai_drivers/desktop"))
data = base64.b64decode("{tarball_b64}")
with tarfile.open(fileobj=io.BytesIO(data), mode='r:gz') as tar:
    tar.extractall()
print("RDDS_EXTRACTED")
"""

_SCRIPT_UPLOAD_INTERFACES = """
import base64, tarfile, io, os
os.chdir(os.path.expanduser("~/hai/hai_drivers"))
data = base64.b64decode("{tarball_b64}")
with tarfile.open(fileobj=io.BytesIO(data), mode='r:gz') as tar:
    tar.extractall()
print("INTERFACES_EXTRACTED")
"""

_SCRIPT_PATCH_PYPROJECTS = """
import os

rdds_toml = os.path.expanduser("~/hai/hai_drivers/desktop/remote_desktop_driver_server/pyproject.toml")
with open(rdds_toml) as f:
    content = f.read()
content = content.replace(
    'hai-drivers-interfaces = {workspace = true}',
    'hai-drivers-interfaces = {path = "../../interfaces"}',
)
with open(rdds_toml, "w") as f:
    f.write(content)

iface_toml = os.path.expanduser("~/hai/hai_drivers/interfaces/pyproject.toml")
with open(iface_toml) as f:
    content = f.read()
content = content.replace('readme = "README.md"' + chr(10), '')
with open(iface_toml, "w") as f:
    f.write(content)

print("PYPROJECTS_PATCHED")
"""

_SCRIPT_UV_SYNC = """
import subprocess, os
home = os.path.expanduser("~")
uv_path = f"{home}/.local/bin/uv"
venv_path = f"{home}/.venv"
project_path = f"{home}/hai/hai_drivers/desktop/remote_desktop_driver_server"
log_path = "/tmp/uv_sync_stderr.log"
env = os.environ.copy()
env["UV_PROJECT_ENVIRONMENT"] = venv_path
env["PATH"] = f"{home}/.local/bin:" + env.get("PATH", "")
with open(log_path, "w") as stderr_file:
    result = subprocess.run([uv_path, "sync", "--project", project_path], stdout=subprocess.PIPE, stderr=stderr_file, text=True, timeout=180, env=env)
print("UV_SYNC_RC:", result.returncode)
if result.returncode != 0:
    with open(log_path) as f:
        lines = f.read()
    print("UV_SYNC_STDERR_TAIL:", lines[-2000:])
"""

_SCRIPT_OPEN_FIREWALL = """
import subprocess
result = subprocess.run("echo 'osworld-public-evaluation' | sudo -S ufw allow {port}/tcp",
                        shell=True, capture_output=True, text=True)
if result.returncode == 0:
    print("FIREWALL_OPENED", result.stdout.strip())
else:
    print("FIREWALL_FAILED", f"rc={{result.returncode}}", result.stdout.strip(), result.stderr.strip())
"""

_SCRIPT_KILL_PORT = """
import subprocess
subprocess.run(["fuser", "-k", "{port}/tcp"], capture_output=True, text=True)
print("PORT_CLEARED")
"""

_SCRIPT_START_SERVER = '''
import subprocess, os, stat
home = os.path.expanduser("~")
uvicorn_path = f"{{home}}/.venv/bin/uvicorn"
rdds_dir = f"{{home}}/hai/hai_drivers/desktop/remote_desktop_driver_server"
launcher = f"""#!/bin/bash
export DISPLAY=:0
export HOME={{home}}
cd {{rdds_dir}}
nohup {{uvicorn_path}} remote_desktop_driver_server:app --host 0.0.0.0 --port {port} > /tmp/rdds.log 2>&1 &
disown
"""
with open("/tmp/start_rdds.sh", "w") as f:
    f.write(launcher)
os.chmod("/tmp/start_rdds.sh", stat.S_IRWXU)
subprocess.run(["bash", "/tmp/start_rdds.sh"], timeout=10)
print("SERVER_STARTED")
'''

_SCRIPT_DIAGNOSE_RDDS = """
import subprocess, os
result = subprocess.run(["pgrep", "-af", "uvicorn"], capture_output=True, text=True)
print("UVICORN_PROCS:", result.stdout.strip() or "<none>")
result = subprocess.run(["ss", "-tlnp"], capture_output=True, text=True)
print("LISTENING_PORTS:", result.stdout.strip())
log_path = "/tmp/rdds.log"
if os.path.exists(log_path):
    with open(log_path) as f:
        content = f.read()
    print("RDDS_LOG:", content[-2000:])
else:
    print("RDDS_LOG: <file not found>")
"""


def _create_tarball(source_path: Path, package_name: str, exclude_patterns: list[str]) -> bytes:
    def should_exclude(tarinfo: tarfile.TarInfo) -> tarfile.TarInfo | None:
        name = tarinfo.name
        if "__pycache__" in name:
            return None
        for pattern in exclude_patterns:
            if pattern.startswith("*.") and name.endswith(pattern[1:]):
                return None
            elif pattern in name:
                return None
        return tarinfo

    buffer = BytesIO()
    with tarfile.open(fileobj=buffer, mode="w:gz") as tar:
        tar.add(source_path, arcname=package_name, filter=should_exclude)
    buffer.seek(0)
    return buffer.read()


def _run_on_vm(env: DesktopEnv, code: str, desc: str, expected_marker: str = "") -> str:
    logger.info(f"RDDS setup: {desc}")
    result = env.controller.execute_python_command(code)
    if isinstance(result, dict):
        output = result["output"] if "output" in result else (result["message"] if "message" in result else str(result))
    else:
        output = str(result)
    logger.info(f"  -> {output[:200]}")
    if expected_marker:
        assert expected_marker in output, f"RDDS step failed ({desc}): expected '{expected_marker}' in output: {output[:500]}"
    return output


def _resolve_public_ip(instance_id: str, region: str) -> str:
    """Get the public IP of an EC2 instance. AGP runs outside the VPC so it
    needs the internet-routable address, not the private 10.x.x.x."""
    ec2 = boto3.client("ec2", region_name=region)
    resp = ec2.describe_instances(InstanceIds=[instance_id])
    instance = resp["Reservations"][0]["Instances"][0]
    public_ip = instance["PublicIpAddress"]
    logger.info(f"Resolved public IP for {instance_id}: {public_ip}")
    return public_ip


def install_rdds(env: DesktopEnv, rdds_port: int) -> str:
    """Install RDDS on a running OSWorld VM and return the server URL.

    Returns the public-IP based URL so AGP agents (outside the VPC) can reach it.
    """
    logger.info(
        f"[RDDS-DEBUG] install_rdds called: vm_ip={env.vm_ip}, "
        f"path_to_vm={env.path_to_vm}, region={env.region}, port={rdds_port}"
    )

    assert RDDS_PATH.exists(), f"RDDS source not found at {RDDS_PATH}"
    assert INTERFACES_PATH.exists(), f"Interfaces source not found at {INTERFACES_PATH}"

    _run_on_vm(env, _SCRIPT_INSTALL_UV, "Installing uv...", expected_marker="UV_INSTALLED")
    _run_on_vm(env, _SCRIPT_CREATE_DIRS, "Creating directories...", expected_marker="DIRS_CREATED")

    logger.info(f"Packaging RDDS from {RDDS_PATH}...")
    rdds_b64 = base64.b64encode(_create_tarball(RDDS_PATH, "remote_desktop_driver_server", [])).decode()
    _run_on_vm(env, _SCRIPT_UPLOAD_RDDS.format(tarball_b64=rdds_b64), "Uploading RDDS...", expected_marker="RDDS_EXTRACTED")

    logger.info(f"Packaging interfaces from {INTERFACES_PATH}...")
    interfaces_b64 = base64.b64encode(
        _create_tarball(INTERFACES_PATH, "interfaces", ["*.pdf", "assets"])
    ).decode()
    _run_on_vm(env, _SCRIPT_UPLOAD_INTERFACES.format(tarball_b64=interfaces_b64), "Uploading interfaces...", expected_marker="INTERFACES_EXTRACTED")

    _run_on_vm(env, _SCRIPT_PATCH_PYPROJECTS, "Patching pyproject.toml files for VM...", expected_marker="PYPROJECTS_PATCHED")

    uv_output = _run_on_vm(env, _SCRIPT_UV_SYNC, "Running uv sync...")
    assert "UV_SYNC_RC: 0" in uv_output, f"uv sync failed: {uv_output[-2000:]}"

    _run_on_vm(env, _SCRIPT_OPEN_FIREWALL.format(port=rdds_port), f"Opening UFW firewall for port {rdds_port}...", expected_marker="FIREWALL_OPENED")
    _run_on_vm(env, _SCRIPT_KILL_PORT.format(port=rdds_port), f"Killing any process on port {rdds_port}...")
    _run_on_vm(env, _SCRIPT_START_SERVER.format(port=rdds_port), "Starting RDDS server...", expected_marker="SERVER_STARTED")

    time.sleep(3)

    # Health-check via private IP (reachable from the benchmark runner in the VPC)
    private_url = f"http://{env.vm_ip}:{rdds_port}"
    _wait_for_rdds(env, private_url, rdds_port)
    logger.info(f"[RDDS-DEBUG] Health check passed at {private_url}")

    # Return public IP URL — AGP agents run outside the VPC
    public_ip = _resolve_public_ip(env.path_to_vm, env.region)
    public_url = f"http://{public_ip}:{rdds_port}"
    logger.info(
        f"[RDDS-DEBUG] FINAL URL: private={private_url} | public={public_url} | "
        f"vm_ip_is_private={'10.' in env.vm_ip or '172.' in env.vm_ip}"
    )
    return public_url


def _wait_for_rdds(env: DesktopEnv, server_url: str, rdds_port: int, timeout: int = 300) -> None:
    logger.info(f"Waiting for RDDS at {server_url}/health ...")
    start = time.time()
    diagnosed = False
    while time.time() - start < timeout:
        elapsed = int(time.time() - start)
        try:
            resp = requests.get(f"{server_url}/health", timeout=5)
            if resp.status_code == 200:
                logger.info(f"RDDS ready after {elapsed}s!")
                return
        except Exception:
            pass
        # Dump VM-side diagnostics once after 30s to catch early failures
        if elapsed > 30 and not diagnosed:
            diagnosed = True
            diag = _run_on_vm(env, _SCRIPT_DIAGNOSE_RDDS, "Diagnosing RDDS failure...")
            logger.error(f"[RDDS-DIAG] after {elapsed}s:\n{diag}")
        logger.info(f"  Not ready yet ({elapsed}s elapsed)...")
        time.sleep(5)
    # Final diagnostic dump before raising
    diag = _run_on_vm(env, _SCRIPT_DIAGNOSE_RDDS, "Final RDDS diagnostics...")
    logger.error(f"[RDDS-DIAG] final:\n{diag}")
    raise RuntimeError(f"RDDS not available after {timeout}s at {server_url}")


def wait_for_flask(env: DesktopEnv, timeout: int = 300) -> None:
    """Wait for the Flask server on :5000 to be ready after VM boot."""
    server_url = f"http://{env.vm_ip}:5000"
    logger.info(f"Waiting for Flask at {server_url}/screenshot ...")
    start = time.time()
    while time.time() - start < timeout:
        elapsed = int(time.time() - start)
        try:
            resp = requests.get(f"{server_url}/screenshot", timeout=5)
            if resp.status_code == 200:
                logger.info(f"Flask ready after {elapsed}s!")
                return
        except Exception:
            pass
        logger.info(f"  Not ready yet ({elapsed}s elapsed)...")
        time.sleep(10)
    raise RuntimeError(f"Flask not available after {timeout}s at {server_url}")
