#!/usr/bin/env python3
"""Atlas CLI — dispatch jobs and control the mesh.

Usage:
    atlas run <node> "<command>" [--agent shell|claude|python] [--cwd PATH] [--timeout SEC]
    atlas status                        — show heartbeats from all mesh nodes
    atlas where                         — alias for status with extra detail
    atlas wake <node>                   — Wake-on-LAN via tailscale
    atlas notify "<text>" [--target N] [--title T] [--push]
    atlas todo "<text>" [--due YYYY-MM-DD] [--category CAT]
    atlas brain "<question>"            — one-shot ask the brain (Claude with mesh context)

Examples:
    atlas run atlasworkspace "df -h"
    atlas notify "deploy passed" --push
    atlas todo "renew Hetzner SSL" --due 2026-06-01 --category infra
    atlas brain "what overdue invoices do we have?"
    atlas where
"""
from __future__ import annotations
import argparse
import asyncio
import datetime as dt
import json
import os
import socket
import subprocess
import sys
import time
import uuid
from pathlib import Path

try:
    import nats
except ImportError:
    print("Installing nats-py...", file=sys.stderr)
    subprocess.check_call([sys.executable, "-m", "pip", "install", "--quiet", "nats-py"])
    import nats

def load_config():
    env = dict(os.environ)
    cfg_file = Path.home() / ".config/atlas/orchestrator.env"
    if cfg_file.exists():
        for line in cfg_file.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            if "=" in line:
                k, v = line.split("=", 1)
                env.setdefault(k, v)
    # Pushover keys come from the same env file or atlas-secrets unlocked
    secrets = Path.home() / ".config/atlas-secrets/unlocked/atlas-api-secrets.env"
    if secrets.exists():
        for line in secrets.read_text(errors="ignore").splitlines():
            line = line.strip()
            if not line or line.startswith("#") or line.startswith("["):
                continue
            if "=" in line:
                k, v = line.split("=", 1)
                env.setdefault(k.strip(), v.strip())
    return env

CFG = load_config()
NATS_URL = CFG.get("ATLAS_NATS_URL", "nats://atlas-01:4222")
NATS_TOKEN = CFG.get("ATLAS_NATS_TOKEN", "")
PUSHOVER_USER = CFG.get("PUSHOVER_USER_KEY", "")
PUSHOVER_TOKEN = CFG.get("PUSHOVER_APP_TOKEN", "")
HOSTNAME = socket.gethostname().lower()
DAILY_DIR = Path.home() / "Sync/atlas-artifacts/atlas-os/daily"
DAILY_DIR.mkdir(parents=True, exist_ok=True)

# ----------------------------- helpers -----------------------------

async def connect():
    opts = {"servers": [NATS_URL], "reconnect_time_wait": 2}
    if NATS_TOKEN:
        opts["token"] = NATS_TOKEN
    return await nats.connect(**opts)

def push_phone(title: str, body: str, priority: int = 0):
    """Pushover one-way alert. priority=1 bypasses quiet hours."""
    if not (PUSHOVER_USER and PUSHOVER_TOKEN):
        return False
    try:
        import urllib.request, urllib.parse
        data = urllib.parse.urlencode({
            "token": PUSHOVER_TOKEN, "user": PUSHOVER_USER,
            "title": title, "message": body, "priority": str(priority),
        }).encode()
        urllib.request.urlopen("https://api.pushover.net/1/messages.json", data, timeout=10)
        return True
    except Exception as e:
        print(f"pushover failed: {e}", file=sys.stderr)
        return False

def toast_local(title: str, body: str):
    """Windows toast via BurntToast. Best-effort, no-op on non-Windows."""
    if sys.platform != "win32":
        return
    cmd = ["powershell", "-NoProfile", "-Command",
           f"Import-Module BurntToast -ErrorAction SilentlyContinue; "
           f"if (Get-Command New-BurntToastNotification -ErrorAction SilentlyContinue) "
           f"{{ New-BurntToastNotification -Text {json.dumps(title)},{json.dumps(body)} }}"]
    subprocess.run(cmd, capture_output=True, timeout=15)

