"""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 queue
import shlex
import subprocess
import tempfile
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Callable

from rich.console import Console
import yaml

from obliteratus import __version__
from obliteratus.remote_contracts import (
    parse_remote_target,
    remote_python_command,
    remote_scp_spec,
    validate_remote_settings,
)

console = Console()


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

    host: str
    user: str = "root"
    port: int = 22
    ssh_key: str | None = None
    known_hosts_file: 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"
    install_source: str = "git+https://github.com/elder-plinius/OBLITERATUS.git"

    def __post_init__(self) -> None:
        self.gpus = validate_remote_settings(
            host=self.host,
            user=self.user,
            port=self.port,
            remote_dir=self.remote_dir,
            python=self.python,
            gpus=self.gpus,
            install_source=self.install_source,
        )
        if (
            isinstance(self.install_timeout, bool)
            or not isinstance(self.install_timeout, int)
            or self.install_timeout <= 0
        ):
            raise ValueError("remote install timeout must be a positive integer")

    @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."""
        user, host = parse_remote_target(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=yes",
            "-o", "BatchMode=yes",
            "-o", "ConnectTimeout=30",
            "-p", str(self.config.port),
        ]
        if self.config.known_hosts_file:
            known_hosts = os.path.expanduser(self.config.known_hosts_file)
            cmd.extend(["-o", f"UserKnownHostsFile={known_hosts}"])
        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=yes",
            "-o", "BatchMode=yes",
            "-P", str(self.config.port),
            "-r",
        ]
        if self.config.known_hosts_file:
            known_hosts = os.path.expanduser(self.config.known_hosts_file)
            cmd.extend(["-o", f"UserKnownHostsFile={known_hosts}"])
        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,
            )
            if proc.stdout is None:
                proc.kill()
                proc.wait()
                raise RuntimeError("remote process did not expose a stdout stream")

            output: queue.Queue[object] = queue.Queue()
            finished = object()

            def read_output() -> None:
                try:
                    for line in proc.stdout:
                        if not isinstance(line, str):
                            raise RuntimeError("remote process emitted non-text output")
                        output.put(line.rstrip("\n"))
                except BaseException as exc:
                    output.put(exc)
                finally:
                    output.put(finished)

            threading.Thread(target=read_output, daemon=True).start()
            deadline = time.monotonic() + timeout if timeout is not None else None
            try:
                while True:
                    remaining = None if deadline is None else max(0, deadline - time.monotonic())
                    try:
                        item = output.get(timeout=remaining)
                    except queue.Empty as exc:
                        raise subprocess.TimeoutExpired(cmd, timeout) from exc
                    if item is finished:
                        break
                    if isinstance(item, BaseException):
                        raise item
                    self.on_log(item)
                remaining = None if deadline is None else max(0, deadline - time.monotonic())
                proc.wait(timeout=remaining)
            except subprocess.TimeoutExpired:
                proc.kill()
                proc.wait()
                self.on_log("[red]Remote command timed out[/]")
                return 124
            except KeyboardInterrupt:
                proc.kill()
                proc.wait()
                self.on_log("[yellow]Remote command cancelled[/]")
                raise
            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 ensure_obliteratus(self) -> bool:
        """Install or update obliteratus on the remote if needed."""
        version_command = remote_python_command(
            self.config.python,
            ["-c", "import obliteratus; print(obliteratus.__version__)"],
        )
        check = self.run_ssh(
            version_command,
            timeout=30,
        )
        if (
            isinstance(check, subprocess.CompletedProcess)
            and check.returncode == 0
            and check.stdout.strip() == __version__
        ):
            self.on_log(f"Obliteratus {__version__} already installed on remote")
            return True

        self.on_log(f"Installing Obliteratus {__version__} on remote...")
        install_cmd = remote_python_command(
            self.config.python,
            ["-m", "pip", "install", "--quiet", "--upgrade", self.config.install_source],
        )
        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

        verify = self.run_ssh(version_command, timeout=30)
        if not (
            isinstance(verify, subprocess.CompletedProcess)
            and verify.returncode == 0
            and verify.stdout.strip() == __version__
        ):
            self.on_log("[red]Installed Obliteratus version could not be verified[/]")
            return False
        self.on_log(f"Obliteratus {__version__} 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() + [
            remote_scp_spec(self.config.ssh_target, remote_output_dir, directory=True),
            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,
        gpu_memory_utilization: float | 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,
        min_layer_fraction: float | None = None,
        max_layer_fraction: float | None = None,
        harmless_pc_count: int | None = None,
        shield_concept_count: int | None = None,
        shield_ridge: float | None = None,
        shield_residualize: bool | None = None,
        shield_layer_penalty: float | None = None,
        projection_target: str | None = None,
        projection_row_fraction: float | None = None,
        refusal_max_tokens: int | None = None,
    ) -> str:
        """Build the remote obliteratus CLI command."""
        remote_output = output_dir or f"{self.config.remote_dir}/output/{model.replace('/', '_')}"

        parts = [
            "-m", "obliteratus",
            "obliterate", model,
            "--output-dir", remote_output,
            "--method", method,
            "--device", device,
            "--dtype", dtype,
        ]
        if quantization:
            parts.extend(["--quantization", quantization])
        if gpu_memory_utilization is not None:
            parts.extend(["--gpu-memory-utilization", str(gpu_memory_utilization)])
        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)])
        optional_values = (
            ("--min-layer-fraction", min_layer_fraction),
            ("--max-layer-fraction", max_layer_fraction),
            ("--harmless-pc-count", harmless_pc_count),
            ("--shield-concept-count", shield_concept_count),
            ("--shield-ridge", shield_ridge),
            ("--shield-layer-penalty", shield_layer_penalty),
            ("--projection-target", projection_target),
            ("--projection-row-fraction", projection_row_fraction),
            ("--refusal-max-tokens", refusal_max_tokens),
        )
        for flag, value in optional_values:
            if value is not None:
                parts.extend([flag, str(value)])
        if shield_residualize:
            parts.append("--shield-residualize")

        return remote_python_command(self.config.python, parts, gpus=self.config.gpus)

    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 = [
            "-m", "obliteratus",
            "run", remote_config_path,
        ]
        if output_dir:
            parts.extend(["--output-dir", output_dir])
        if preset:
            parts.extend(["--preset", preset])
        return remote_python_command(self.config.python, parts, gpus=self.config.gpus)

    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 = [
            "-m", "obliteratus",
            "tourney", model,
            "--output-dir", 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 remote_python_command(self.config.python, parts, gpus=self.config.gpus)

    def upload_config(self, local_config_path: str) -> str:
        """Upload a YAML config without recursively redispatching remotely."""
        remote_path = f"{self.config.remote_dir}/config.yaml"
        self.run_ssh(shlex.join(["mkdir", "-p", self.config.remote_dir]))

        source = Path(local_config_path)
        payload = yaml.safe_load(source.read_text(encoding="utf-8"))
        if not isinstance(payload, dict):
            raise ValueError("remote study config must contain a YAML mapping")
        payload.pop("remote", None)

        with tempfile.TemporaryDirectory(prefix="obliteratus-remote-config-") as temp_dir:
            upload_path = Path(temp_dir) / "config.yaml"
            upload_path.write_text(yaml.safe_dump(payload, sort_keys=False), encoding="utf-8")
            cmd = self._scp_base_cmd()
            # scp uses -P not -p, already handled in _scp_base_cmd
            cmd += [str(upload_path), remote_scp_spec(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(shlex.join(["mkdir", "-p", 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