app.py

import asyncio
import base64
import os
import re
import shlex
import shutil
import signal
import socket
import subprocess
import tempfile
import time
import uuid
from pathlib import Path
from typing import Any, Dict, Optional

import httpx
from fastapi import Depends, FastAPI, File, Header, HTTPException, UploadFile, status
from fastapi.responses import FileResponse
from pydantic import BaseModel


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

CLUSTER_LOGIN_HOST = os.environ.get("CLUSTER_LOGIN_HOST", "kraken")
CLUSTER_USER = os.environ.get("CLUSTER_USER", "manolo")
CLUSTER_SSH_PORT = int(os.environ.get("CLUSTER_SSH_PORT", "22"))
CLUSTER_MODEL_DIR = os.environ.get("CLUSTER_MODEL_DIR", "/home/manolo/GeoGlyphFinal/model")
WORKER_INFO_DIR = os.environ.get("WORKER_INFO_DIR", "/home/manolo/geoglyph_workers")
SBATCH_SCRIPT = os.environ.get("SBATCH_SCRIPT", "run_geoglyph_worker_env.sbatch")
WORKER_JOB_NAME = os.environ.get("WORKER_JOB_NAME", "geoglyph-worker")
WORKER_PORT = int(os.environ.get("WORKER_PORT", "8000"))
GATEWAY_LOCAL_PORT = int(os.environ.get("GATEWAY_LOCAL_PORT", "18001"))
GATEWAY_API_TOKEN = os.environ.get("GATEWAY_API_TOKEN", "")

SSH_PRIVATE_KEY = os.environ.get("SSH_PRIVATE_KEY", "")
SSH_PRIVATE_KEY_BASE64 = os.environ.get("SSH_PRIVATE_KEY_BASE64", "")
SSH_PASSWORD = os.environ.get("SSH_PASSWORD", "")
SSH_KNOWN_HOSTS = os.environ.get("SSH_KNOWN_HOSTS", "")

RESULTS_DIR = Path(os.environ.get("RESULTS_DIR", "/tmp/geoglyph_gateway_results"))
RESULTS_DIR.mkdir(parents=True, exist_ok=True)

REQUEST_TIMEOUT = float(os.environ.get("REQUEST_TIMEOUT", "30"))
SEGMENT_TIMEOUT_SECONDS = int(os.environ.get("SEGMENT_TIMEOUT_SECONDS", "1800"))
STATUS_POLL_SECONDS = float(os.environ.get("STATUS_POLL_SECONDS", "2"))
WORKER_START_TIMEOUT_SECONDS = int(os.environ.get("WORKER_START_TIMEOUT_SECONDS", "300"))

SSH_TUNNEL_RETRIES = int(os.environ.get("SSH_TUNNEL_RETRIES", "5"))
SSH_TUNNEL_RETRY_SLEEP_SECONDS = float(os.environ.get("SSH_TUNNEL_RETRY_SLEEP_SECONDS", "4"))

ENSURE_WORKER_RETRY_SECONDS = int(os.environ.get("ENSURE_WORKER_RETRY_SECONDS", "900"))
ENSURE_WORKER_RETRY_SLEEP_SECONDS = float(os.environ.get("ENSURE_WORKER_RETRY_SLEEP_SECONDS", "10"))

CLUSTER_SSH_REFUSED_COOLDOWN_SECONDS = int(
    os.environ.get("CLUSTER_SSH_REFUSED_COOLDOWN_SECONDS", "620")
)

# ---------------------------------------------------------------------------
# Runtime state
# ---------------------------------------------------------------------------

app = FastAPI(
    title="GeoGlyph Gateway API",
    description="Gateway API that tunnels to a Slurm-hosted GeoGlyph SAM2 worker.",
    version="0.1.0",
)

STATE: Dict[str, Any] = {
    "worker": {
        "status": "not_started",
        "slurm_job_id": None,
        "node": None,
        "worker_port": WORKER_PORT,
        "local_port": GATEWAY_LOCAL_PORT,
        "tunnel_pid": None,
        "last_error": None,
        "started_at": None,
    },
    "tunnel_proc": None,
}

TASKS: Dict[str, Dict[str, Any]] = {}
TASK_LOCK = asyncio.Lock()
WORKER_ENSURE_LOCK = asyncio.Lock()


# ---------------------------------------------------------------------------
# Auth
# ---------------------------------------------------------------------------

def require_token(authorization: Optional[str] = Header(None)):
    """Protect the public gateway if GATEWAY_API_TOKEN is configured."""
    if not GATEWAY_API_TOKEN:
        return

    expected = f"Bearer {GATEWAY_API_TOKEN}"
    if authorization != expected:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or missing bearer token.",
        )


# ---------------------------------------------------------------------------
# SSH helpers
# ---------------------------------------------------------------------------

def parse_remote_worker_info(output: str) -> dict:
    data = {}

    for line in output.strip().splitlines():
        line = line.strip()
        if "=" in line:
            key, value = line.split("=", 1)
            data[key.strip()] = value.strip()

    if "JOB_ID" not in data or "NODE" not in data or "PORT" not in data:
        raise RuntimeError(f"Could not parse worker info from remote output:\n{output}")

    return {
        "slurm_job_id": data["JOB_ID"],
        "node": data["NODE"],
        "worker_port": int(data["PORT"]),
        "local_port": GATEWAY_LOCAL_PORT,
    }


