"""
HyperTower distributed job client daemon.

Registers with the server, polls for jobs, syncs code, runs training,
uploads results, and loops. Parses stdout to stream live status.

Usage:
    python -m v4.distributed.client \\
        --server http://apollo:8765 \\
        --token <secret> \\
        --server-ssh rpotter@apollo \\
        --server-path /home/rpotter/hypertower \\
        [--local-path ~/hypertower] \\
        [--poll-interval 15]

Compatibility test (verify GPU env, 1-fold dry-run):
    python -m v4.distributed.client ... --test --config v4/configs/ensemble_fused.json
"""
from __future__ import annotations

import argparse
import os
import re
import shutil
import signal
import socket
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Optional

import requests

from .protocol import (
    JobResult,
    JobSpec,
    PollResponse,
    RegisterRequest,
    RegisterResponse,
    StatusPush,
)

# ──────────────────────────────────────────────────────────────
# Server HTTP wrapper
# ──────────────────────────────────────────────────────────────

class _Server:
    def __init__(self, base_url: str, token: str):
        self.base_url = base_url.rstrip("/")
        self._h = {"x-token": token}
        self.client_id: str = ""

    def _post(self, path: str, **kw) -> dict:
        r = requests.post(f"{self.base_url}{path}", headers=self._h, timeout=15, **kw)
        r.raise_for_status()
        return r.json()

    def register(self, hostname: str, gpu_info: str) -> str:
        data = self._post("/register",
                          json={"hostname": hostname, "gpu_info": gpu_info})
        self.client_id = data["client_id"]
        self.hostname = hostname
        self.gpu_info = gpu_info
        return self.client_id

    def _reregister(self):
        try:
            self._post("/register",
                       json={"hostname": self.hostname, "gpu_info": self.gpu_info},
                       params={"reuse_id": self.client_id})
        except Exception:
            pass

    def poll(self) -> Optional[JobSpec]:
        data = self._post("/poll", params={"client_id": self.client_id})
        if data.get("please_reregister"):
            self._reregister()
        return JobSpec(**data["job"]) if data.get("job") else None

    def push_status(self, status: StatusPush):
        try:
            r = requests.post(
                f"{self.base_url}/status/{self.client_id}",
                json=status.model_dump(),
                headers=self._h,
                timeout=5,
            )
            if r.ok and r.json().get("please_reregister"):
                self._reregister()
        except Exception:
            pass  # don't crash job on status push failure

    def complete(self, job_id: str, success: bool, error_msg: Optional[str] = None):
        self._post("/complete",
                   json={"job_id": job_id, "success": success, "error_msg": error_msg})


# ──────────────────────────────────────────────────────────────
# GPU info
# ──────────────────────────────────────────────────────────────

def _gpu_precheck() -> tuple[bool, str]:
    """Quick allocation + op test on every visible CUDA/HIP device.

    Catches the case where torch.cuda.is_available() returns True but the
    device is busy (e.g. another process holding it — Ollama, a forgotten
    notebook, an OS-level driver issue). Returns (ok, message); when ok is
    False, the caller should defer polling rather than accept a job that will
    crash on the first .to(device) call.
    """
    try:
        import torch
    except Exception as e:
        return False, f"torch import failed: {e}"
    if not torch.cuda.is_available():
        return True, "no-cuda (cpu-only client)"
    try:
        for i in range(torch.cuda.device_count()):
            x = torch.zeros(1024, device=f"cuda:{i}")
            _ = (x + 1).sum().item()  # forces actual kernel launch
            del x
        torch.cuda.empty_cache()
        return True, "ok"
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"