# ----------------------------- run -----------------------------

async def cmd_run(args):
    nc = await connect()
    job = {
        "id": str(uuid.uuid4()),
        "agent": args.agent,
        "command": args.command,
        "cwd": args.cwd,
        "timeout_s": args.timeout,
        "env": {},
        "originator": HOSTNAME,
    }
    topic = f"atlas.jobs.{args.node}"
    print(f"-> {topic} ({job['id']})", file=sys.stderr)
    inbox = "_INBOX." + uuid.uuid4().hex
    fut = asyncio.Future()
    async def on_reply(msg):
        if not fut.done():
            fut.set_result(msg)
    sub = await nc.subscribe(inbox, cb=on_reply)
    await nc.publish(topic, json.dumps(job).encode(), reply=inbox)
    try:
        msg = await asyncio.wait_for(fut, timeout=args.timeout + 10)
        result = json.loads(msg.data.decode())
        if result.get("stdout"):
            print(result["stdout"], end="" if result["stdout"].endswith("\n") else "\n")
        if result.get("stderr"):
            print(result["stderr"], file=sys.stderr, end="" if result["stderr"].endswith("\n") else "\n")
        return result.get("exit_code", 0) or 0
    except asyncio.TimeoutError:
        print(f"TIMEOUT — {args.node} did not reply within {args.timeout + 10}s", file=sys.stderr)
        return 124
    finally:
        await sub.unsubscribe()
        await nc.drain()

# ----------------------------- status / where -----------------------------

async def cmd_status(args):
    nc = await connect()
    nodes = {}
    async def on_hb(msg):
        try:
            d = json.loads(msg.data.decode())
            nodes[d["node"]] = d
        except Exception:
            pass
    sub = await nc.subscribe("atlas.heartbeat.>", cb=on_hb)
    window = 35 if not getattr(args, "fast", False) else 5
    await asyncio.sleep(window)
    await sub.unsubscribe()
    await nc.drain()
    if not nodes:
        print(f"(no heartbeats heard in {window}s)")
        return 1
    print(f"{'node':<20} {'platform':<10} {'last seen':<12} {'version'}")
    print("-" * 60)
    for n, d in sorted(nodes.items()):
        ago = time.time() - d.get("ts", 0)
        print(f"{n:<20} {d.get('platform','?'):<10} {ago:>4.0f}s ago    {d.get('version','?')}")
    return 0

# ----------------------------- notify -----------------------------

async def cmd_notify(args):
    title = args.title or "Atlas"
    body = args.text
    target = args.target or HOSTNAME

    pushed = False
    if args.push:
        pushed = push_phone(title, body, priority=1 if args.urgent else 0)

    if target == HOSTNAME or target == "local":
        toast_local(title, body)
        print(f"toast: local{' + push' if pushed else ''}")
        return 0

    # Remote target via NATS
    nc = await connect()
    try:
        cmd = (
            "Import-Module BurntToast -ErrorAction SilentlyContinue; "
            f"New-BurntToastNotification -Text {json.dumps(title)},{json.dumps(body)}"
        )
        job = {
            "id": str(uuid.uuid4()),
            "agent": "shell",
            "command": cmd,
            "timeout_s": 15,
            "originator": HOSTNAME,
        }
        await nc.publish(f"atlas.jobs.{target}", json.dumps(job).encode())
        print(f"toast: dispatched to {target}{' + push' if pushed else ''}")
    finally:
        await nc.drain()
    return 0

# ----------------------------- todo -----------------------------

async def cmd_todo(args):
    today = dt.date.today().isoformat()
    f = DAILY_DIR / f"{today}.md"
    section = f"## [{dt.datetime.now().strftime('%H:%M')}] todo: {args.category or 'general'}\n"
    body = f"- [ ] {args.text}"
    if args.due:
        body += f"  *(due {args.due})*"
    block = f"\n{section}\n{body}\n"
    if not f.exists():
        f.write_text(f"# {today}\n")
    with f.open("a", encoding="utf-8") as fh:
        fh.write(block)
    print(f"appended to {f}")
    return 0