async def ensure_or_start_worker_job_remote() -> dict:
    """
    Use ONE ssh connection to:
    1. Check for a running worker job.
    2. If absent, submit sbatch.
    3. Wait on the cluster side for node/port.
    4. Return JOB_ID/NODE/PORT.
    """

    remote_script = f"""
set -u

MODEL_DIR={shlex.quote(CLUSTER_MODEL_DIR)}
INFO_DIR={shlex.quote(WORKER_INFO_DIR)}
JOB_NAME={shlex.quote(WORKER_JOB_NAME)}
USER_NAME={shlex.quote(CLUSTER_USER)}
SBATCH_SCRIPT={shlex.quote(SBATCH_SCRIPT)}
DEFAULT_PORT={WORKER_PORT}
TIMEOUT={WORKER_START_TIMEOUT_SECONDS}

existing="$(squeue -u "$USER_NAME" -n "$JOB_NAME" -t PENDING,RUNNING -h -o "%A|%T|%N|%R" | head -n 1 || true)"

if [ -n "$existing" ]; then
    JOB_ID="$(printf '%s\n' "$existing" | cut -d'|' -f1)"
    EXISTING_STATE="$(printf '%s\n' "$existing" | cut -d'|' -f2)"
    NODE="$(printf '%s\n' "$existing" | cut -d'|' -f3)"
    PORT="$DEFAULT_PORT"

    if [ "$EXISTING_STATE" = "RUNNING" ] && [ -n "$NODE" ] && [ "$NODE" != "(null)" ]; then
        if [ -f "$INFO_DIR/$JOB_ID/port.txt" ]; then
            PORT="$(cat "$INFO_DIR/$JOB_ID/port.txt" | tr -d '\r\n')"
        fi

        echo "JOB_ID=$JOB_ID"
        echo "NODE=$NODE"
        echo "PORT=$PORT"
        exit 0
    fi

    # Existing job is pending or has no usable node yet.
    # Do not submit a duplicate. Continue below and wait for this JOB_ID.
else
    sbatch_output="$(cd "$MODEL_DIR" && sbatch "$SBATCH_SCRIPT")"
    JOB_ID="$(printf '%s\n' "$sbatch_output" | sed -n 's/^Submitted batch job //p' | head -n 1)"

    if [ -z "$JOB_ID" ]; then
        echo "ERROR=Could not parse sbatch output: $sbatch_output"
        exit 2
    fi
fi

deadline=$((SECONDS + TIMEOUT))
NODE=""
PORT="$DEFAULT_PORT"

while [ "$SECONDS" -lt "$deadline" ]; do
    if [ -f "$INFO_DIR/$JOB_ID/node.txt" ]; then
        NODE="$(cat "$INFO_DIR/$JOB_ID/node.txt" | tr -d '\\r\\n')"
    fi

    if [ -f "$INFO_DIR/$JOB_ID/port.txt" ]; then
        PORT="$(cat "$INFO_DIR/$JOB_ID/port.txt" | tr -d '\\r\\n')"
    fi

    if [ -z "$NODE" ]; then
        sline="$(squeue -j "$JOB_ID" -h -o "%T|%N|%R" | head -n 1 || true)"
        state="$(printf '%s' "$sline" | cut -d'|' -f1)"
        snode="$(printf '%s' "$sline" | cut -d'|' -f2)"

        if [ "$state" = "RUNNING" ] && [ -n "$snode" ] && [ "$snode" != "(null)" ]; then
            NODE="$snode"
        fi

        if [ "$state" = "FAILED" ] || [ "$state" = "CANCELLED" ] || [ "$state" = "TIMEOUT" ] || [ "$state" = "OUT_OF_MEMORY" ]; then
            echo "ERROR=Slurm job $JOB_ID ended with state $state"
            exit 3
        fi
    fi

    if [ -n "$NODE" ]; then
        echo "JOB_ID=$JOB_ID"
        echo "NODE=$NODE"
        echo "PORT=$PORT"
        exit 0
    fi

    sleep 5
done

echo "ERROR=Timed out waiting for worker node for job $JOB_ID"
exit 4
"""

    remote_command = "bash -lc " + shlex.quote(remote_script)

    output = await asyncio.to_thread(
        ssh_run,
        remote_command,
        WORKER_START_TIMEOUT_SECONDS + 90,
    )

    if "ERROR=" in output and "JOB_ID=" not in output:
        raise RuntimeError(f"Remote worker startup failed:\n{output}")

    return parse_remote_worker_info(output)

def is_retryable_cluster_ssh_error(error_text: str) -> bool:
    text = (error_text or "").lower()
    return any(
        needle in text
        for needle in [
            "connection refused",
            "connection timed out",
            "connection reset",
            "connection closed",
            "temporary failure",
            "no route to host",
            "timed out after",
            "timeoutexpired",
        ]
    )

async def ensure_worker_running_with_retries() -> dict:
    deadline = time.time() + ENSURE_WORKER_RETRY_SECONDS
    attempt = 0
    last_error = None

    while time.time() < deadline:
        attempt += 1

        try:
            STATE["worker"].update(
                {
                    "status": "ensuring_worker",
                    "last_error": last_error,
                    "ensure_attempt": attempt,
                }
            )
            return await ensure_worker_running()

        except Exception as exc:
            last_error = str(exc)

            STATE["worker"].update(
                {
                    "status": "waiting_for_cluster_ssh",
                    "last_error": last_error,
                    "ensure_attempt": attempt,
                }
            )

            if not is_retryable_cluster_ssh_error(last_error):
                raise

            lower_error = (last_error or "").lower()

            if "connection refused" in lower_error:
                sleep_seconds = CLUSTER_SSH_REFUSED_COOLDOWN_SECONDS
            else:
                sleep_seconds = ENSURE_WORKER_RETRY_SLEEP_SECONDS

            STATE["worker"]["retry_sleep_seconds"] = sleep_seconds

            await asyncio.sleep(sleep_seconds)

    raise RuntimeError(
        f"Could not ensure worker after {ENSURE_WORKER_RETRY_SECONDS} seconds. "
        f"Last error: {last_error}"
    )

