"""Run execution backends for local and remote experiment execution."""

from __future__ import annotations

import json
import shlex
import shutil
import subprocess
import time
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path, PurePosixPath
from typing import Any, ClassVar

from rich.console import Console

from abreka.config import project_root
from abreka.models import Config, RemoteRunConfig


@dataclass
class RunExecutionResult:
    stdout: str
    stderr: str
    returncode: int
    timed_out: bool = False


@dataclass
class LocalRunExecutor:
    project_root: Path
    console: Console

    def execute(self, *, run_py: Path, output_path: Path, timeout_seconds: float) -> RunExecutionResult:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        timed_out = False

        with open(output_path, "w", encoding="utf-8") as output_file:
            proc = subprocess.Popen(
                ["uv", "run", str(run_py)],
                cwd=self.project_root,
                stdout=output_file,
                stderr=subprocess.STDOUT,
                text=True,
            )

            try:
                proc.wait(timeout=timeout_seconds)
            except subprocess.TimeoutExpired:
                timed_out = True
                proc.terminate()
                try:
                    proc.wait(timeout=10)
                except subprocess.TimeoutExpired:
                    proc.kill()
                    proc.wait()
                self.console.print(
                    f"  [yellow]Timed out after {timeout_seconds / 60:.0f}m — capturing partial results[/yellow]"
                )

            output_file.flush()

        return RunExecutionResult(
            stdout="",
            stderr="",
            returncode=proc.returncode or 0,
            timed_out=timed_out,
        )