def _gpu_info() -> str:
    # NVIDIA
    try:
        out = subprocess.check_output(
            ["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader"],
            text=True, stderr=subprocess.DEVNULL,
        ).strip()
        if out:
            return " | ".join(out.splitlines())
    except Exception:
        pass
    # AMD
    try:
        out = subprocess.check_output(
            ["rocm-smi", "--showproductname", "--csv"],
            text=True, stderr=subprocess.DEVNULL,
        ).strip().splitlines()
        names = [l for l in out if l and not l.startswith("device")]
        if names:
            return "AMD: " + " | ".join(names)
    except Exception:
        pass
    # AMD fallback
    try:
        out = subprocess.check_output(
            ["rocminfo"],
            text=True, stderr=subprocess.DEVNULL,
        )
        names = [l.split(":", 1)[1].strip() for l in out.splitlines()
                 if "Marketing Name:" in l]
        if names:
            return "AMD: " + " | ".join(names)
    except Exception:
        pass
    return "no-gpu"


# ──────────────────────────────────────────────────────────────
# rsync helpers
# ──────────────────────────────────────────────────────────────

def _rsync(src: str, dst: str, delete: bool = False,
           excludes: list[str] | None = None):
    cmd = ["rsync", "-az", "--info=progress2"]
    if delete:
        cmd.append("--delete")
    for pat in (excludes or []):
        cmd.append(f"--exclude={pat}")
    cmd += [src, dst]
    subprocess.run(cmd, check=True)


def _sync_code(server_ssh: str, server_path: str, local_path: str):
    """Pull v4/ source from server → local (overwrites local changes).

    results/ is excluded so the client never overwrites or deletes its own
    per-job output directory, and so it never pulls down the full corpus of
    historical results from the server.
    """
    src = f"{server_ssh}:{server_path}/v4/"
    dst = f"{local_path}/v4/"
    Path(dst).mkdir(parents=True, exist_ok=True)
    _rsync(src, dst, delete=True, excludes=["results/"])


def _upload_results(server_ssh: str, server_path: str, local_path: str,
                    run_name: str, output_dir: str):
    src = f"{local_path}/{output_dir}/{run_name}/"
    dst = f"{server_ssh}:{server_path}/{output_dir}/{run_name}/"
    remote_parent = f"{server_path}/{output_dir}/{Path(run_name).parent}"
    subprocess.run(["ssh", server_ssh, f"mkdir -p '{remote_parent}'"], check=True)
    _rsync(src, dst)


def _clean_local(local_path: str, run_name: str, output_dir: str):
    target = Path(local_path) / output_dir / run_name
    if target.exists():
        shutil.rmtree(target)
        print(f"[client] cleaned {target}", flush=True)


# ──────────────────────────────────────────────────────────────
# Stdout parsers (match v4_hypertower.py / fusion.py print format)
# ──────────────────────────────────────────────────────────────

# "  fold2 [nt] ep004/036 [fused_warmup  ]  loss=0.4321  acc=0.876  val_auc=0.7654"
_EP_RE = re.compile(
    r"fold(\d+)\s+\[([^\]]+)\]\s+ep(\d+)/(\d+).*?val_auc=([0-9.nan]+)",
    re.IGNORECASE,
)
# "── fold 2/5  train_groups=..."
_FOLD_RE = re.compile(r"fold\s+(\d+)/\d+", re.IGNORECASE)


def _parse_line(line: str) -> dict:
    """Return any structured fields found in a stdout line."""
    out = {}
    m = _FOLD_RE.search(line)
    if m:
        out["fold"] = int(m.group(1))
    m = _EP_RE.search(line)
    if m:
        out["fold"]         = int(m.group(1))
        out["stage"]        = m.group(2)
        out["epoch"]        = int(m.group(3))
        out["total_epochs"] = int(m.group(4))
        try:
            out["last_val_auc"] = float(m.group(5))
        except ValueError:
            pass
    return out


# ──────────────────────────────────────────────────────────────
# Core job runner
# ──────────────────────────────────────────────────────────────