def _write_ssh_files() -> tuple[Optional[str], Optional[str]]:
    """Write SSH private key and known_hosts from env secrets to /tmp."""
    key_path = None
    known_hosts_path = None

    key_text = ""
    if SSH_PRIVATE_KEY_BASE64:
        try:
            key_text = base64.b64decode(SSH_PRIVATE_KEY_BASE64).decode("utf-8")
        except Exception as exc:
            raise RuntimeError(f"Could not decode SSH_PRIVATE_KEY_BASE64: {exc}")
    elif SSH_PRIVATE_KEY:
        key_text = SSH_PRIVATE_KEY.replace("\\n", "\n")

    if key_text.strip():
        key_path = "/tmp/geoglyph_gateway_ssh_key"
        Path(key_path).write_text(key_text.strip() + "\n", encoding="utf-8")
        os.chmod(key_path, 0o600)

    if SSH_KNOWN_HOSTS:
        known_hosts_path = "/tmp/geoglyph_gateway_known_hosts"
        known_hosts_text = SSH_KNOWN_HOSTS.replace("\\n", "\n").strip() + "\n"
        Path(known_hosts_path).write_text(known_hosts_text, encoding="utf-8")
        os.chmod(known_hosts_path, 0o644)

    return key_path, known_hosts_path


def _ssh_env() -> Optional[dict]:
    """Use sshpass via SSHPASS when password auth is enabled."""
    if not SSH_PASSWORD:
        return None
    env = os.environ.copy()
    env["SSHPASS"] = SSH_PASSWORD
    return env


def _ssh_base_cmd() -> list[str]:
    key_path, known_hosts_path = _write_ssh_files()
    use_password = bool(SSH_PASSWORD and not key_path)

    cmd = []
    if use_password:
        cmd.extend(["sshpass", "-e"])

    cmd.extend([
        "ssh",
        "-p",
        str(CLUSTER_SSH_PORT),
        "-o",
        "ServerAliveInterval=30",
        "-o",
        "ServerAliveCountMax=3",
    ])

    if use_password:
        cmd.extend([
            "-o",
            "BatchMode=no",
            "-o",
            "PreferredAuthentications=password,keyboard-interactive",
            "-o",
            "PubkeyAuthentication=no",
        ])
    else:
        cmd.extend(["-o", "BatchMode=yes"])

    if key_path:
        cmd.extend(["-i", key_path])

    if known_hosts_path:
        cmd.extend([
            "-o",
            "StrictHostKeyChecking=yes",
            "-o",
            f"UserKnownHostsFile={known_hosts_path}",
        ])
    else:
        # Prototype-friendly. For production, provide SSH_KNOWN_HOSTS.
        cmd.extend([
            "-o",
            "StrictHostKeyChecking=no",
            "-o",
            "UserKnownHostsFile=/tmp/geoglyph_gateway_known_hosts_auto",
        ])

    return cmd


def ssh_run(remote_command: str, timeout: int = 60) -> str:
    cmd = _ssh_base_cmd() + [f"{CLUSTER_USER}@{CLUSTER_LOGIN_HOST}", remote_command]
    proc = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout, env=_ssh_env())
    if proc.returncode != 0:
        raise RuntimeError(
            f"SSH command failed with code {proc.returncode}.\n"
            f"STDOUT: {proc.stdout}\nSTDERR: {proc.stderr}\nCOMMAND: {remote_command}"
        )
    return proc.stdout.strip()


def _open_tunnel_once(node: str, worker_port: int, local_port: int) -> int:
    """Open one SSH tunnel attempt. Raise with stderr if ssh exits immediately."""
    close_tunnel()

    cmd = _ssh_base_cmd() + [
        "-N",
        "-L",
        f"127.0.0.1:{local_port}:{node}:{worker_port}",
        f"{CLUSTER_USER}@{CLUSTER_LOGIN_HOST}",
    ]

    proc = subprocess.Popen(
        cmd,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        env=_ssh_env(),
        preexec_fn=os.setsid if hasattr(os, "setsid") else None,
    )

    STATE["tunnel_proc"] = proc
    STATE["worker"]["tunnel_pid"] = proc.pid
    STATE["worker"]["local_port"] = local_port

    time.sleep(1.0)

    if proc.poll() is not None:
        stdout = ""
        stderr = ""

        try:
            stdout = proc.stdout.read() if proc.stdout else ""
        except Exception:
            pass

        try:
            stderr = proc.stderr.read() if proc.stderr else ""
        except Exception:
            pass

        STATE["tunnel_proc"] = None
        STATE["worker"]["tunnel_pid"] = None

        raise RuntimeError(
            "SSH tunnel process exited immediately.\n"
            f"local_port={local_port}, target={node}:{worker_port}\n"
            f"exit_code={proc.returncode}\n"
            f"STDOUT: {stdout}\n"
            f"STDERR: {stderr}"
        )

    return proc.pid


