# *************************************************************************
#
#  Copyright (c) 2026 - Datatailr Inc.
#  All Rights Reserved.
#
#  This file is part of Datatailr and subject to the terms and conditions
#  defined in 'LICENSE.txt'. Unauthorized copying and/or distribution
#  of this file, in parts or full, via any medium is strictly prohibited.
# *************************************************************************

from __future__ import annotations

import json
import os
import platform
import re
import shlex
import shutil
import stat
import subprocess
import sys
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional
from urllib.parse import urlparse

import requests

from datatailr._remote_client.utils import CONFIG_DIR, Config, username_from_jwt

_WORKSTATION_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_-]*$")
_WEBSOCAT_GITHUB_URL = "https://github.com/vi/websocat"
_WEBSOCAT_RELEASES_URL = f"{_WEBSOCAT_GITHUB_URL}/releases"
_WEBSOCAT_LATEST_DOWNLOAD_URL = f"{_WEBSOCAT_RELEASES_URL}/latest/download"
_API_CALL = Callable[[Config, Dict[str, Any], str, str, List[str], int], Any]


def validate_workstation_name(workstation: str) -> str:
    name = workstation.strip()
    if not name or not _WORKSTATION_RE.match(name):
        raise RuntimeError(
            "Invalid workstation name. Allowed pattern: [A-Za-z0-9][A-Za-z0-9_-]*"
        )
    return name


def _job_rows_from_ls_result(result: Any) -> List[Dict[str, Any]]:
    if isinstance(result, str):
        text = result.strip()
        if not text:
            return []
        try:
            parsed = json.loads(text)
        except Exception:
            return []
        return _job_rows_from_ls_result(parsed)

    if isinstance(result, list):
        return [row for row in result if isinstance(row, dict)]
    if isinstance(result, dict):
        for key in ("jobs", "items", "data", "result"):
            rows = result.get(key)
            if isinstance(rows, list):
                return [row for row in rows if isinstance(row, dict)]
    return []


def _format_memory_bytes(value: Any) -> str:
    if isinstance(value, (int, float)):
        memory = float(value)
    elif isinstance(value, str) and value.strip():
        try:
            memory = float(value.strip())
        except ValueError:
            return value
    else:
        return "unknown"

    gib = memory / (1024.0**3)
    return f"{gib:.2f} GiB"


def _normalize_requirements(value: Any) -> List[str]:
    if isinstance(value, list):
        return [str(item).strip() for item in value if str(item).strip()]
    if not isinstance(value, str):
        return []

    text = value.strip()
    if not text:
        return []

    parts: List[str] = []
    for line in text.splitlines():
        for token in line.split(","):
            item = token.strip()
            if item:
                parts.append(item)
    return parts


def _format_requirements_display(requirements: List[str], max_length: int = 120) -> str:
    if not requirements:
        return "none"

    joined = ", ".join(requirements)
    if len(joined) <= max_length:
        return joined

    if max_length <= 3:
        return "..."

    return joined[: max_length - 3].rstrip() + "..."


def _list_workstation_candidates(
    cfg: Config, api_json: Dict[str, Any], username: str, api_call: _API_CALL
) -> List[Dict[str, Any]]:
    workstation_names: Dict[str, None] = {}
    suffix = f"-{username}"

    result = api_call(
        cfg,
        api_json,
        "job",
        "ls",
        ["-j", "--filter", "type=workstation,environment=dev"],
        30,
    )
    rows = _job_rows_from_ls_result(result)
    for row in rows:
        job_name = row.get("name")
        if not isinstance(job_name, str) or not job_name:
            continue

        if job_name.endswith(suffix) and len(job_name) > len(suffix):
            workstation_name = job_name[: -len(suffix)]
        else:
            workstation_name = job_name

        workstation_names[workstation_name] = None

    candidates: List[Dict[str, Any]] = []
    for workstation_name in sorted(workstation_names):
        details: Dict[str, Any] = {
            "workstation_name": workstation_name,
            "display_name": workstation_name,
            "cpu": None,
            "memory": None,
            "python_requirements": [],
        }

        try:
            result = api_call(
                cfg,
                api_json,
                "job",
                "get",
                [workstation_name, "-j", "-p", "-e", "dev"],
                30,
            )
            rows = _job_rows_from_ls_result(result)
            if rows:
                row = rows[0]
                if row.get("per_user_job") is not True:
                    continue

                image = row.get("image")
                requirements_value = None
                if isinstance(image, dict):
                    requirements_value = image.get("python_requirements")

                display_name = row.get("display_name")
                details["display_name"] = (
                    display_name
                    if isinstance(display_name, str) and display_name
                    else workstation_name
                )
                details["cpu"] = row.get("num_cpus")
                details["memory"] = row.get("memory")
                details["python_requirements"] = _normalize_requirements(
                    requirements_value
                )
        except Exception:
            pass

        candidates.append(details)

    return candidates