def _run_job(job: JobSpec, server: _Server,
             server_ssh: str, server_path: str, local_path: str,
             no_sync: bool = False,
             extra_args: list[str] | None = None) -> bool:
    extra_args = extra_args or []
    # 1. Sync code
    if no_sync:
        print(f"[client] skipping sync (--no-sync)", flush=True)
    else:
        print(f"[client] syncing v4/ from server...", flush=True)
        server.push_status(StatusPush(state="syncing", job_id=job.job_id, run_name=job.run_name))
        _sync_code(server_ssh, server_path, local_path)

    # 2. Launch training subprocess
    cmd = [sys.executable, "-m", job.module] + job.args + extra_args
    print(f"[client] running: {' '.join(cmd)}", flush=True)
    server.push_status(StatusPush(state="running", job_id=job.job_id, run_name=job.run_name))

    log_dir = Path(local_path) / "v4" / "distributed" / "logs"
    log_dir.mkdir(parents=True, exist_ok=True)
    log_file = log_dir / f"job_{job.job_id}.log"
    ctx: dict = {}
    last_push = [time.time()]  # mutable holder so _tail and _heartbeat share it

    def _push():
        server.push_status(StatusPush(
            state="running",
            job_id=job.job_id,
            run_name=job.run_name,
            fold=ctx.get("fold"),
            stage=ctx.get("stage"),
            epoch=ctx.get("epoch"),
            total_epochs=ctx.get("total_epochs"),
            last_val_auc=ctx.get("last_val_auc"),
        ))
        last_push[0] = time.time()

    def _tail(path: Path):
        with open(path, "r") as f:
            while True:
                raw = f.readline()
                if raw:
                    print(raw, end="", flush=True)
                    info = _parse_line(raw)
                    ctx.update(info)
                    if "epoch" in info:
                        _push()
                elif proc.poll() is not None:
                    for raw in f:
                        print(raw, end="", flush=True)
                    break
                else:
                    time.sleep(0.05)

    def _heartbeat():
        # Push a status update every ~30s even when no log line is parsed.
        # Prevents the server's reaper from declaring this client stale during
        # long deterministic blocks (UNet fine-tune, data load, etc.).
        while proc.poll() is None:
            time.sleep(5)
            if time.time() - last_push[0] >= 30:
                try:
                    _push()
                except Exception:
                    pass

    with open(log_file, "w") as logf:
        proc = subprocess.Popen(
            cmd,
            stdout=logf,
            stderr=logf,
            cwd=local_path,
            start_new_session=True,
        )

    tailer    = threading.Thread(target=_tail,      args=(log_file,), daemon=True)
    heartbeat = threading.Thread(target=_heartbeat,                     daemon=True)
    tailer.start()
    heartbeat.start()
    proc.wait()
    tailer.join(timeout=5)
    heartbeat.join(timeout=5)

    success = proc.returncode == 0
    if success:
        log_file.unlink(missing_ok=True)
    else:
        failed_path = log_file.with_suffix(".failed.log")
        try:
            log_file.replace(failed_path)
            print(f"[client] preserved failure log at {failed_path}", flush=True)
        except Exception:
            pass

    if not success:
        try:
            os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
        except ProcessLookupError:
            pass

    if not success:
        print(f"[client] job FAILED (rc={proc.returncode})", flush=True)
        return False

    # 3. Upload results
    if no_sync:
        print(f"[client] skipping upload (--no-sync, results already local)", flush=True)
    else:
        print(f"[client] uploading results...", flush=True)
        server.push_status(StatusPush(state="uploading", job_id=job.job_id, run_name=job.run_name))
        _upload_results(server_ssh, server_path, local_path, job.run_name, job.output_dir)
        _clean_local(local_path, job.run_name, job.output_dir)

    print(f"[client] job {job.job_id} complete.", flush=True)
    return True


# ──────────────────────────────────────────────────────────────
# Compatibility test
# ──────────────────────────────────────────────────────────────