def open_tunnel(node: str, worker_port: int = WORKER_PORT, local_port: int = GATEWAY_LOCAL_PORT) -> int:
    """Open local_port and forward it to node:worker_port through the login node."""
    last_error = None

    for attempt in range(1, SSH_TUNNEL_RETRIES + 1):
        try:
            STATE["worker"]["last_error"] = (
                f"Opening SSH tunnel attempt {attempt}/{SSH_TUNNEL_RETRIES} "
                f"via {CLUSTER_LOGIN_HOST}:{CLUSTER_SSH_PORT}, "
                f"local_port={local_port}, target={node}:{worker_port}"
            )
            return _open_tunnel_once(node, worker_port, local_port)

        except Exception as exc:
            last_error = str(exc)
            STATE["worker"]["last_error"] = last_error
            close_tunnel()

            # If the local port is busy, retrying same port is useless.
            # Let connect_to_worker decide whether to try another local port.
            if is_local_port_bind_error(last_error):
                raise

            # If kraken refused/timed out, retry same tunnel after backoff.
            if is_retryable_ssh_connection_error(last_error) and attempt < SSH_TUNNEL_RETRIES:
                time.sleep(SSH_TUNNEL_RETRY_SLEEP_SECONDS)
                continue

            raise

    raise RuntimeError(f"Could not open SSH tunnel. Last error: {last_error}")


def close_tunnel():
    proc = STATE.get("tunnel_proc")
    if proc is None:
        return

    if proc.poll() is None:
        try:
            if hasattr(os, "killpg"):
                os.killpg(os.getpgid(proc.pid), signal.SIGTERM)
            else:
                proc.terminate()
            proc.wait(timeout=5)
        except Exception:
            try:
                proc.kill()
            except Exception:
                pass

    STATE["tunnel_proc"] = None
    STATE["worker"]["tunnel_pid"] = None


def parse_sbatch_job_id(output: str) -> str:
    match = re.search(r"Submitted batch job\s+(\d+)", output)
    if not match:
        match = re.search(r"\b(\d+)\b", output)
    if not match:
        raise RuntimeError(f"Could not parse Slurm job id from sbatch output: {output!r}")
    return match.group(1)


async def worker_base_url() -> str:
    local_port = STATE["worker"].get("local_port") or GATEWAY_LOCAL_PORT
    return f"http://127.0.0.1:{local_port}"


async def wait_for_worker_health(timeout_seconds: int = 120) -> dict:
    deadline = time.time() + timeout_seconds
    url = f"{await worker_base_url()}/health"
    last_error = None

    async with httpx.AsyncClient(timeout=REQUEST_TIMEOUT) as client:
        while time.time() < deadline:
            try:
                response = await client.get(url)
                if response.status_code == 200:
                    return response.json()
                last_error = f"HTTP {response.status_code}: {response.text[:300]}"
            except Exception as exc:
                last_error = str(exc)

            await asyncio.sleep(2)

    raise RuntimeError(f"Worker did not become healthy in time. Last error: {last_error}")


async def get_worker_health_once(timeout_seconds: float = 5) -> Optional[dict]:
    """Return worker /health if the current tunnel works; otherwise return None."""
    tunnel_proc = STATE.get("tunnel_proc")
    if tunnel_proc is None or tunnel_proc.poll() is not None:
        return None

    url = f"{await worker_base_url()}/health"
    try:
        async with httpx.AsyncClient(timeout=timeout_seconds) as client:
            response = await client.get(url)
        if response.status_code == 200:
            return response.json()
    except Exception:
        return None
    return None

async def get_health_on_local_port(local_port: int, timeout_seconds: float = 5) -> Optional[dict]:
    """Return worker /health through a specific local port, independent of STATE."""
    url = f"http://127.0.0.1:{local_port}/health"

    try:
        async with httpx.AsyncClient(timeout=timeout_seconds) as client:
            response = await client.get(url)

        if response.status_code == 200:
            return response.json()
    except Exception:
        return None

    return None


def candidate_local_ports(preferred_port: int, allow_fallback: bool = True) -> list[int]:
    """Try preferred port first; optionally try a few nearby ports if it is stale/busy."""
    if not allow_fallback:
        return [preferred_port]

    ports = [preferred_port]
    for port in range(preferred_port + 1, preferred_port + 6):
        ports.append(port)

    return ports

def is_local_port_bind_error(error_text: str) -> bool:
    text = (error_text or "").lower()
    return any(
        needle in text
        for needle in [
            "address already in use",
            "bind",
            "cannot listen",
            "port is already allocated",
        ]
    )


def is_retryable_ssh_connection_error(error_text: str) -> bool:
    text = (error_text or "").lower()
    return any(
        needle in text
        for needle in [
            "connection refused",
            "connection timed out",
            "connection reset",
            "connection closed",
            "temporary failure",
            "no route to host",
        ]
    )


def parse_squeue_worker(output: str) -> Optional[dict]:
    """Parse one RUNNING Slurm worker line: job_id|node."""
    line = output.strip().splitlines()[0].strip() if output.strip() else ""
    if not line:
        return None

    parts = line.split("|", 1)
    if len(parts) != 2:
        raise RuntimeError(f"Unexpected squeue output for worker discovery: {output!r}")

    job_id, node = parts[0].strip(), parts[1].strip()
    if not job_id or not node or node in {"(null)", "None", "n/a"}:
        return None

    return {
        "slurm_job_id": job_id,
        "node": node,
        "worker_port": WORKER_PORT,
        "local_port": GATEWAY_LOCAL_PORT,
    }


