#!/usr/bin/env python3
"""Cluster orchestrator for the Russell 3000 extension (runs on the M1M32 gateway).

    python3 orchestrate.py sync | downloads | ship | pipeline | pull | merge | status | all

State: every stage is idempotent and polls marker files on the nodes, so the
script can be stopped and restarted at any stage.  Logs: ~/edgar-orchestrate/.
"""

from __future__ import annotations

import json
import os
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime
from pathlib import Path

HOME = Path.home()
CODE = HOME / "UQO_RESEARCH_DB-code"          # code + metadata mirror on the gateway
WORK = HOME / "edgar-orchestrate"
WORK.mkdir(exist_ok=True)
LOG = WORK / "orchestrate.log"

MASTER = "M3U96b"
BHS64 = "ubuntu@51.161.112.61"
BHS64_DL = "edgar-dl"                       # ~/edgar-dl/<node>/{raw,submissions,download_log.jsonl,DONE}
REMOTE_ROOT = "UQO_RESEARCH_DB"

# node: (lan_ip, parse_workers)
NODES = {
    "M3U96b": ("192.168.2.85", 26),
    "M3U96a": ("192.168.2.80", 24),
    "M2U64":  ("192.168.2.61", 20),
    "M5M36":  ("192.168.2.104", 14),
    "M4M64a": ("192.168.2.70", 12),
    "M4M64c": ("192.168.2.103", 14),
    "M4BP48": ("192.168.2.114", 12),
    "M4BP36": ("192.168.2.113", 10),
    "m4mm":   ("192.168.2.92", 10),
    "m4mn":   ("192.168.2.93", 10),
}
LAN_RATE_TOTAL = 9.0       # SEC: 10 req/s per IP; all LAN nodes share one public IP
BHS64_RATE_TOTAL = 9.0

SSH_OPTS = ["-o", "BatchMode=yes", "-o", "ConnectTimeout=20", "-o", "StrictHostKeyChecking=accept-new",
            "-o", "ServerAliveInterval=30"]


def log(msg: str) -> None:
    line = f"{datetime.now():%F %T} {msg}"
    print(line, flush=True)
    with open(LOG, "a") as f:
        f.write(line + "\n")


def ssh(host: str, cmd: str, agent: bool = False, timeout: int = 600, check: bool = False) -> subprocess.CompletedProcess:
    args = ["ssh", "-n", *SSH_OPTS] + (["-A"] if agent else []) + [host, cmd]
    r = subprocess.run(args, capture_output=True, text=True, timeout=timeout)
    if check and r.returncode != 0:
        raise RuntimeError(f"ssh {host} failed rc={r.returncode}: {r.stderr.strip()[:500]}")
    return r


def node_host(node: str) -> str:
    return NODES[node][0]


def ensure_agent() -> None:
    """Start an ssh-agent holding the gateway key so nodes can rsync from BHS64 / the master from nodes."""
    if os.environ.get("SSH_AUTH_SOCK") and subprocess.run(["ssh-add", "-l"], capture_output=True).returncode == 0:
        return
    out = subprocess.run(["ssh-agent", "-s"], capture_output=True, text=True, check=True).stdout
    for part in out.replace("\n", ";").split(";"):
        part = part.strip()
        if part.startswith("SSH_AUTH_SOCK=") or part.startswith("SSH_AGENT_PID="):
            k, v = part.split("=", 1)
            os.environ[k] = v
    subprocess.run(["ssh-add", str(HOME / ".ssh" / "id_ed25519")], check=True, capture_output=True)
    log("ssh-agent started with id_ed25519")