def select_workstation_interactively(
    cfg: Config, api_json: Dict[str, Any], username: str, api_call: _API_CALL
) -> str:
    workstations = _list_workstation_candidates(cfg, api_json, username, api_call)
    if not workstations:
        raise RuntimeError(
            "Could not find any workstation jobs in dev for your account. "
            "Pass workstation explicitly: 'datatailr setup-ssh <workstation>'."
        )

    choose_label = "Choose a workstation:"
    if sys.stdout.isatty():
        print(f"\033[4m{choose_label}\033[0m")
    else:
        print(choose_label)
    for index, workstation in enumerate(workstations, start=1):
        display_name = workstation["display_name"]
        name = workstation["workstation_name"]
        cpu = workstation.get("cpu")
        cpu_display = str(cpu) if cpu is not None else "unknown"
        memory_display = _format_memory_bytes(workstation.get("memory"))
        reqs = workstation.get("python_requirements") or []
        reqs_display = _format_requirements_display(reqs)

        print(
            f"  {index}. {display_name} ({name}) | "
            f"cpu={cpu_display} Core | memory={memory_display} | "
            f"python_requirements=[{reqs_display}]"
        )

    while True:
        raw = input(
            f"Select workstation [1-{len(workstations)}], c to cancel: "
        ).strip()
        if not raw:
            continue
        if raw.lower() in {"c", "cancel", "q", "quit", "exit"}:
            raise RuntimeError("Workstation selection cancelled by user.")
        if raw.isdigit():
            choice = int(raw)
            if 1 <= choice <= len(workstations):
                selected = str(workstations[choice - 1]["workstation_name"])
                return selected
        print("Invalid selection. Enter a valid number from the list.")


def _command_exists(command: str) -> bool:
    return shutil.which(command) is not None


def _download_file(url: str, path: Path) -> None:
    resp = requests.get(url, timeout=60)
    if resp.status_code >= 400:
        raise RuntimeError(
            f"Failed to download {url} (HTTP {resp.status_code}): {resp.text[:200]}"
        )
    path.write_bytes(resp.content)


def _resolve_websocat_path() -> Optional[Path]:
    bin_path = shutil.which("websocat")
    if bin_path:
        return Path(bin_path).resolve()

    candidates = [
        Path("/usr/local/bin/websocat"),
        Path.home() / ".local" / "bin" / "websocat",
    ]
    for candidate in candidates:
        if candidate.exists():
            return candidate.resolve()
    return None


def _confirm_websocat_install(method: str, source_url: str) -> None:
    if not sys.stdin.isatty():
        raise RuntimeError(
            "websocat is not installed and setup is running non-interactively. "
            f"Install it manually from {_WEBSOCAT_RELEASES_URL} and rerun 'datatailr setup-ssh'."
        )

    print("websocat is required to tunnel SSH over WebSocket.")
    print(f"Install method: {method}")
    print(f"Source: {source_url}")
    answer = input("Proceed with websocat installation? [y/N]: ").strip().lower()
    if answer not in {"y", "yes"}:
        raise RuntimeError(
            "Cancelled by user. Install websocat manually from "
            f"{_WEBSOCAT_RELEASES_URL} and rerun 'datatailr setup-ssh'."
        )