async def discover_running_worker_job() -> Optional[dict]:
    """Find one RUNNING Slurm job named WORKER_JOB_NAME, if it exists."""
    job_name = shlex.quote(WORKER_JOB_NAME)
    user = shlex.quote(CLUSTER_USER)
    # We intentionally request RUNNING jobs only. Your assumption is: running or absent, never pending.
    cmd = f'squeue -u {user} -n {job_name} -t RUNNING -h -o "%A|%N" | head -n 1'
    output = await asyncio.to_thread(ssh_run, cmd, 30)
    return parse_squeue_worker(output)


async def connect_to_worker(
    worker: dict,
    health_timeout_seconds: int = 120,
    allow_port_fallback: bool = True,
) -> dict:
    """Open tunnel to the given worker node and verify /health."""
    node = worker["node"]
    worker_port = int(worker.get("worker_port") or WORKER_PORT)
    preferred_local_port = int(worker.get("local_port") or GATEWAY_LOCAL_PORT)
    job_id = worker.get("slurm_job_id")

    last_error = None

    for local_port in candidate_local_ports(preferred_local_port, allow_fallback=allow_port_fallback):
        STATE["worker"].update(
            {
                "status": "checking_existing_tunnel",
                "slurm_job_id": job_id,
                "node": node,
                "worker_port": worker_port,
                "local_port": local_port,
                "last_error": None,
            }
        )

        # Reuse an already healthy tunnel instead of killing it.
        existing_health = await get_health_on_local_port(local_port, timeout_seconds=5)
        if existing_health is not None:
            STATE["worker"].update(
                {
                    "status": "running",
                    "slurm_job_id": job_id,
                    "node": node,
                    "worker_port": worker_port,
                    "local_port": local_port,
                    "last_error": None,
                    "started_at": time.time(),
                }
            )
            return {**STATE["worker"], "worker_health": existing_health}

        try:
            STATE["worker"].update(
                {
                    "status": "opening_tunnel",
                    "slurm_job_id": job_id,
                    "node": node,
                    "worker_port": worker_port,
                    "local_port": local_port,
                    "last_error": None,
                }
            )

            pid = await asyncio.to_thread(open_tunnel, node, worker_port, local_port)

            STATE["worker"]["tunnel_pid"] = pid
            STATE["worker"].update(
                {
                    "status": "waiting_for_health",
                    "started_at": time.time(),
                }
            )

            health_payload = await wait_for_worker_health(timeout_seconds=health_timeout_seconds)

            STATE["worker"].update(
                {
                    "status": "running",
                    "last_error": None,
                    "local_port": local_port,
                    "worker_port": worker_port,
                    "node": node,
                    "slurm_job_id": job_id,
                }
            )

            return {**STATE["worker"], "worker_health": health_payload}

        except Exception as exc:
            last_error = str(exc)
            STATE["worker"].update(
                {
                    "status": "failed",
                    "last_error": last_error,
                    "local_port": local_port,
                }
            )

            close_tunnel()

            # Only try the next local port if the failure is actually about
            # binding/listening on the local port.
            if allow_port_fallback and is_local_port_bind_error(last_error):
                continue

            # SSH-to-kraken errors will not be fixed by changing local_port.
            break

    raise RuntimeError(
        f"Could not connect to worker {node}:{worker_port}. "
        f"Preferred local port was {preferred_local_port}. "
        f"Last error: {last_error}"
    )