# ----------------------------------------------------------------------------- stages
def stage_sync() -> None:
    excl = ["--exclude", ".DS_Store", "--exclude", "__pycache__", "--exclude", ".venv", "--exclude", "data/",
            "--exclude", "project-db/", "--exclude", "analysis/", "--exclude", "notebooks/",
            "--exclude", "outputs/", "--exclude", "UQ_WP_template/", "--exclude", "*.zip", "--exclude", ".git/"]
    meta = ["universe.json", "shards.json", "russell3000_constituents.json", "sp500_constituents.json", "cik_mapping.json"]

    def push(target: str) -> str:
        subprocess.run(["rsync", "-a", *excl, "-e", "ssh " + " ".join(SSH_OPTS), f"{CODE}/", f"{target}:{REMOTE_ROOT}/"],
                       check=True, capture_output=True)
        subprocess.run(["ssh", *SSH_OPTS, target, f"mkdir -p {REMOTE_ROOT}/data/metadata/shards {REMOTE_ROOT}/data/logs"],
                       check=True, capture_output=True)
        subprocess.run(["rsync", "-a", "-e", "ssh " + " ".join(SSH_OPTS),
                        *[str(CODE / "data" / "metadata" / m) for m in meta if (CODE / "data" / "metadata" / m).exists()],
                        f"{target}:{REMOTE_ROOT}/data/metadata/"], check=True, capture_output=True)
        subprocess.run(["rsync", "-a", "-e", "ssh " + " ".join(SSH_OPTS), f"{CODE}/data/metadata/shards/",
                        f"{target}:{REMOTE_ROOT}/data/metadata/shards/"], check=True, capture_output=True)
        return target

    with ThreadPoolExecutor(8) as ex:
        for t in ex.map(push, [node_host(n) for n in NODES] + [BHS64]):
            log(f"sync -> {t} ok")


def stage_downloads() -> None:
    shards = json.loads((CODE / "data" / "metadata" / "shards.json").read_text())["nodes"]
    lan_nodes = [n for n in NODES if shards.get(n, {}).get("self")]
    rate = LAN_RATE_TOTAL / max(1, len(lan_nodes))
    for n in lan_nodes:
        r = ssh(node_host(n), f"ls {REMOTE_ROOT}/data/logs/stage.download_self.done 2>/dev/null; "
                              f"pgrep -f 'shards/[{n[0]}]{n[1:]}.self.txt' >/dev/null && echo RUNNING")
        if "done" in r.stdout or "RUNNING" in r.stdout:
            log(f"downloads: {n} already {'done' if 'done' in r.stdout else 'running'}")
            continue
        ssh(node_host(n), f"cd {REMOTE_ROOT}; RATE={rate:.3f} DL_WORKERS=6 nohup scripts/cluster/run_shard.sh {n} download_self "
                          f"> data/logs/nohup.download_self.out 2>&1 < /dev/null &", check=True)
        log(f"downloads: {n} started self list ({len(shards[n]['self'])} CIKs @ {rate:.2f} req/s)")

    bhs_nodes = [n for n in NODES if shards.get(n, {}).get("bhs64")]
    brate = BHS64_RATE_TOTAL / max(1, len(bhs_nodes))
    for n in bhs_nodes:
        r = ssh(BHS64, f"cat {BHS64_DL}/{n}/DONE 2>/dev/null; pgrep -f 'shards/[{n[0]}]{n[1:]}.bhs64.txt' >/dev/null && echo RUNNING")
        if "rc=" in r.stdout or "RUNNING" in r.stdout:
            log(f"downloads: BHS64/{n} already {'done' if 'rc=' in r.stdout else 'running'}")
            continue
        cmd = (f"mkdir -p {BHS64_DL}/{n}; cd {REMOTE_ROOT}; PYTHONUNBUFFERED=1 nohup sh -c "
               f"'.venv/bin/python scripts/download_filings.py --ciks-file data/metadata/shards/{n}.bhs64.txt "
               f"--rate {brate:.3f} --workers 6 --raw-dir ~/{BHS64_DL}/{n}/raw --log ~/{BHS64_DL}/{n}/download_log.jsonl "
               f"--submissions-dir ~/{BHS64_DL}/{n}/submissions; echo rc=$? > ~/{BHS64_DL}/{n}/DONE' "
               f"> ~/{BHS64_DL}/{n}/download.log 2>&1 < /dev/null &")
        ssh(BHS64, cmd, check=True)
        log(f"downloads: BHS64 started for {n} ({len(shards[n]['bhs64'])} CIKs @ {brate:.2f} req/s)")


def downloads_status() -> dict:
    shards = json.loads((CODE / "data" / "metadata" / "shards.json").read_text())["nodes"]
    st = {}
    for n in NODES:
        r = ssh(node_host(n), f"cd {REMOTE_ROOT}/data/logs 2>/dev/null && ls stage.download_self.* 2>/dev/null; "
                              f"tail -1 {REMOTE_ROOT}/data/logs/download_self.log 2>/dev/null | cut -c1-160")
        st[n] = r.stdout.strip().replace("\n", " | ")
    for n in NODES:
        if not shards.get(n, {}).get("bhs64"):
            continue
        r = ssh(BHS64, f"cat {BHS64_DL}/{n}/DONE 2>/dev/null; echo \"$(wc -l < {BHS64_DL}/{n}/download_log.jsonl 2>/dev/null) filings\"; tail -1 {BHS64_DL}/{n}/download.log 2>/dev/null | cut -c1-160")
        st[f"BHS64/{n}"] = r.stdout.strip().replace("\n", " | ")
    return st