def install_websocat() -> Path:
    existing = _resolve_websocat_path()
    if existing:
        return existing

    system = platform.system()
    if system == "Darwin":
        if not _command_exists("brew"):
            raise RuntimeError(
                "websocat is not installed and Homebrew is not available. Install Homebrew, "
                f"or install websocat manually from {_WEBSOCAT_RELEASES_URL}, then rerun 'datatailr setup-ssh'."
            )
        _confirm_websocat_install(
            method="brew install websocat",
            source_url=_WEBSOCAT_RELEASES_URL,
        )
        subprocess.run(["brew", "install", "websocat"], check=True)
        installed = _resolve_websocat_path()
        if installed:
            return installed
        raise RuntimeError(
            "websocat install completed but binary was not found on PATH. "
            f"Install it manually from {_WEBSOCAT_RELEASES_URL} and rerun 'datatailr setup-ssh'."
        )

    if system == "Linux":
        arch = platform.machine().lower()
        if arch in {"x86_64", "amd64"}:
            ws_arch = "x86_64-unknown-linux-musl"
        elif arch in {"aarch64", "arm64"}:
            ws_arch = "aarch64-unknown-linux-musl"
        else:
            raise RuntimeError(
                f"Unsupported Linux architecture for websocat: {arch}. "
                f"Install websocat manually from {_WEBSOCAT_RELEASES_URL} and rerun 'datatailr setup-ssh'."
            )

        url = f"{_WEBSOCAT_LATEST_DOWNLOAD_URL}/websocat.{ws_arch}"
        _confirm_websocat_install(
            method=f"Direct download ({ws_arch})",
            source_url=url,
        )

        install_candidates = [Path("/usr/local/bin"), Path.home() / ".local" / "bin"]
        last_error = None
        for install_dir in install_candidates:
            try:
                install_dir.mkdir(parents=True, exist_ok=True)
                target = install_dir / "websocat"
                _download_file(url, target)
                target.chmod(
                    target.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH
                )
                return target.resolve()
            except Exception as exc:  # keep trying fallback location
                last_error = exc

        raise RuntimeError(
            f"Failed to install websocat: {last_error}. "
            f"Install it manually from {_WEBSOCAT_RELEASES_URL} and rerun 'datatailr setup-ssh'."
        )

    if system == "Windows":
        if _command_exists("winget"):
            _confirm_websocat_install(
                method="winget install websocat",
                source_url=_WEBSOCAT_RELEASES_URL,
            )
            subprocess.run(["winget", "install", "websocat"], check=True)
            installed = _resolve_websocat_path()
            if installed:
                return installed
            raise RuntimeError(
                "websocat install completed but binary was not found on PATH. "
                f"Install it manually from {_WEBSOCAT_RELEASES_URL} and rerun 'datatailr setup-ssh'."
            )
        if _command_exists("scoop"):
            _confirm_websocat_install(
                method="scoop install websocat",
                source_url=_WEBSOCAT_RELEASES_URL,
            )
            subprocess.run(["scoop", "install", "websocat"], check=True)
            installed = _resolve_websocat_path()
            if installed:
                return installed
            raise RuntimeError(
                "websocat install completed but binary was not found on PATH. "
                f"Install it manually from {_WEBSOCAT_RELEASES_URL} and rerun 'datatailr setup-ssh'."
            )
        raise RuntimeError(
            "websocat is not installed. Install it with winget or scoop, or manually from "
            f"{_WEBSOCAT_RELEASES_URL}, then rerun 'datatailr setup-ssh'."
        )

    raise RuntimeError(
        f"Unsupported platform for websocat installation: {system}. "
        f"Install websocat manually from {_WEBSOCAT_RELEASES_URL} and rerun 'datatailr setup-ssh'."
    )


