"""Remote execution support for Obliteratus.

Run abliteration pipelines on remote GPU nodes via SSH. The remote machine
must have CUDA-capable GPUs and a Python environment. Obliteratus will be
auto-installed if not present.

Usage (CLI):
    obliteratus obliterate meta-llama/Llama-3.1-8B-Instruct \
        --remote user@gpu-node \
        --ssh-key ~/.ssh/id_rsa

Usage (YAML config):
    remote:
      host: gpu-node
      user: root
      ssh_key: ~/.ssh/id_rsa
      remote_dir: /tmp/obliteratus_run
"""

from __future__ import annotations

import os
import shlex
import subprocess
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

from rich.console import Console

console = Console()


@dataclass
class RemoteConfig:
    """SSH connection and remote execution settings."""

    host: str
    user: str = "root"
    port: int = 22
    ssh_key: str | None = None
    remote_dir: str = "/tmp/obliteratus_run"
    install_timeout: int = 600  # seconds
    python: str = "python3"  # remote python binary
    sync_results: bool = True
    gpus: str | None = None  # comma-separated GPU IDs or "all"

    @property
    def ssh_target(self) -> str:
        return f"{self.user}@{self.host}"

    @classmethod
    def from_cli_args(cls, remote_str: str, **kwargs) -> RemoteConfig:
        """Parse 'user@host' or just 'host' from CLI --remote flag."""
        if "@" in remote_str:
            user, host = remote_str.rsplit("@", 1)
        else:
            user = "root"
            host = remote_str
        return cls(host=host, user=user, **kwargs)

    @classmethod
    def from_dict(cls, d: dict) -> RemoteConfig:
        return cls(**{k: v for k, v in d.items() if k in cls.__dataclass_fields__})