async def start_worker_job_and_connect() -> dict:
    """Submit the Slurm worker job, get its node, open tunnel, and verify /health."""
    STATE["worker"].update(
        {
            "status": "submitting",
            "last_error": None,
            "node": None,
            "tunnel_pid": None,
        }
    )

    sbatch_cmd = f"cd {shlex.quote(CLUSTER_MODEL_DIR)} && sbatch {shlex.quote(SBATCH_SCRIPT)}"
    sbatch_output = await asyncio.to_thread(ssh_run, sbatch_cmd, 60)
    job_id = parse_sbatch_job_id(sbatch_output)

    STATE["worker"].update(
        {
            "status": "waiting_for_slurm",
            "slurm_job_id": job_id,
            "node": None,
            "worker_port": WORKER_PORT,
            "local_port": GATEWAY_LOCAL_PORT,
            "last_error": None,
        }
    )

    deadline = time.time() + WORKER_START_TIMEOUT_SECONDS
    node = None
    worker_port = WORKER_PORT
    last_error = None

    while time.time() < deadline:
        try:
            # First source of truth: Slurm.
            # Format: STATE|NODELIST|REASON
            squeue_cmd = f'squeue -j {shlex.quote(job_id)} -h -o "%T|%N|%R"'
            squeue_output = await asyncio.to_thread(ssh_run, squeue_cmd, 20)

            if not squeue_output.strip():
                # Job disappeared from squeue; ask sacct for final state.
                sacct_cmd = (
                    f"sacct -j {shlex.quote(job_id)} "
                    "--format=JobID,JobName,State,ExitCode,Elapsed,NodeList,Reason "
                    "--noheader"
                )
                try:
                    sacct_output = await asyncio.to_thread(ssh_run, sacct_cmd, 20)
                except Exception as sacct_exc:
                    sacct_output = f"sacct failed: {sacct_exc}"

                raise RuntimeError(
                    f"Slurm job {job_id} is no longer visible in squeue. "
                    f"sacct output: {sacct_output}"
                )

            first_line = squeue_output.strip().splitlines()[0]
            parts = first_line.split("|")
            slurm_state = parts[0].strip() if len(parts) > 0 else ""
            slurm_node = parts[1].strip() if len(parts) > 1 else ""
            slurm_reason = parts[2].strip() if len(parts) > 2 else ""

            STATE["worker"].update(
                {
                    "status": "waiting_for_slurm",
                    "slurm_job_id": job_id,
                    "node": slurm_node if slurm_node not in {"", "(null)", "None", "n/a"} else None,
                    "last_error": f"Slurm state={slurm_state}, reason={slurm_reason}",
                }
            )

            if slurm_state in {"FAILED", "CANCELLED", "TIMEOUT", "OUT_OF_MEMORY", "NODE_FAIL"}:
                raise RuntimeError(
                    f"Slurm worker job {job_id} ended with state={slurm_state}, reason={slurm_reason}"
                )

            if slurm_state != "RUNNING":
                await asyncio.sleep(3)
                continue

            # At this point Slurm says the job is running.
            if slurm_node and slurm_node not in {"(null)", "None", "n/a"}:
                node = slurm_node

            # Optional: prefer node.txt if it exists, but do not depend on it.
            try:
                node_cmd = (
                    f"test -f {shlex.quote(WORKER_INFO_DIR)}/{shlex.quote(job_id)}/node.txt "
                    f"&& cat {shlex.quote(WORKER_INFO_DIR)}/{shlex.quote(job_id)}/node.txt "
                    f"|| true"
                )
                node_file = await asyncio.to_thread(ssh_run, node_cmd, 20)
                if node_file.strip():
                    node = node_file.strip()
            except Exception as exc:
                last_error = f"Could not read node.txt, using squeue node if available: {exc}"

            # Optional: read port.txt, but default to WORKER_PORT.
            try:
                port_cmd = (
                    f"test -f {shlex.quote(WORKER_INFO_DIR)}/{shlex.quote(job_id)}/port.txt "
                    f"&& cat {shlex.quote(WORKER_INFO_DIR)}/{shlex.quote(job_id)}/port.txt "
                    f"|| true"
                )
                port_text = await asyncio.to_thread(ssh_run, port_cmd, 20)
                if port_text.strip():
                    worker_port = int(port_text.strip())
            except Exception as exc:
                last_error = f"Could not read port.txt, using WORKER_PORT={WORKER_PORT}: {exc}"
                worker_port = WORKER_PORT

            if node:
                break

            last_error = (
                f"Job {job_id} is RUNNING but no node was available. "
                f"squeue output: {squeue_output!r}"
            )
            STATE["worker"]["last_error"] = last_error
            await asyncio.sleep(3)

        except Exception as exc:
            last_error = str(exc)
            STATE["worker"].update(
                {
                    "status": "waiting_for_slurm",
                    "slurm_job_id": job_id,
                    "last_error": last_error,
                }
            )
            await asyncio.sleep(3)

    if not node:
        raise RuntimeError(
            f"Worker node was not discovered for job {job_id}. "
            f"Last error: {last_error}"
        )

    STATE["worker"].update(
        {
            "status": "connecting",
            "slurm_job_id": job_id,
            "node": node,
            "worker_port": worker_port,
            "local_port": GATEWAY_LOCAL_PORT,
            "last_error": None,
        }
    )

    return await connect_to_worker(
        {
            "slurm_job_id": job_id,
            "node": node,
            "worker_port": worker_port,
            "local_port": GATEWAY_LOCAL_PORT,
        },
        health_timeout_seconds=120,
    )


async def ensure_worker_running() -> dict:
    """
    Ensure that /process can forward to a healthy worker.

    Cold start uses one remote SSH command to avoid triggering cluster SSH bans.
    """
    async with WORKER_ENSURE_LOCK:
        health_payload = await get_worker_health_once(timeout_seconds=5)
        if health_payload is not None:
            STATE["worker"].update({"status": "running", "last_error": None})
            return {**STATE["worker"], "worker_health": health_payload}

        try:
            STATE["worker"].update(
                {
                    "status": "ensuring_slurm_worker",
                    "last_error": None,
                }
            )

            worker = await ensure_or_start_worker_job_remote()

            STATE["worker"].update(
                {
                    "status": "slurm_worker_ready",
                    "slurm_job_id": worker["slurm_job_id"],
                    "node": worker["node"],
                    "worker_port": worker["worker_port"],
                    "local_port": worker["local_port"],
                    "last_error": None,
                }
            )

            return await connect_to_worker(worker, health_timeout_seconds=120)

        except Exception as exc:
            STATE["worker"].update({"status": "failed", "last_error": str(exc)})
            raise

@app.get("/cluster/ping", dependencies=[Depends(require_token)])
async def cluster_ping():
    """
    Minimal diagnostic endpoint.

    It only checks whether the Hugging Face Gateway can SSH into the cluster login node.
    It does not use Slurm, does not open a tunnel, and does not contact the SAM2 worker.
    """
    started = time.time()

    remote_command = (
        "echo GEOGLYPH_SSH_OK && "
        "hostname && "
        "whoami && "
        "date"
    )

    try:
        output = await asyncio.to_thread(ssh_run, remote_command, 20)

        return {
            "status": "ok",
            "cluster_login_host": CLUSTER_LOGIN_HOST,
            "cluster_user": CLUSTER_USER,
            "cluster_ssh_port": CLUSTER_SSH_PORT,
            "elapsed_seconds": round(time.time() - started, 3),
            "output": output,
        }

    except Exception as exc:
        return {
            "status": "failed",
            "cluster_login_host": CLUSTER_LOGIN_HOST,
            "cluster_user": CLUSTER_USER,
            "cluster_ssh_port": CLUSTER_SSH_PORT,
            "elapsed_seconds": round(time.time() - started, 3),
            "error": str(exc),
        }

# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------

class WorkerStartResponse(BaseModel):
    status: str
    slurm_job_id: Optional[str]
    node: Optional[str]
    worker_port: int
    local_port: int
    tunnel_pid: Optional[int]
    worker_health: Optional[dict] = None


class WorkerConnectRequest(BaseModel):
    node: str
    worker_port: int = WORKER_PORT
    local_port: int = GATEWAY_LOCAL_PORT


# ---------------------------------------------------------------------------
# Gateway endpoints
# ---------------------------------------------------------------------------

@app.on_event("shutdown")
def shutdown_event():
    close_tunnel()


@app.get("/health")
async def health():
    tunnel_proc = STATE.get("tunnel_proc")
    tunnel_alive = bool(tunnel_proc and tunnel_proc.poll() is None)
    return {
        "status": "ok",
        "gateway": "running",
        "worker_job_name": WORKER_JOB_NAME,
        "worker": STATE["worker"],
        "tunnel_alive": tunnel_alive,
        "tasks": len(TASKS),
    }


@app.post("/worker/connect", response_model=WorkerStartResponse, dependencies=[Depends(require_token)])
async def worker_connect(req: WorkerConnectRequest):
    """Connect the gateway to an already-running Slurm worker node."""
    try:
        async with WORKER_ENSURE_LOCK:
            worker_payload = await connect_to_worker(
                {
                    "slurm_job_id": None,
                    "node": req.node,
                    "worker_port": req.worker_port,
                    "local_port": req.local_port,
                },
                health_timeout_seconds=60,
                allow_port_fallback=False,
            )

        return WorkerStartResponse(
            **STATE["worker"],
            worker_health=worker_payload.get("worker_health"),
        )

    except Exception as exc:
        STATE["worker"].update({"status": "failed", "last_error": str(exc)})
        raise HTTPException(status_code=500, detail=str(exc))
    

@app.post("/worker/start", response_model=WorkerStartResponse, dependencies=[Depends(require_token)])
async def worker_start():
    """Explicit admin endpoint: reuse a healthy worker, connect to an existing job, or start one."""
    try:
        worker_payload = await ensure_worker_running_with_retries()
        return WorkerStartResponse(**STATE["worker"], worker_health=worker_payload.get("worker_health"))
    except Exception as exc:
        STATE["worker"].update({"status": "failed", "last_error": str(exc)})
        raise HTTPException(status_code=500, detail=str(exc))


@app.post("/worker/stop", dependencies=[Depends(require_token)])
async def worker_stop():
    """Close the tunnel. Optionally cancel the Slurm job if there is a known job id."""
    job_id = STATE["worker"].get("slurm_job_id")
    close_tunnel()

    cancel_error = None
    if job_id:
        try:
            await asyncio.to_thread(ssh_run, f"scancel {job_id}", 30)
        except Exception as exc:
            cancel_error = str(exc)

    STATE["worker"].update(
        {
            "status": "stopped",
            "tunnel_pid": None,
            "last_error": cancel_error,
        }
    )
    return {"status": "stopped", "slurm_job_id": job_id, "cancel_error": cancel_error}


@app.get("/worker/health", dependencies=[Depends(require_token)])
async def worker_health():
    try:
        health_payload = await wait_for_worker_health(timeout_seconds=10)
        return {"status": "ok", "worker_health": health_payload, "worker": STATE["worker"]}
    except Exception as exc:
        raise HTTPException(status_code=503, detail=str(exc))
    
async def ensure_worker_running_for_process_task(gateway_task_id: str) -> dict:
    """
    Robust worker ensure for /process.

    Temporary SSH failures to kraken should not immediately fail the task.
    Instead, the task stays alive and retries.
    """
    deadline = time.time() + ENSURE_WORKER_RETRY_SECONDS
    attempt = 0
    last_error = None

    while time.time() < deadline:
        attempt += 1

        try:
            async with TASK_LOCK:
                TASKS[gateway_task_id].update(
                    {
                        "status": "ensuring_worker",
                        "ensure_attempt": attempt,
                        "worker_gateway_error": last_error,
                    }
                )

            return await ensure_worker_running()

        except Exception as exc:
            last_error = str(exc)

            async with TASK_LOCK:
                TASKS[gateway_task_id].update(
                    {
                        "status": "waiting_for_cluster_ssh",
                        "ensure_attempt": attempt,
                        "worker_gateway_error": last_error,
                        "error": None,
                    }
                )

            STATE["worker"].update(
                {
                    "status": "waiting_for_cluster_ssh",
                    "last_error": last_error,
                }
            )

            # If this is not a temporary cluster SSH problem, fail normally.
            if not is_retryable_cluster_ssh_error(last_error):
                raise

            lower_error = (last_error or "").lower()

            if "connection refused" in lower_error:
                sleep_seconds = CLUSTER_SSH_REFUSED_COOLDOWN_SECONDS
            else:
                sleep_seconds = ENSURE_WORKER_RETRY_SLEEP_SECONDS

            async with TASK_LOCK:
                TASKS[gateway_task_id].update(
                    {
                        "status": "waiting_for_cluster_ssh",
                        "retry_sleep_seconds": sleep_seconds,
                        "worker_gateway_error": last_error,
                        "error": None,
                    }
                )

            await asyncio.sleep(sleep_seconds)

    raise RuntimeError(
        f"Could not connect to worker after {ENSURE_WORKER_RETRY_SECONDS} seconds. "
        f"Last error: {last_error}"
    )