def _compat_test(server: _Server, server_ssh: str, server_path: str,
                 local_path: str, config: str, extra_args: list[str],
                 no_sync: bool = False):
    """Run 1 fold to verify the env works end-to-end."""
    print("[client] === compatibility test ===", flush=True)
    job = JobSpec(
        job_id="compat-test",
        run_name="_compat_test",
        module="v4.classes.v4_hypertower",
        args=["--config", config, "--device", "cpu"] + extra_args,
        output_dir="v4/results",
    )
    ok = _run_job(job, server, server_ssh, server_path, local_path,
                  no_sync=no_sync)
    _clean_local(local_path, "_compat_test", "v4/results")
    if ok:
        print("[client] compatibility test PASSED ✓", flush=True)
    else:
        print("[client] compatibility test FAILED ✗", flush=True)
    return ok


# ──────────────────────────────────────────────────────────────
# Main daemon
# ──────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    ap.add_argument("--server", required=True,
                    help="Server URL, e.g. http://apollo:8765")
    ap.add_argument("--token", default=os.environ.get("HT_TOKEN", ""),
                    help="Shared secret (or set HT_TOKEN env var)")
    ap.add_argument("--server-ssh", required=True,
                    help="SSH target for rsync, e.g. rpotter@apollo")
    ap.add_argument("--server-path", required=True,
                    help="Absolute path to hypertower root on server")
    ap.add_argument("--local-path",
                    default=str(Path.home() / "hypertower"),
                    help="Absolute path to hypertower root on this machine")
    ap.add_argument("--poll-interval", type=int, default=15,
                    help="Seconds to wait between polls when idle")
    ap.add_argument("--extra-args", nargs=argparse.REMAINDER, default=[],
                    help="Extra args appended to every job on this client. "
                         "Use -- to separate: --extra-args -- --device cpu")
    ap.add_argument("--no-sync", action="store_true",
                    help="Skip rsync of v4/ before each job (use when client IS the server)")
    ap.add_argument("--test", action="store_true",
                    help="Run 1-fold compatibility test and exit")
    ap.add_argument("--config", default="v4/configs/ensemble_fused.json",
                    help="Config path for --test mode")
    args = ap.parse_args()

    if not args.token:
        ap.error("--token is required (or set HT_TOKEN)")

    hostname = socket.gethostname()
    gpu_info = _gpu_info()
    server = _Server(args.server, args.token)

    client_id = server.register(hostname, gpu_info)
    print(f"[client] registered as {client_id} ({hostname} | {gpu_info})", flush=True)

    if args.test:
        extra = [a for a in args.extra_args if a != "--"]
        sys.exit(0 if _compat_test(
            server, args.server_ssh, args.server_path,
            args.local_path, args.config, extra,
            no_sync=args.no_sync,
        ) else 1)

    print(f"[client] polling every {args.poll_interval}s...", flush=True)
    last_precheck_msg = ""
    while True:
        try:
            ok, msg = _gpu_precheck()
            if not ok:
                if msg != last_precheck_msg:
                    print(f"[client] gpu precheck failed ({msg}) — deferring polls",
                          flush=True)
                    last_precheck_msg = msg
                server.push_status(StatusPush(state="gpu_busy"))
                time.sleep(args.poll_interval)
                continue
            if last_precheck_msg:
                print(f"[client] gpu precheck recovered — resuming polls", flush=True)
                last_precheck_msg = ""

            job = server.poll()
            if job is None:
                server.push_status(StatusPush(state="idle"))
                time.sleep(args.poll_interval)
                continue

            success = _run_job(
                job, server,
                args.server_ssh, args.server_path, args.local_path,
                no_sync=args.no_sync,
                extra_args=[a for a in args.extra_args if a != "--"],
            )
            server.complete(job.job_id, success,
                            error_msg=None if success else "non-zero exit code")
            server.push_status(StatusPush(state="idle"))

        except KeyboardInterrupt:
            print("\n[client] shutting down", flush=True)
            break
        except Exception as exc:
            print(f"[client] error: {exc}", flush=True)
            time.sleep(args.poll_interval)


if __name__ == "__main__":
    main()