def wait_downloads(poll: int = 300) -> None:
    shards = json.loads((CODE / "data" / "metadata" / "shards.json").read_text())["nodes"]
    while True:
        pending = []
        for n in NODES:
            if shards.get(n, {}).get("self"):
                r = ssh(node_host(n), f"ls {REMOTE_ROOT}/data/logs/stage.download_self.done {REMOTE_ROOT}/data/logs/stage.download_self.failed 2>/dev/null")
                if "done" not in r.stdout and "failed" not in r.stdout:
                    pending.append(n)
                elif "failed" in r.stdout:
                    log(f"WARNING {n} download_self FAILED (partial data kept; rerun 'downloads' later)")
            if shards.get(n, {}).get("bhs64"):
                r = ssh(BHS64, f"cat {BHS64_DL}/{n}/DONE 2>/dev/null")
                if "rc=" not in r.stdout:
                    pending.append(f"BHS64/{n}")
        if not pending:
            log("downloads: all finished")
            return
        log(f"downloads: waiting on {pending}")
        time.sleep(poll)


def stage_ship() -> None:
    """Move the BHS64-downloaded halves onto their nodes (node pulls over SSH with the gateway agent)."""
    ensure_agent()
    shards = json.loads((CODE / "data" / "metadata" / "shards.json").read_text())["nodes"]
    targets = [n for n in NODES if shards.get(n, {}).get("bhs64")]

    def pull(n: str) -> str:
        r = ssh(node_host(n), f"ls {REMOTE_ROOT}/data/logs/stage.ship.done 2>/dev/null")
        if r.returncode == 0:
            return f"{n}: already shipped"
        cmd = (f"set -e; cd {REMOTE_ROOT}; mkdir -p data/raw/sec data/metadata/submissions; "
               f"rsync -a --partial -e 'ssh -o StrictHostKeyChecking=accept-new' {BHS64}:{BHS64_DL}/{n}/raw/ data/raw/sec/ && "
               f"rsync -a -e 'ssh -o StrictHostKeyChecking=accept-new' {BHS64}:{BHS64_DL}/{n}/submissions/ data/metadata/submissions/ && "
               f"rsync -a -e 'ssh -o StrictHostKeyChecking=accept-new' {BHS64}:{BHS64_DL}/{n}/download_log.jsonl data/metadata/download_log.bhs64.jsonl && "
               f"touch data/logs/stage.ship.done")
        r = ssh(node_host(n), cmd, agent=True, timeout=6 * 3600)
        return f"{n}: ship rc={r.returncode} {r.stderr.strip()[-200:]}"

    with ThreadPoolExecutor(len(targets) or 1) as ex:
        for msg in ex.map(pull, targets):
            log(msg)


def stage_pipeline() -> None:
    while True:  # the master needs the S&P 500 corpus (rsync'ed from the laptop) before it can parse
        r = ssh(node_host(MASTER), f"ls {REMOTE_ROOT}/data/logs/stage.master_corpus.done 2>/dev/null")
        if r.returncode == 0:
            break
        log("pipeline: waiting for stage.master_corpus.done on the master (S&P 500 corpus transfer)")
        time.sleep(300)
    for n in NODES:
        r = ssh(node_host(n), f"ls {REMOTE_ROOT}/data/logs/stage.pipeline.done {REMOTE_ROOT}/data/logs/stage.pipeline.running 2>/dev/null")
        if "done" in r.stdout or "running" in r.stdout:
            log(f"pipeline: {n} already {'done' if 'done' in r.stdout else 'running'}")
            continue
        workers = NODES[n][1]
        stage = "pipeline_master" if n == MASTER else "pipeline"
        ssh(node_host(n), f"cd {REMOTE_ROOT}; WORKERS={workers} nohup scripts/cluster/run_shard.sh {n} {stage} "
                          f"> data/logs/nohup.pipeline.out 2>&1 < /dev/null &", check=True)
        log(f"pipeline: {n} started ({workers} parse workers)")


