"""Wait for SSH and run remote commands."""

from __future__ import annotations

from collections.abc import Callable
from pathlib import Path
import shlex
import socket
import subprocess
import time

import paramiko

from gpu_rent.config import Config
from gpu_rent.errors import CloudError


def split_ssh_stream(buf: str) -> tuple[list[tuple[str, str]], str]:
    """Split SSH/PTY stdout into (kind, text) events.

    kind is ``line`` (ended with \\n or \\r\\n) or ``progress`` (ended with lone \\r).
    Returns (events, remainder).
    """
    events: list[tuple[str, str]] = []
    i = 0
    start = 0
    n = len(buf)
    while i < n:
        ch = buf[i]
        if ch == "\r":
            if i + 1 < n and buf[i + 1] == "\n":
                events.append(("line", buf[start:i]))
                i += 2
                start = i
                continue
            events.append(("progress", buf[start:i]))
            i += 1
            start = i
            continue
        if ch == "\n":
            events.append(("line", buf[start:i]))
            i += 1
            start = i
            continue
        i += 1
    return events, buf[start:]


def feed_ssh_log(
    log: Callable[[str], None] | None,
    kind: str,
    text: str,
) -> None:
    """Dispatch a stream event to the CLI log callback."""
    if not log or not text:
        # Empty progress is a no-op; empty line still logs blank via caller if needed.
        if log and kind == "line" and text == "":
            log("")
        return
    if kind == "progress":
        log("\r" + text)
    else:
        log(text)


def _stream_pty_output(
    stdout,
    *,
    log: Callable[[str], None] | None,
    chunks: list[str],
) -> None:
    """Read PTY stdout in chunks; honor \\r progress and \\n lines."""
    buf = ""
    channel = getattr(stdout, "channel", None)
    if channel is None:
        while True:
            piece = stdout.read(4096)
            if not piece:
                break
            text = piece.decode("utf-8", errors="replace") if isinstance(piece, bytes) else piece
            chunks.append(text)
            buf += text
            events, buf = split_ssh_stream(buf)
            for kind, part in events:
                feed_ssh_log(log, kind, part)
        if buf:
            feed_ssh_log(log, "line", buf.rstrip("\r"))
        return

    while True:
        if channel.recv_ready():
            data = channel.recv(4096)
            if not data:
                break
            text = data.decode("utf-8", errors="replace")
            chunks.append(text)
            buf += text
            events, buf = split_ssh_stream(buf)
            for kind, part in events:
                feed_ssh_log(log, kind, part)
            continue
        if channel.exit_status_ready():
            while channel.recv_ready():
                data = channel.recv(4096)
                if not data:
                    break
                text = data.decode("utf-8", errors="replace")
                chunks.append(text)
                buf += text
                events, buf = split_ssh_stream(buf)
                for kind, part in events:
                    feed_ssh_log(log, kind, part)
            break
        time.sleep(0.05)
    if buf:
        feed_ssh_log(log, "line", buf.rstrip("\r"))


def wait_tcp(host: str, port: int, timeout: float = 300.0) -> None:
    deadline = time.time() + timeout
    last = None
    while time.time() < deadline:
        try:
            with socket.create_connection((host, port), timeout=8):
                return
        except OSError as exc:
            last = exc
            time.sleep(4)
    raise CloudError(f"TCP {host}:{port} не открылся за {int(timeout)} с ({last})")


def probe_ssh(cfg: Config, host: str, attempts: int = 3) -> str:
    """Quick SSH check. Returns 'ok' | 'auth' | 'down'."""
    import logging

    logging.getLogger("paramiko").setLevel(logging.CRITICAL)
    logging.getLogger("paramiko.transport").setLevel(logging.CRITICAL)
    try:
        with socket.create_connection((host, 22), timeout=8):
            pass
    except OSError:
        return "down"

    key = str(cfg.ssh_private_key_path)
    auth_seen = False
    for _ in range(max(1, attempts)):
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            client.connect(
                hostname=host,
                username=cfg.ssh_user,
                key_filename=key,
                timeout=12,
                banner_timeout=20,
                auth_timeout=15,
                allow_agent=False,
                look_for_keys=False,
            )
            client.close()
            return "ok"
        except Exception as exc:
            name = type(exc).__name__
            if "Authentication" in name or "Authentication" in str(exc):
                auth_seen = True
            try:
                client.close()
            except Exception:
                pass
            time.sleep(2)
    return "auth" if auth_seen else "down"