# ----------------------------- brain -----------------------------

async def cmd_brain(args):
    """One-shot ask: dispatch job to atlas-01 brain runner with a question."""
    nc = await connect()
    try:
        job = {
            "id": str(uuid.uuid4()),
            "agent": "shell",
            "command": (
                "set -a; . /etc/atlas/orchestrator.env 2>/dev/null; "
                "[ -f /etc/atlas/brain.env ] && . /etc/atlas/brain.env; set +a; "
                f"/opt/atlas/orchestrator/.venv/bin/python /opt/atlas/brain/atlas-brain.py --once-prompt {json.dumps(args.question)}"
            ),
            "timeout_s": 60,
            "originator": HOSTNAME,
        }
        inbox = "_INBOX." + uuid.uuid4().hex
        fut = asyncio.Future()
        async def on_reply(msg):
            if not fut.done():
                fut.set_result(msg)
        sub = await nc.subscribe(inbox, cb=on_reply)
        await nc.publish("atlas.jobs.atlas-01", json.dumps(job).encode(), reply=inbox)
        try:
            m = await asyncio.wait_for(fut, timeout=70)
            r = json.loads(m.data.decode())
            print(r.get("stdout", "(no output)"))
            return 0
        except asyncio.TimeoutError:
            print("brain: timeout (atlas-01 did not reply in 70s)", file=sys.stderr)
            return 124
        finally:
            await sub.unsubscribe()
    finally:
        await nc.drain()

# ----------------------------- main -----------------------------

def main():
    p = argparse.ArgumentParser(prog="atlas")
    sub = p.add_subparsers(dest="cmd", required=True)

    pr = sub.add_parser("run", help="run a command on a node")
    pr.add_argument("node")
    pr.add_argument("command")
    pr.add_argument("--agent", default="shell", choices=["shell", "claude", "claude-code", "python"])
    pr.add_argument("--cwd")
    pr.add_argument("--timeout", type=int, default=300)
    pr.set_defaults(func=cmd_run)

    ps = sub.add_parser("status", help="show node heartbeats")
    ps.add_argument("--fast", action="store_true", help="only wait 5s instead of 35s")
    ps.set_defaults(func=cmd_status)
    pw = sub.add_parser("where", help="alias for status")
    pw.add_argument("--fast", action="store_true")
    pw.set_defaults(func=cmd_status)
    pn = sub.add_parser("nodes", help="alias for status")
    pn.add_argument("--fast", action="store_true")
    pn.set_defaults(func=cmd_status)

    pwk = sub.add_parser("wake", help="wake a node")
    pwk.add_argument("node")
    pwk.set_defaults(func=lambda a: cmd_run(argparse.Namespace(
        node=a.node, command="echo woke", agent="shell", cwd=None, timeout=10)))

    pno = sub.add_parser("notify", help="push a notification (toast + optional phone)")
    pno.add_argument("text")
    pno.add_argument("--title", default=None)
    pno.add_argument("--target", default=None, help="remote node hostname; default local")
    pno.add_argument("--push", action="store_true", help="also send Pushover phone alert")
    pno.add_argument("--urgent", action="store_true", help="priority=1 on phone push")
    pno.set_defaults(func=cmd_notify)

    pt = sub.add_parser("todo", help="append a todo item to today's daily.md")
    pt.add_argument("text")
    pt.add_argument("--due", default=None)
    pt.add_argument("--category", default=None)
    pt.set_defaults(func=cmd_todo)

    pb = sub.add_parser("brain", help="ask the brain a one-shot question")
    pb.add_argument("question")
    pb.set_defaults(func=cmd_brain)

    args = p.parse_args()
    rc = asyncio.run(args.func(args))
    sys.exit(rc or 0)

if __name__ == "__main__":
    main()