def wait_pipeline(poll: int = 300) -> None:
    while True:
        pending, failed = [], []
        for n in NODES:
            r = ssh(node_host(n), f"ls {REMOTE_ROOT}/data/logs/ 2>/dev/null | grep '^stage\\.' | tr '\\n' ' '")
            marks = r.stdout
            if "stage.pipeline.done" in marks:
                continue
            if "stage.pipeline.failed" in marks:
                failed.append(n)
            else:
                pending.append(n)
        if not pending:
            log(f"pipeline: finished (failed: {failed})")
            return
        log(f"pipeline: waiting on {pending} failed={failed}")
        time.sleep(poll)


def stage_pull() -> None:
    """The master pulls every worker's export + embeddings with its own key (authorised on the workers)."""
    workers = [n for n in NODES if n != MASTER]
    e = "ssh -i ~/.ssh/id_ed25519 -o IdentitiesOnly=yes -o StrictHostKeyChecking=accept-new"

    def pull(n: str) -> str:
        ip = node_host(n)
        cmd = (f"set -e; cd {REMOTE_ROOT}; mkdir -p data/shards/{n}/export data/shards/{n}/embeddings; "
               f"rsync -a -e '{e}' {ip}:{REMOTE_ROOT}/data/export/ data/shards/{n}/export/ && "
               f"rsync -a -e '{e}' {ip}:{REMOTE_ROOT}/data/parquet/embeddings/ data/shards/{n}/embeddings/ && "
               f"test -f data/shards/{n}/export/EXPORT_DONE")
        r = ssh(node_host(MASTER), cmd, timeout=6 * 3600)
        return f"pull {n}: rc={r.returncode} {r.stderr.strip()[-200:]}"

    with ThreadPoolExecutor(4) as ex:
        for msg in ex.map(pull, workers):
            log(msg)


def stage_merge() -> None:
    m = node_host(MASTER)
    ssh(m, f"rm -f {REMOTE_ROOT}/data/logs/stage.merge.done {REMOTE_ROOT}/data/logs/stage.merge.failed")
    ssh(m, f"cd {REMOTE_ROOT}; nohup sh -c '.venv/bin/python scripts/merge_shards.py && "
           f".venv/bin/python scripts/build_factors.py && touch data/logs/stage.merge.done || touch data/logs/stage.merge.failed' "
           f"> data/logs/merge.log 2>&1 < /dev/null &", check=True)
    log("merge: started on master (merge_shards + build_factors)")
    while True:
        r = ssh(m, f"ls {REMOTE_ROOT}/data/logs/stage.merge.* 2>/dev/null; tail -1 {REMOTE_ROOT}/data/logs/merge.log | cut -c1-200")
        if "merge.done" in r.stdout or "merge.failed" in r.stdout:
            log(f"merge: {r.stdout.strip()}")
            return
        log(f"merge: running | {r.stdout.strip().splitlines()[-1] if r.stdout.strip() else ''}")
        time.sleep(300)


def stage_status() -> None:
    for n, (ip, _) in NODES.items():
        r = ssh(ip, f"cd {REMOTE_ROOT}/data/logs 2>/dev/null && ls | grep '^stage\\.' | tr '\\n' ' '; "
                    f"df -h / | tail -1 | awk '{{print $4\" free\"}}'; "
                    f"echo \"self-log: $(cat ~/{REMOTE_ROOT}/data/metadata/download_log.self.jsonl 2>/dev/null | wc -l | tr -d ' ') filings\"; "
                    f"for f in download_self parse embed export merge; do [ -f {REMOTE_ROOT}/data/logs/$f.log ] && tail -1 {REMOTE_ROOT}/data/logs/$f.log | cut -c1-150; done")
        print(f"--- {n}\n{r.stdout.strip()}")
    for k, v in downloads_status().items():
        if k.startswith("BHS64"):
            print(f"--- {k}: {v}")


STAGES = {
    "sync": stage_sync, "downloads": stage_downloads, "wait-downloads": wait_downloads, "ship": stage_ship,
    "pipeline": stage_pipeline, "wait-pipeline": wait_pipeline, "pull": stage_pull, "merge": stage_merge,
    "status": stage_status,
}


def main() -> None:
    if len(sys.argv) < 2 or sys.argv[1] not in STAGES and sys.argv[1] != "all":
        print(__doc__)
        sys.exit(1)
    what = sys.argv[1]
    if what == "all":
        for s in ["sync", "downloads", "wait-downloads", "ship", "pipeline", "wait-pipeline", "pull", "merge"]:
            log(f"##### stage {s}")
            STAGES[s]()
        log("##### ALL DONE")
    else:
        STAGES[what]()


if __name__ == "__main__":
    main()