class RemoteRunner:
    """Execute Obliteratus commands on a remote machine via SSH."""

    def __init__(
        self,
        config: RemoteConfig,
        on_log: Callable[[str], None] | None = None,
    ):
        self.config = config
        self.on_log = on_log or (lambda msg: console.print(f"[dim][remote][/] {msg}"))

    def _ssh_base_cmd(self) -> list[str]:
        """Build base SSH command with common options."""
        cmd = [
            "ssh",
            "-o", "StrictHostKeyChecking=no",
            "-o", "BatchMode=yes",
            "-o", "ConnectTimeout=30",
            "-p", str(self.config.port),
        ]
        if self.config.ssh_key:
            key_path = os.path.expanduser(self.config.ssh_key)
            cmd.extend(["-i", key_path])
        cmd.append(self.config.ssh_target)
        return cmd

    def _scp_base_cmd(self) -> list[str]:
        """Build base SCP command."""
        cmd = [
            "scp",
            "-o", "StrictHostKeyChecking=no",
            "-o", "BatchMode=yes",
            "-P", str(self.config.port),
            "-r",
        ]
        if self.config.ssh_key:
            key_path = os.path.expanduser(self.config.ssh_key)
            cmd.extend(["-i", key_path])
        return cmd

    def run_ssh(self, remote_cmd: str, stream: bool = False, timeout: int | None = None) -> subprocess.CompletedProcess | int:
        """Run a command on the remote host.

        If stream=True, streams stdout/stderr in real-time and returns the
        exit code. Otherwise returns CompletedProcess.
        """
        cmd = self._ssh_base_cmd() + [remote_cmd]

        if stream:
            proc = subprocess.Popen(
                cmd,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT,
                text=True,
                bufsize=1,
            )
            try:
                for line in proc.stdout:
                    line = line.rstrip("\n")
                    self.on_log(line)
                proc.wait(timeout=timeout)
            except subprocess.TimeoutExpired:
                proc.kill()
                self.on_log("[red]Remote command timed out[/]")
                return 1
            return proc.returncode
        else:
            return subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=timeout,
            )

    def check_connection(self) -> bool:
        """Verify SSH connectivity."""
        self.on_log(f"Testing SSH connection to {self.config.ssh_target}...")
        result = self.run_ssh("echo ok", timeout=30)
        if isinstance(result, subprocess.CompletedProcess) and result.returncode == 0:
            self.on_log("SSH connection OK")
            return True
        self.on_log("[red]SSH connection failed[/]")
        return False

    def check_gpu(self) -> str | None:
        """Check for CUDA GPUs on remote. Returns nvidia-smi output or None."""
        result = self.run_ssh(
            "nvidia-smi --query-gpu=index,name,memory.total,memory.free --format=csv,noheader",
            timeout=30,
        )
        if isinstance(result, subprocess.CompletedProcess) and result.returncode == 0:
            gpu_info = result.stdout.strip()
            lines = gpu_info.split("\n")
            self.on_log(f"Remote GPUs ({len(lines)} detected):")
            for line in lines:
                self.on_log(f"  {line.strip()}")
            if self.config.gpus and self.config.gpus.lower() != "all":
                self.on_log(f"  Selected GPUs: {self.config.gpus}")
            else:
                self.on_log(f"  Using: all {len(lines)} GPUs")
            return gpu_info
        self.on_log("[yellow]No GPUs detected on remote (nvidia-smi failed)[/]")
        return None

    def _env_prefix(self) -> str:
        """Build environment variable prefix for remote commands (e.g. CUDA_VISIBLE_DEVICES)."""
        parts = []
        if self.config.gpus and self.config.gpus.lower() != "all":
            parts.append(f"CUDA_VISIBLE_DEVICES={self.config.gpus}")
        return " ".join(parts) + " " if parts else ""

    def ensure_obliteratus(self) -> bool:
        """Install or update obliteratus on the remote if needed."""
        # Check if already installed
        check = self.run_ssh(
            f"{self.config.python} -c \"import obliteratus; print(obliteratus.__version__)\"",
            timeout=30,
        )
        if isinstance(check, subprocess.CompletedProcess) and check.returncode == 0:
            version = check.stdout.strip()
            self.on_log(f"Obliteratus {version} already installed on remote")
            return True

        # Install from PyPI or git
        self.on_log("Installing obliteratus on remote...")
        install_cmd = (
            f"{self.config.python} -m pip install --quiet "
            f"git+https://github.com/StellaAthena/OBLITERATUS.git"
        )
        rc = self.run_ssh(install_cmd, stream=True, timeout=self.config.install_timeout)
        if rc != 0:
            self.on_log("[red]Failed to install obliteratus on remote[/]")
            return False

        self.on_log("Obliteratus installed successfully")
        return True

    def sync_results_back(self, remote_output_dir: str, local_output_dir: str) -> bool:
        """Copy results from remote back to local machine via scp."""
        local_path = Path(local_output_dir)
        local_path.mkdir(parents=True, exist_ok=True)

        self.on_log(f"Syncing results: {self.config.ssh_target}:{remote_output_dir} -> {local_output_dir}")

        cmd = self._scp_base_cmd() + [
            f"{self.config.ssh_target}:{remote_output_dir}/",
            str(local_path),
        ]

        result = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
        if result.returncode == 0:
            self.on_log(f"Results synced to {local_output_dir}")
            return True
        else:
            self.on_log(f"[red]SCP failed: {result.stderr}[/]")
            return False

    def build_obliterate_command(
        self,
        model: str,
        output_dir: str | None = None,
        method: str = "advanced",
        device: str = "auto",
        dtype: str = "float16",
        quantization: str | None = None,
        n_directions: int | None = None,
        direction_method: str | None = None,
        regularization: float | None = None,
        refinement_passes: int | None = None,
        large_model: bool = False,
        verify_sample_size: int | None = None,
    ) -> str:
        """Build the remote obliteratus CLI command."""
        remote_output = output_dir or f"{self.config.remote_dir}/output/{model.replace('/', '_')}"

        parts = [
            self._env_prefix() + self.config.python, "-m", "obliteratus",
            "obliterate", shlex.quote(model),
            "--output-dir", shlex.quote(remote_output),
            "--method", method,
            "--device", device,
            "--dtype", dtype,
        ]
        if quantization:
            parts.extend(["--quantization", quantization])
        if n_directions is not None:
            parts.extend(["--n-directions", str(n_directions)])
        if direction_method:
            parts.extend(["--direction-method", direction_method])
        if regularization is not None:
            parts.extend(["--regularization", str(regularization)])
        if refinement_passes is not None:
            parts.extend(["--refinement-passes", str(refinement_passes)])
        if large_model:
            parts.append("--large-model")
        if verify_sample_size is not None:
            parts.extend(["--verify-sample-size", str(verify_sample_size)])

        return " ".join(parts)

    def build_run_command(self, remote_config_path: str, output_dir: str | None = None, preset: str | None = None) -> str:
        """Build remote 'obliteratus run' command."""
        parts = [
            self._env_prefix() + self.config.python, "-m", "obliteratus",
            "run", shlex.quote(remote_config_path),
        ]
        if output_dir:
            parts.extend(["--output-dir", shlex.quote(output_dir)])
        if preset:
            parts.extend(["--preset", preset])
        return " ".join(parts)

    def build_tourney_command(
        self,
        model: str,
        output_dir: str | None = None,
        device: str = "auto",
        dtype: str = "float16",
        quantization: str | None = None,
        methods: list[str] | None = None,
        hub_org: str | None = None,
        hub_repo: str | None = None,
        dataset: str = "builtin",
    ) -> str:
        """Build remote 'obliteratus tourney' command."""
        remote_output = output_dir or f"{self.config.remote_dir}/tourney/{model.replace('/', '_')}"

        parts = [
            self._env_prefix() + self.config.python, "-m", "obliteratus",
            "tourney", shlex.quote(model),
            "--output-dir", shlex.quote(remote_output),
            "--device", device,
            "--dtype", dtype,
            "--dataset", dataset,
        ]
        if quantization:
            parts.extend(["--quantization", quantization])
        if hub_org:
            parts.extend(["--hub-org", hub_org])
        if hub_repo:
            parts.extend(["--hub-repo", hub_repo])
        if methods:
            parts.extend(["--methods"] + methods)
        return " ".join(parts)

    def upload_config(self, local_config_path: str) -> str:
        """Upload a YAML config file to the remote."""
        remote_path = f"{self.config.remote_dir}/config.yaml"
        self.run_ssh(f"mkdir -p {shlex.quote(self.config.remote_dir)}")

        cmd = self._scp_base_cmd()
        # scp uses -P not -p, already handled in _scp_base_cmd
        cmd += [local_config_path, f"{self.config.ssh_target}:{remote_path}"]

        result = subprocess.run(cmd, capture_output=True, text=True, timeout=60)
        if result.returncode != 0:
            raise RuntimeError(f"Failed to upload config: {result.stderr}")
        self.on_log(f"Config uploaded to {remote_path}")
        return remote_path

    def run_obliterate(
        self,
        model: str,
        local_output_dir: str | None = None,
        **kwargs,
    ) -> str | None:
        """Full remote obliteration: setup, run, sync results.

        Returns local path to results, or None on failure.
        """
        # 1. Verify connection
        if not self.check_connection():
            return None

        # 2. Check GPUs
        self.check_gpu()

        # 3. Ensure obliteratus is installed
        if not self.ensure_obliteratus():
            return None

        # 4. Create remote working directory
        self.run_ssh(f"mkdir -p {shlex.quote(self.config.remote_dir)}")

        # 5. Build and run the command
        remote_output = f"{self.config.remote_dir}/output/{model.replace('/', '_')}"
        cmd = self.build_obliterate_command(model, output_dir=remote_output, **kwargs)
        self.on_log(f"Running: {cmd}")

        rc = self.run_ssh(cmd, stream=True)
        if rc != 0:
            self.on_log(f"[red]Remote obliteration failed (exit code {rc})[/]")
            return None

        # 6. Sync results back
        if self.config.sync_results:
            local_output = local_output_dir or f"abliterated/{model.replace('/', '_')}"
            if self.sync_results_back(remote_output, local_output):
                return local_output
            return None

        self.on_log(f"Results on remote: {remote_output}")
        return remote_output

    def run_config(
        self,
        local_config_path: str,
        local_output_dir: str | None = None,
        preset: str | None = None,
    ) -> str | None:
        """Upload config, run study remotely, sync results."""
        if not self.check_connection():
            return None
        self.check_gpu()
        if not self.ensure_obliteratus():
            return None

        # Upload config
        remote_config = self.upload_config(local_config_path)

        # Determine remote output dir
        remote_output = f"{self.config.remote_dir}/results"
        cmd = self.build_run_command(remote_config, output_dir=remote_output, preset=preset)
        self.on_log(f"Running: {cmd}")

        rc = self.run_ssh(cmd, stream=True)
        if rc != 0:
            self.on_log(f"[red]Remote run failed (exit code {rc})[/]")
            return None

        if self.config.sync_results:
            local_output = local_output_dir or "results"
            if self.sync_results_back(remote_output, local_output):
                return local_output
            return None

        return remote_output

    def run_tourney(
        self,
        model: str,
        local_output_dir: str | None = None,
        **kwargs,
    ) -> str | None:
        """Run tournament remotely, sync results."""
        if not self.check_connection():
            return None
        self.check_gpu()
        if not self.ensure_obliteratus():
            return None

        remote_output = f"{self.config.remote_dir}/tourney/{model.replace('/', '_')}"
        cmd = self.build_tourney_command(model, output_dir=remote_output, **kwargs)
        self.on_log(f"Running: {cmd}")

        rc = self.run_ssh(cmd, stream=True)
        if rc != 0:
            self.on_log(f"[red]Remote tourney failed (exit code {rc})[/]")
            return None

        if self.config.sync_results:
            local_output = local_output_dir or f"/tmp/obliteratus_tourney/{model.replace('/', '_')}"
            if self.sync_results_back(remote_output, local_output):
                return local_output
            return None

        return remote_output