def ensure_ssh_key() -> Path:
    if not _command_exists("ssh-keygen"):
        raise RuntimeError(
            "ssh-keygen is not available on this machine. Install OpenSSH client tools first."
        )

    ssh_dir = Path.home() / ".ssh"
    ssh_dir.mkdir(parents=True, exist_ok=True)
    try:
        os.chmod(ssh_dir, 0o700)
    except Exception:
        pass

    private_key = ssh_dir / "datatailr_id_ed25519"
    public_key = ssh_dir / "datatailr_id_ed25519.pub"

    if not private_key.exists():
        subprocess.run(
            ["ssh-keygen", "-t", "ed25519", "-N", "", "-f", str(private_key)],
            check=True,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
    elif not public_key.exists():
        with public_key.open("w", encoding="utf-8") as out_f:
            subprocess.run(
                ["ssh-keygen", "-y", "-f", str(private_key)],
                check=True,
                stdout=out_f,
                stderr=subprocess.DEVNULL,
            )

    try:
        os.chmod(private_key, 0o600)
    except Exception:
        pass
    return private_key


def _write_proxy_helper() -> Path:
    helper_path = CONFIG_DIR / "ssh_proxy_command.py"
    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
    source = Path(__file__).with_name("ssh_proxy_command.py")
    shutil.copy2(source, helper_path)
    helper_path.chmod(0o700)
    return helper_path


def _build_ws_url(base_url: str, workstation: str, username: str) -> str:
    parsed = urlparse(base_url)
    if not parsed.netloc:
        raise RuntimeError(
            f"Invalid base URL in config: '{base_url}'. Run 'datatailr login' to refresh."
        )
    ws_scheme = "wss" if parsed.scheme == "https" else "ws"
    per_user_job_name = f"{workstation}-{username}"
    return f"{ws_scheme}://{parsed.netloc}/workstation/dev/{per_user_job_name}/ssh"


def _upsert_managed_host_block(config_path: Path, workstation: str, block: str) -> None:
    marker_start = f"# >>> Datatailr {workstation} (managed)"
    marker_end = f"# <<< Datatailr {workstation} (managed)"
    host_alias = f"dt-{workstation}"

    current = config_path.read_text(encoding="utf-8") if config_path.exists() else ""

    # Remove previously managed block for this workstation.
    pattern = re.compile(
        rf"(?ms)^\s*{re.escape(marker_start)}\n.*?^\s*{re.escape(marker_end)}\n?"
    )
    cleaned = re.sub(pattern, "", current)

    # Also remove any existing Host block matching dt-<workstation>, even if not managed.
    lines = cleaned.splitlines(keepends=True)
    out_lines: list[str] = []
    i = 0
    while i < len(lines):
        line = lines[i]
        match = re.match(r"^\s*Host\s+(.+?)\s*$", line, flags=re.IGNORECASE)
        if not match:
            out_lines.append(line)
            i += 1
            continue

        host_values = [v.lower() for v in re.split(r"\s+", match.group(1).strip())]
        if host_alias.lower() not in host_values:
            out_lines.append(line)
            i += 1
            continue

        i += 1
        while i < len(lines):
            if re.match(r"^\s*Host\s+.+$", lines[i], flags=re.IGNORECASE):
                break
            i += 1

    cleaned = "".join(out_lines).rstrip("\n")
    new_content = cleaned + ("\n\n" if cleaned else "") + block.strip() + "\n"

    config_path.parent.mkdir(parents=True, exist_ok=True)
    config_path.write_text(new_content, encoding="utf-8")
    try:
        os.chmod(config_path, 0o600)
    except Exception:
        pass


def setup_ssh_config(cfg: Config, workstation: str) -> str:
    workstation = validate_workstation_name(workstation)
    username = username_from_jwt(cfg.jwt_cookie)
    host_alias = f"dt-{workstation}"

    install_websocat()
    key_path = ensure_ssh_key()
    helper_path = _write_proxy_helper()
    ws_url = _build_ws_url(cfg.base_url, workstation, username)
    python_exe = Path(sys.executable).resolve()
    proxy_parts = [
        os.fspath(python_exe),
        os.fspath(helper_path),
        ws_url,
    ]
    if platform.system() == "Windows":
        proxy_command = subprocess.list2cmdline(proxy_parts)
    else:
        proxy_command = " ".join(shlex.quote(part) for part in proxy_parts)

    config_path = Path.home() / ".ssh" / "config"
    block = f"""
# >>> Datatailr {workstation} (managed)
Host {host_alias}
    User {username}
    IdentityFile {key_path}
    IdentitiesOnly yes
    SetEnv LANG=C.UTF-8 LC_CTYPE=C.UTF-8
    ProxyCommand {proxy_command}
    StrictHostKeyChecking no
    UserKnownHostsFile /dev/null
    ServerAliveInterval 30
    ServerAliveCountMax 3
    LogLevel ERROR
# <<< Datatailr {workstation} (managed)
"""
    _upsert_managed_host_block(config_path, workstation, block)
    return host_alias