@dataclass
class RemoteRunExecutor:
    project_root: Path
    remote: RemoteRunConfig
    console: Console
    max_poll_errors: int = 3

    PROJECT_SYNC_RULES: ClassVar[tuple[tuple[str, str], ...]] = (
        ("exclude", ".git/"),
        ("exclude", ".venv/"),
        ("exclude", ".pytest_cache/"),
        ("exclude", ".mypy_cache/"),
        ("exclude", ".ruff_cache/"),
        ("exclude", "__pycache__/"),
        ("include", "abreka/artifacts/"),
        ("include", "abreka/artifacts/shared/"),
        ("include", "abreka/artifacts/shared/**"),
        ("exclude", "abreka/artifacts/**"),
        ("exclude", "abreka/experiments/*/artifacts/"),
    )

    def execute(self, *, run_py: Path, output_path: Path, timeout_seconds: float) -> RunExecutionResult:
        self._check_local_dependencies()

        run_relative_path = run_py.relative_to(self.project_root)
        remote_paths = self._build_remote_paths(run_relative_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)
        metadata_path = output_path.parent / "remote_run.json"

        metadata: dict[str, Any] = {
            "mode": "remote",
            "state": "starting",
            "host": self.remote.host,
            "remote_project_dir": self.remote.project_dir,
            "remote_paths": remote_paths,
            "run_relative_path": run_relative_path.as_posix(),
            "timeout_seconds": timeout_seconds,
            "poll_seconds": self.remote.poll_seconds,
            "bootstrap_command": self.remote.bootstrap_command,
            "startup_command": self.remote.startup_command,
            "teardown_command": self.remote.teardown_command,
            "started_at": _now_iso(),
        }
        self._write_remote_run_metadata(metadata_path, metadata)

        launched = False
        synced_back = False

        try:
            self._run_local_hook("startup", self.remote.startup_command)
            self._sync_project_to_remote()
            self._run_remote_bootstrap()

            remote_pid = self._launch_remote_process(run_relative_path, timeout_seconds, remote_paths)
            launched = True
            metadata.update({
                "state": "running",
                "remote_pid": remote_pid,
                "launched_at": _now_iso(),
            })
            self._write_remote_run_metadata(metadata_path, metadata)

            status = self._poll_remote_run(remote_paths, output_path)
            self._sync_artifacts_from_remote(remote_paths, output_path.parent)
            synced_back = True

            stdout = output_path.read_text(errors="replace") if output_path.is_file() else ""
            metadata.update({
                "state": "completed",
                "finished_at": _now_iso(),
                "last_sync_at": _now_iso(),
                "status": status,
            })
            self._write_remote_run_metadata(metadata_path, metadata)
            return RunExecutionResult(
                stdout=stdout,
                stderr="",
                returncode=int(status.get("returncode", 1)),
                timed_out=bool(status.get("timed_out", False)),
            )
        except Exception as e:
            metadata.update({
                "state": "interrupted" if launched else "failed",
                "finished_at": _now_iso(),
                "error": str(e),
            })
            self._write_remote_run_metadata(metadata_path, metadata)
            raise
        finally:
            if not launched or synced_back:
                self._run_local_hook("teardown", self.remote.teardown_command)

    def inspect_recovery(self, *, remote_paths: dict[str, str], remote_pid: str = "") -> dict[str, Any]:
        self._check_local_dependencies()
        info: dict[str, Any] = {
            "remote_reachable": False,
            "pid_alive": False,
            "status": None,
            "status_exists": False,
            "remote_log_size": 0,
        }
        try:
            info["remote_log_size"] = self._get_remote_file_size(remote_paths["log_path"])
            status = self._read_remote_status(remote_paths["status_path"])
            info["status"] = status
            info["status_exists"] = status is not None
            if remote_pid:
                info["pid_alive"] = self._remote_pid_alive(remote_pid)
            info["remote_reachable"] = True
        except Exception as e:
            info["error"] = str(e)
        return info

    def sync_recovery_artifacts(
        self,
        *,
        remote_paths: dict[str, str],
        local_artifacts_dir: Path,
        metadata_path: Path,
    ) -> None:
        self._check_local_dependencies()
        self._sync_artifacts_from_remote(remote_paths, local_artifacts_dir)
        metadata = self._load_remote_run_metadata(metadata_path)
        metadata["last_sync_at"] = _now_iso()
        self._write_remote_run_metadata(metadata_path, metadata)

    def resume_recovery(
        self,
        *,
        output_path: Path,
        metadata_path: Path,
        remote_paths: dict[str, str],
    ) -> RunExecutionResult:
        self._check_local_dependencies()
        local_status_path = output_path.parent / "remote_status.json"
        local_status = self._read_local_status(local_status_path)
        if local_status is not None:
            return self._finalize_local_recovery_result(output_path, metadata_path, local_status)

        synced_back = False
        metadata = self._load_remote_run_metadata(metadata_path)

        try:
            self._run_local_hook("startup", self.remote.startup_command)

            status = self._read_remote_status(remote_paths["status_path"])
            if status is None:
                start_offset = output_path.stat().st_size if output_path.is_file() else 0
                status = self._poll_remote_run(
                    remote_paths,
                    output_path,
                    start_offset=start_offset,
                    reset_output=not output_path.is_file(),
                )

            self._sync_artifacts_from_remote(remote_paths, output_path.parent)
            synced_back = True
            return self._finalize_local_recovery_result(output_path, metadata_path, status)
        except Exception as e:
            metadata.update({
                "state": "interrupted",
                "finished_at": _now_iso(),
                "error": str(e),
            })
            self._write_remote_run_metadata(metadata_path, metadata)
            raise
        finally:
            if synced_back:
                self._run_local_hook("teardown", self.remote.teardown_command)

    def stop_recovery(
        self,
        *,
        remote_pid: str,
        remote_paths: dict[str, str],
        local_artifacts_dir: Path,
        metadata_path: Path,
        kill: bool = False,
    ) -> dict[str, Any]:
        self._check_local_dependencies()
        if remote_pid:
            self._send_signal(remote_pid, kill=kill)
            if not kill:
                time.sleep(1)

        self._sync_artifacts_from_remote(remote_paths, local_artifacts_dir)

        local_status_path = local_artifacts_dir / "remote_status.json"
        local_status = self._read_local_status(local_status_path)
        if local_status is None:
            local_status = {
                "returncode": 137 if kill else 143,
                "timed_out": False,
                "stopped": True,
            }
            local_status_path.write_text(json.dumps(local_status) + "\n")

        metadata = self._load_remote_run_metadata(metadata_path)
        metadata.update({
            "state": "stopped",
            "finished_at": _now_iso(),
            "last_sync_at": _now_iso(),
            "status": local_status,
            "stop_signal": "SIGKILL" if kill else "SIGTERM",
        })
        self._write_remote_run_metadata(metadata_path, metadata)
        self._run_local_hook("teardown", self.remote.teardown_command)
        return local_status

    def _check_local_dependencies(self) -> None:
        for binary in ("ssh", "rsync"):
            if not shutil.which(binary):
                raise RuntimeError(f"{binary} not found on PATH")

    def _run_local_hook(self, label: str, command: str) -> None:
        if not command:
            return
        self.console.print(f"  [dim]remote {label}: {command}[/dim]")
        subprocess.run(command, cwd=self.project_root, shell=True, check=True)

    def _sync_project_to_remote(self) -> None:
        self._ssh(f"mkdir -p {shlex.quote(self.remote.project_dir)}")
        command = [
            "rsync",
            "-az",
            "--delete",
            *[
                arg
                for action, pattern in self.PROJECT_SYNC_RULES
                for arg in (f"--{action}", pattern)
            ],
            f"{self.project_root}/",
            f"{self.remote.host}:{self.remote.project_dir.rstrip('/')}/",
        ]
        subprocess.run(command, check=True)

    def _run_remote_bootstrap(self) -> None:
        if not self.remote.bootstrap_command:
            return
        self.console.print(f"  [dim]remote bootstrap: {self.remote.bootstrap_command}[/dim]")
        self._ssh(
            f"cd {shlex.quote(self.remote.project_dir)} && {self.remote.bootstrap_command}"
        )

    def _launch_remote_process(
        self,
        run_relative_path: Path,
        timeout_seconds: float,
        remote_paths: dict[str, str],
    ) -> str:
        timeout = max(1, int(timeout_seconds))
        wrapper = "\n".join([
            "set -e",
            f"cd {shlex.quote(self.remote.project_dir)}",
            f"mkdir -p {shlex.quote(remote_paths['artifacts_dir'])}",
            "set +e",
            f"timeout --signal=TERM {timeout} uv run {shlex.quote(run_relative_path.as_posix())}",
            "status=$?",
            "set -e",
            'if [ "$status" -eq 124 ]; then timed_out=true; else timed_out=false; fi',
            (
                "printf '{\"returncode\": %s, \"timed_out\": %s}\\n' "
                f'"$status" "$timed_out" > {shlex.quote(remote_paths["status_path"])}'
            ),
        ])
        command = (
            f"mkdir -p {shlex.quote(self.remote.project_dir)} {shlex.quote(remote_paths['artifacts_dir'])} && "
            f"nohup setsid bash -lc {shlex.quote(wrapper)} > {shlex.quote(remote_paths['log_path'])} 2>&1 < /dev/null & echo $!"
        )
        result = self._ssh(command, capture_output=True, text=True)
        remote_pid = result.stdout.strip()
        if not remote_pid:
            raise RuntimeError("Failed to launch remote run process")
        return remote_pid

    def _poll_remote_run(
        self,
        remote_paths: dict[str, str],
        output_path: Path,
        *,
        start_offset: int = 0,
        reset_output: bool = True,
    ) -> dict[str, Any]:
        offset = start_offset
        poll_errors = 0
        if reset_output or not output_path.is_file():
            output_path.write_bytes(b"")

        while True:
            try:
                chunk = self._read_remote_log_chunk(remote_paths["log_path"], offset)
                if chunk:
                    with open(output_path, "ab") as f:
                        f.write(chunk)
                    self.console.file.write(chunk.decode("utf-8", errors="replace"))
                    self.console.file.flush()
                    offset += len(chunk)

                status = self._read_remote_status(remote_paths["status_path"])
                if status is not None:
                    final_chunk = self._read_remote_log_chunk(remote_paths["log_path"], offset)
                    if final_chunk:
                        with open(output_path, "ab") as f:
                            f.write(final_chunk)
                        self.console.file.write(final_chunk.decode("utf-8", errors="replace"))
                        self.console.file.flush()
                    return status

                poll_errors = 0
                time.sleep(max(1, self.remote.poll_seconds))
            except Exception as e:
                poll_errors += 1
                if poll_errors >= self.max_poll_errors:
                    raise RuntimeError(
                        "Lost contact with remote run while polling. "
                        "See artifacts/remote_run.json for recovery details."
                    ) from e
                time.sleep(max(1, self.remote.poll_seconds))

    def _read_remote_log_chunk(self, log_path: str, offset: int) -> bytes:
        command = (
            f"if [ -f {shlex.quote(log_path)} ]; then "
            f"tail -c +{offset + 1} {shlex.quote(log_path)}; fi"
        )
        result = self._ssh(command, capture_output=True, text=False)
        return result.stdout or b""

    def _read_remote_status(self, status_path: str) -> dict[str, Any] | None:
        command = f"if [ -f {shlex.quote(status_path)} ]; then cat {shlex.quote(status_path)}; fi"
        result = self._ssh(command, capture_output=True, text=True)
        payload = result.stdout.strip()
        if not payload:
            return None
        return json.loads(payload)

    def _sync_artifacts_from_remote(self, remote_paths: dict[str, str], local_artifacts_dir: Path) -> None:
        local_artifacts_dir.mkdir(parents=True, exist_ok=True)
        subprocess.run(
            [
                "rsync",
                "-az",
                f"{self.remote.host}:{remote_paths['artifacts_dir'].rstrip('/')}/",
                f"{local_artifacts_dir}/",
            ],
            check=True,
        )

    def _write_remote_run_metadata(self, metadata_path: Path, metadata: dict[str, Any]) -> None:
        metadata_path.write_text(json.dumps(metadata, indent=2) + "\n")

    def _load_remote_run_metadata(self, metadata_path: Path) -> dict[str, Any]:
        return json.loads(metadata_path.read_text())

    def _build_remote_paths(self, run_relative_path: Path) -> dict[str, str]:
        remote_project_dir = PurePosixPath(self.remote.project_dir)
        remote_run_path = remote_project_dir / PurePosixPath(run_relative_path.as_posix())
        remote_artifacts_dir = remote_run_path.parent / "artifacts"
        return {
            "artifacts_dir": remote_artifacts_dir.as_posix(),
            "log_path": (remote_artifacts_dir / "run_output.txt").as_posix(),
            "status_path": (remote_artifacts_dir / "remote_status.json").as_posix(),
        }

    def _finalize_local_recovery_result(
        self,
        output_path: Path,
        metadata_path: Path,
        status: dict[str, Any],
    ) -> RunExecutionResult:
        stdout = output_path.read_text(errors="replace") if output_path.is_file() else ""
        metadata = self._load_remote_run_metadata(metadata_path)
        metadata.update({
            "state": "completed",
            "finished_at": _now_iso(),
            "last_sync_at": _now_iso(),
            "status": status,
        })
        self._write_remote_run_metadata(metadata_path, metadata)
        return RunExecutionResult(
            stdout=stdout,
            stderr="",
            returncode=int(status.get("returncode", 1)),
            timed_out=bool(status.get("timed_out", False)),
        )

    def _read_local_status(self, status_path: Path) -> dict[str, Any] | None:
        if not status_path.is_file():
            return None
        payload = status_path.read_text().strip()
        if not payload:
            return None
        return json.loads(payload)

    def _get_remote_file_size(self, path: str) -> int:
        result = self._ssh(
            f"if [ -f {shlex.quote(path)} ]; then wc -c < {shlex.quote(path)}; else echo 0; fi",
            capture_output=True,
            text=True,
        )
        try:
            return int(result.stdout.strip() or "0")
        except ValueError:
            return 0

    def _remote_pid_alive(self, remote_pid: str) -> bool:
        pid = self._normalize_pid(remote_pid)
        result = self._ssh(
            f"if kill -0 -- -{pid} 2>/dev/null || kill -0 {pid} 2>/dev/null; then echo true; else echo false; fi",
            capture_output=True,
            text=True,
        )
        return result.stdout.strip() == "true"

    def _send_signal(self, remote_pid: str, *, kill: bool = False) -> None:
        signal_name = "KILL" if kill else "TERM"
        pid = self._normalize_pid(remote_pid)
        self._ssh(
            f"kill -{signal_name} -- -{pid} 2>/dev/null || kill -{signal_name} {pid} 2>/dev/null || true",
            capture_output=True,
            text=True,
        )

    def _normalize_pid(self, remote_pid: str) -> int:
        return int(remote_pid)

    def _ssh(self, command: str, *, capture_output: bool = True, text: bool = False) -> subprocess.CompletedProcess:
        return subprocess.run(
            ["ssh", self.remote.host, "bash", "-lc", command],
            check=True,
            capture_output=capture_output,
            text=text,
        )


def build_run_executor(abreka_dir: Path, config: Config, console: Console) -> LocalRunExecutor | RemoteRunExecutor:
    mode = (config.run.mode or "local").lower()
    proj_root = project_root(abreka_dir)

    if mode == "local":
        return LocalRunExecutor(project_root=proj_root, console=console)
    if mode == "remote":
        if not config.run.remote.host or not config.run.remote.project_dir:
            raise ValueError("remote run mode requires both run.remote.host and run.remote.project_dir")
        return RemoteRunExecutor(project_root=proj_root, remote=config.run.remote, console=console)
    raise ValueError(f"Unknown run mode: {config.run.mode!r}. Expected 'local' or 'remote'.")


def _now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()