def wait_ssh(
    cfg: Config,
    host: str,
    timeout: float = 900.0,
    log: Callable[[str], None] | None = None,
) -> None:
    """Wait until sshd accepts our key. Paramiko banner noise is muted.

    AuthenticationException is normal while cloud-init injects keys: keep
    retrying. Only give up early after AUTH_GIVE_UP seconds of *continuous*
    auth rejection (sshd up, key still wrong) so recreate-with-user_data can run.
    """
    import logging

    logging.getLogger("paramiko").setLevel(logging.CRITICAL)
    logging.getLogger("paramiko.transport").setLevel(logging.CRITICAL)

    emit = log if log is not None else print

    # sshd may answer before authorized_keys is ready (boot-from-volume / user_data).
    auth_give_up = 180.0

    wait_tcp(host, 22, timeout=min(timeout, 300))
    deadline = time.time() + timeout
    key = str(cfg.ssh_private_key_path)
    last = None
    attempt = 0
    auth_streak_started: float | None = None
    while time.time() < deadline:
        attempt += 1
        client = paramiko.SSHClient()
        client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        try:
            client.connect(
                hostname=host,
                username=cfg.ssh_user,
                key_filename=key,
                timeout=20,
                banner_timeout=60,
                auth_timeout=30,
                allow_agent=False,
                look_for_keys=False,
            )
            client.close()
            return
        except Exception as exc:
            last = exc
            name = type(exc).__name__
            is_auth = "Authentication" in name or "Authentication" in str(exc)
            if is_auth:
                if auth_streak_started is None:
                    auth_streak_started = time.time()
            else:
                auth_streak_started = None
            try:
                client.close()
            except Exception:
                pass
            if attempt == 1 or attempt % 6 == 0:
                emit(f"жду SSH {cfg.ssh_user}@{host}… ({name})")
            if (
                is_auth
                and auth_streak_started is not None
                and (time.time() - auth_streak_started) >= auth_give_up
            ):
                raise CloudError(
                    f"SSH {cfg.ssh_user}@{host}: ключ отклонён {int(auth_give_up)}с подряд. "
                    "Nova keypair / user_data не попал в authorized_keys. "
                    "up пересоздаст compute с user_data (Base64), диски оставит."
                ) from exc
            time.sleep(5)
    raise CloudError(
        f"SSH {cfg.ssh_user}@{host} не принял ключ за {int(timeout)} с ({last}). "
        "Часто cloud-init ещё поднимает sshd на GPU-образе. "
        "Если SG узкий (VPN/WARP): в .env GPU_RENT_SSH_CIDR=0.0.0.0/0 и снова up."
    )


def ssh_argv(cfg: Config, host: str, remote: list[str] | None = None) -> list[str]:
    cmd = [
        "ssh",
        "-i",
        str(cfg.ssh_private_key_path),
        "-o",
        "StrictHostKeyChecking=accept-new",
        "-o",
        "IdentitiesOnly=yes",
        f"{cfg.ssh_user}@{host}",
    ]
    if remote:
        cmd.append(" ".join(remote))
    return cmd


def interactive_ssh(cfg: Config, host: str) -> int:
    argv = ssh_argv(cfg, host)
    try:
        return subprocess.call(argv)
    except FileNotFoundError:
        raise CloudError(
            "Нет клиента ssh в PATH. Windows: установи OpenSSH Client "
            f"или зайди так: ssh -i {cfg.ssh_private_key_path} {cfg.ssh_user}@{host}"
        ) from None


def run_ssh(cfg: Config, host: str, command: str, timeout: int = 60, check: bool = True) -> str:
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    try:
        client.connect(
            hostname=host,
            username=cfg.ssh_user,
            key_filename=str(cfg.ssh_private_key_path),
            timeout=15,
        )
        _stdin, stdout, stderr = client.exec_command(command, timeout=timeout)
        out = stdout.read().decode("utf-8", errors="replace")
        err = stderr.read().decode("utf-8", errors="replace")
        code = stdout.channel.recv_exit_status()
        if check and code != 0:
            raise CloudError(f"SSH `{command}` exit {code}: {err or out}")
        return out
    finally:
        client.close()


def run_script_sudo(
    cfg: Config,
    host: str,
    script: str,
    *,
    remote_path: str,
    timeout: int = 1800,
    env: dict[str, str] | None = None,
    log: Callable[[str], None] | None = None,
) -> str:
    """Upload a script and run it with passwordless sudo, streaming stdout."""
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    chunks: list[str] = []
    try:
        client.connect(
            hostname=host,
            username=cfg.ssh_user,
            key_filename=str(cfg.ssh_private_key_path),
            timeout=20,
        )
        sftp = client.open_sftp()
        with sftp.file(remote_path, "w") as fh:
            fh.write(script)
        sftp.chmod(remote_path, 0o755)
        sftp.close()
        env_s = " ".join(
            f"{key}={shlex.quote(str(value))}" for key, value in (env or {}).items()
        )
        prefix = f"sudo -n env {env_s} " if env_s else "sudo -n "
        command = f"{prefix}bash {remote_path}"
        _stdin, stdout, stderr = client.exec_command(command, timeout=timeout, get_pty=True)
        _stream_pty_output(stdout, log=log, chunks=chunks)
        code = stdout.channel.recv_exit_status()
        err = stderr.read().decode("utf-8", errors="replace") if not stdout.channel.closed else ""
        # PTY usually merges stderr; if anything remains, don't drop it past timed log.
        if log and err.strip():
            for line in err.strip().splitlines():
                log(line)
        out = "".join(chunks)
        if code != 0:
            raise CloudError(f"remote script exit {code}: {err or out[-2000:]}")
        return out
    finally:
        client.close()