async def forward_task_to_worker(gateway_task_id: str, crop_path: Path):
    """Forward crop to the worker API, poll until completion, then download GPKG."""
    output_path = RESULTS_DIR / f"{gateway_task_id}.gpkg"
    worker_task_id = None

    try:
        async with TASK_LOCK:
            TASKS[gateway_task_id]["status"] = "forwarding"

        async with TASK_LOCK:
            TASKS[gateway_task_id]["status"] = "ensuring_worker"

        worker_payload = await ensure_worker_running_for_process_task(gateway_task_id)

        async with TASK_LOCK:
            TASKS[gateway_task_id]["worker_gateway"] = worker_payload
            TASKS[gateway_task_id]["status"] = "forwarding"

        base = await worker_base_url()

        async with httpx.AsyncClient(timeout=None) as client:
            with crop_path.open("rb") as f:
                files = {"crop": (crop_path.name, f, "image/tiff")}
                data = {"device": "cuda"}
                response = await client.post(f"{base}/process", files=files, data=data, timeout=REQUEST_TIMEOUT)

            response.raise_for_status()
            process_payload = response.json()
            worker_task_id = process_payload["task_id"]

            async with TASK_LOCK:
                TASKS[gateway_task_id].update(
                    {
                        "status": "processing",
                        "worker_task_id": worker_task_id,
                        "worker_process_response": process_payload,
                    }
                )

            deadline = time.time() + SEGMENT_TIMEOUT_SECONDS
            last_status = None

            while time.time() < deadline:
                status_response = await client.get(f"{base}/status/{worker_task_id}", timeout=REQUEST_TIMEOUT)
                status_response.raise_for_status()
                status_payload = status_response.json()
                last_status = status_payload

                async with TASK_LOCK:
                    TASKS[gateway_task_id]["worker_status"] = status_payload

                if status_payload.get("status") == "completed":
                    if not status_payload.get("output_exists"):
                        async with TASK_LOCK:
                            TASKS[gateway_task_id].update(
                                {
                                    "status": "completed_no_output",
                                    "finished_at": time.time(),
                                    "n_masks": status_payload.get("n_masks"),
                                    "output_exists": False,
                                    "result": status_payload.get("result"),
                                }
                            )
                        return

                    download_response = await client.get(f"{base}/download/{worker_task_id}", timeout=None)
                    download_response.raise_for_status()
                    output_path.write_bytes(download_response.content)

                    async with TASK_LOCK:
                        TASKS[gateway_task_id].update(
                            {
                                "status": "completed",
                                "finished_at": time.time(),
                                "n_masks": status_payload.get("n_masks"),
                                "output_exists": True,
                                "download_url": f"/download/{gateway_task_id}",
                                "result_path": str(output_path),
                                "result": status_payload.get("result"),
                            }
                        )
                    return

                if status_payload.get("status") == "failed":
                    raise RuntimeError(f"Worker task failed: {status_payload.get('error')}")

                await asyncio.sleep(STATUS_POLL_SECONDS)

            raise TimeoutError(f"Worker task timeout. Last status: {last_status}")

    except Exception as exc:
        async with TASK_LOCK:
            TASKS[gateway_task_id].update(
                {
                    "status": "failed",
                    "finished_at": time.time(),
                    "error": str(exc),
                    "worker_task_id": worker_task_id,
                }
            )
    finally:
        try:
            crop_path.unlink(missing_ok=True)
        except Exception:
            pass


@app.post("/process", status_code=status.HTTP_202_ACCEPTED, dependencies=[Depends(require_token)])
async def process(crop: UploadFile = File(...)):
    """Submit a crop. The gateway auto-discovers/starts/connects the Slurm worker in the background."""

    gateway_task_id = uuid.uuid4().hex[:12]
    crop_path = RESULTS_DIR / f"{gateway_task_id}_crop.tif"

    try:
        with crop_path.open("wb") as f:
            shutil.copyfileobj(crop.file, f)
    finally:
        await crop.close()

    if not crop_path.is_file() or crop_path.stat().st_size == 0:
        raise HTTPException(status_code=400, detail="Uploaded crop is empty.")

    async with TASK_LOCK:
        TASKS[gateway_task_id] = {
            "task_id": gateway_task_id,
            "status": "pending",
            "created_at": time.time(),
            "original_filename": crop.filename,
            "worker_task_id": None,
            "worker_status": None,
            "error": None,
            "output_exists": None,
            "download_url": None,
        }

    asyncio.create_task(forward_task_to_worker(gateway_task_id, crop_path))

    return {"task_id": gateway_task_id, "status": "pending", "queue_size": len(TASKS)}


@app.get("/status/{task_id}", dependencies=[Depends(require_token)])
async def get_status(task_id: str):
    if task_id not in TASKS:
        raise HTTPException(status_code=404, detail="Task not found")
    return TASKS[task_id]


@app.get("/download/{task_id}", dependencies=[Depends(require_token)])
async def download(task_id: str):
    if task_id not in TASKS:
        raise HTTPException(status_code=404, detail="Task not found")

    task = TASKS[task_id]
    if task.get("status") != "completed":
        raise HTTPException(status_code=400, detail=f"Task is not completed. Current status: {task.get('status')}")

    result_path = Path(task["result_path"])
    if not result_path.is_file():
        raise HTTPException(status_code=404, detail="Result file is missing.")

    return FileResponse(
        path=str(result_path),
        media_type="application/geopackage+sqlite3",
        filename=f"sam2_result_{task_id}.gpkg",
    )