def run_python(
    cfg: Config,
    host: str,
    script: str,
    *,
    remote_path: str,
    timeout: int = 1800,
    log: Callable[[str], None] | None = None,
) -> str:
    """Upload a Python script and run it as SSH user (not root)."""
    client = _connect(cfg, host)
    chunks: list[str] = []
    try:
        sftp = client.open_sftp()
        with sftp.file(remote_path, "w") as fh:
            fh.write(script)
        sftp.chmod(remote_path, 0o755)
        sftp.close()
        command = f"python3 {shlex.quote(remote_path)}"
        _stdin, stdout, stderr = client.exec_command(command, timeout=timeout, get_pty=True)
        _stream_pty_output(stdout, log=log, chunks=chunks)
        code = stdout.channel.recv_exit_status()
        err = stderr.read().decode("utf-8", errors="replace") if not stdout.channel.closed else ""
        if log and err.strip():
            for line in err.strip().splitlines():
                log(line)
        out = "".join(chunks)
        if code != 0:
            raise CloudError(f"remote python exit {code}: {err or out[-2000:]}")
        return out
    finally:
        client.close()


def _connect(cfg: Config, host: str) -> paramiko.SSHClient:
    client = paramiko.SSHClient()
    client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
    client.connect(
        hostname=host,
        username=cfg.ssh_user,
        key_filename=str(cfg.ssh_private_key_path),
        timeout=20,
    )
    return client


def open_ssh(cfg: Config, host: str) -> paramiko.SSHClient:
    """Public alias for a connected SSH client (caller must close)."""
    return _connect(cfg, host)


def put_text(cfg: Config, host: str, remote_path: str, text: str, mode: int = 0o644) -> None:
    client = _connect(cfg, host)
    try:
        put_text_on(client, remote_path, text, mode=mode)
    finally:
        client.close()


def put_text_on(
    client: paramiko.SSHClient, remote_path: str, text: str, mode: int = 0o644
) -> None:
    sftp = client.open_sftp()
    try:
        _sftp_mkdirs(sftp, str(Path(remote_path).parent).replace("\\", "/"))
        with sftp.file(remote_path, "w") as fh:
            fh.write(text)
        sftp.chmod(remote_path, mode)
    finally:
        sftp.close()


def put_file(cfg: Config, host: str, local: Path, remote_path: str) -> None:
    client = _connect(cfg, host)
    try:
        put_file_on(client, local, remote_path)
    finally:
        client.close()


def put_file_on(client: paramiko.SSHClient, local: Path, remote_path: str) -> None:
    sftp = client.open_sftp()
    try:
        _sftp_mkdirs(sftp, str(Path(remote_path).parent).replace("\\", "/"))
        sftp.put(str(local), remote_path)
    finally:
        sftp.close()


def get_file(cfg: Config, host: str, remote_path: str, local: Path) -> None:
    local.parent.mkdir(parents=True, exist_ok=True)
    client = _connect(cfg, host)
    try:
        get_file_on(client, remote_path, local)
    finally:
        client.close()


def get_file_on(client: paramiko.SSHClient, remote_path: str, local: Path) -> None:
    local.parent.mkdir(parents=True, exist_ok=True)
    sftp = client.open_sftp()
    try:
        sftp.get(remote_path, str(local))
    finally:
        sftp.close()


def remote_exists(cfg: Config, host: str, remote_path: str) -> bool:
    client = _connect(cfg, host)
    try:
        sftp = client.open_sftp()
        try:
            sftp.stat(remote_path)
            return True
        except FileNotFoundError:
            return False
        finally:
            sftp.close()
    finally:
        client.close()


def remote_sha256(cfg: Config, host: str, remote_path: str) -> str | None:
    client = _connect(cfg, host)
    try:
        return remote_sha256_on(client, remote_path)
    finally:
        client.close()


def remote_sha256_on(client: paramiko.SSHClient, remote_path: str) -> str | None:
    _stdin, stdout, stderr = client.exec_command(
        f"sha256sum {shlex.quote(remote_path)} 2>/dev/null | awk '{{print $1}}'",
        timeout=60,
    )
    del stderr
    out = stdout.read().decode("utf-8", errors="replace").strip()
    return out or None


def run_ssh_on(
    client: paramiko.SSHClient, command: str, timeout: int = 60, check: bool = True
) -> str:
    _stdin, stdout, stderr = client.exec_command(command, timeout=timeout)
    out = stdout.read().decode("utf-8", errors="replace")
    err = stderr.read().decode("utf-8", errors="replace")
    code = stdout.channel.recv_exit_status()
    if check and code != 0:
        raise CloudError(f"SSH `{command}` exit {code}: {err or out}")
    return out


def _sftp_mkdirs(sftp, remote_dir: str) -> None:
    if not remote_dir or remote_dir == "/":
        return
    parts = [p for p in remote_dir.split("/") if p]
    cur = ""
    for part in parts:
        cur += "/" + part
        try:
            sftp.stat(cur)
        except FileNotFoundError:
            sftp.mkdir(cur)