#!/usr/bin/env python3
"""
⚡ Zoth Swarm Agent Bus CLI & Coordination Engine
Allows AI agents (Grok, Antigravity, Hermes, Claude, local models) and human operators
to register heartbeats, claim project locks, broadcast messages, and inspect active workers.
"""

import os
import sys
import json
import time
import argparse
from datetime import datetime, timezone

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CLAIMS_DIR = os.path.join(BASE_DIR, "claims")
INBOX_DIR = os.path.join(BASE_DIR, "inbox")
BOARD_DIR = os.path.join(BASE_DIR, "board")
LOGS_DIR = os.path.join(BASE_DIR, "logs")
STATUS_JSON = os.path.join(BOARD_DIR, "status.json")
HEARTBEATS_JSON = os.path.join(BOARD_DIR, "heartbeats.json")
MESSAGES_JSON = os.path.join(BOARD_DIR, "messages.json")

os.makedirs(CLAIMS_DIR, exist_ok=True)
os.makedirs(INBOX_DIR, exist_ok=True)
os.makedirs(BOARD_DIR, exist_ok=True)
os.makedirs(LOGS_DIR, exist_ok=True)

def get_utc_now():
    return datetime.now(timezone.utc).isoformat()

def load_json(filepath, default=None):
    if default is None:
        default = {}
    if not os.path.exists(filepath):
        return default
    try:
        with open(filepath, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return default

def save_json(filepath, data):
    with open(filepath, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)

def log_activity(agent, action, details):
    entry = {
        "timestamp": get_utc_now(),
        "agent": agent,
        "action": action,
        "details": details
    }
    with open(os.path.join(LOGS_DIR, "activity.jsonl"), "a", encoding="utf-8") as f:
        f.write(json.dumps(entry) + "\n")

# Commands
def cmd_heartbeat(agent, task, capabilities="", status="active"):
    heartbeats = load_json(HEARTBEATS_JSON, {})
    heartbeats[agent] = {
        "agent": agent,
        "task": task,
        "status": status,
        "capabilities": capabilities,
        "last_seen": get_utc_now(),
        "pid": os.getpid()
    }
    save_json(HEARTBEATS_JSON, heartbeats)
    log_activity(agent, "heartbeat", {"task": task, "status": status})
    print(f"✅ [Heartbeat] Agent '{agent}' registered: {task} ({status})")

def cmd_post(from_agent, to_agent, message, priority="normal"):
    msg_id = f"{int(time.time() * 1000)}-{from_agent}"
    msg_obj = {
        "id": msg_id,
        "from": from_agent,
        "to": to_agent,
        "message": message,
        "priority": priority,
        "timestamp": get_utc_now()
    }

    # Save to global message feed
    messages = load_json(MESSAGES_JSON, [])
    messages.append(msg_obj)
    # Keep last 200 messages
    if len(messages) > 200:
        messages = messages[-200:]
    save_json(MESSAGES_JSON, messages)

    # Also save markdown to recipient inbox
    target_inbox = os.path.join(INBOX_DIR, f"to-{to_agent}")
    os.makedirs(target_inbox, exist_ok=True)
    md_file = os.path.join(target_inbox, f"{datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')}-{from_agent}.md")
    with open(md_file, "w", encoding="utf-8") as f:
        f.write(f"""---
from: {from_agent}
to: {to_agent}
priority: {priority}
created: {get_utc_now()}
---

{message}
""")

    log_activity(from_agent, "post_message", {"to": to_agent, "message": message[:80]})
    print(f"📨 [Message Sent] From: @{from_agent} -> To: @{to_agent} | ID: {msg_id}")

def cmd_claim(agent, project, scope=None, note=""):
    if scope is None:
        scope = ["all"]
    claim_path = os.path.join(CLAIMS_DIR, f"{project}.json")
    claim_obj = {
        "project": project,
        "owner": agent,
        "scope": scope,
        "status": "in_progress",
        "started_at": get_utc_now(),
        "note": note
    }
    save_json(claim_path, claim_obj)
    log_activity(agent, "claim_project", {"project": project, "note": note})
    print(f"🔒 [Lock Claimed] Agent '@{agent}' claimed project '{project}'")

def cmd_release(agent, project):
    claim_path = os.path.join(CLAIMS_DIR, f"{project}.json")
    if os.path.exists(claim_path):
        claim_data = load_json(claim_path)
        claim_data["status"] = "released"
        claim_data["released_at"] = get_utc_now()
        save_json(claim_path, claim_data)
        log_activity(agent, "release_project", {"project": project})
        print(f"🔓 [Lock Released] Project '{project}' released by '@{agent}'")
    else:
        print(f"⚠️ No active claim found for '{project}'")

def cmd_who():
    heartbeats = load_json(HEARTBEATS_JSON, {})
    print("=" * 60)
    print("🤖 ZOTH SWARM: ACTIVE AGENT REGISTRY")
    print("=" * 60)
    if not heartbeats:
        print("No active agent heartbeats registered.")
        return
    for name, info in heartbeats.items():
        print(f"• @{name.upper()} [{info.get('status', 'unknown')}]")
        print(f"   Task:      {info.get('task', 'N/A')}")
        print(f"   Last Seen: {info.get('last_seen', 'N/A')}")
        if info.get('capabilities'):
            print(f"   Caps:      {info.get('capabilities')}")
        print()

def cmd_messages(limit=15):
    messages = load_json(MESSAGES_JSON, [])
    print("=" * 60)
    print("📬 ZOTH SWARM: LIVE MESSAGE FEED")
    print("=" * 60)
    if not messages:
        print("Inbox is empty.")
        return
    for msg in messages[-limit:]:
        print(f"[{msg['timestamp'][11:19]}] @{msg['from']} ➔ @{msg['to']}: {msg['message']}")

def cmd_claims():
    print("=" * 60)
    print("🔒 ZOTH SWARM: ACTIVE PROJECT CLAIMS")
    print("=" * 60)
    files = [f for f in os.listdir(CLAIMS_DIR) if f.endswith(".json")]
    if not files:
        print("No active project claims.")
        return
    for f in files:
        c = load_json(os.path.join(CLAIMS_DIR, f))
        print(f"• Project: {c.get('project')} | Owner: @{c.get('owner')} | Status: {c.get('status')}")
        print(f"  Note:    {c.get('note', '')}")
        print(f"  Started: {c.get('started_at')}")
        print()

def main():
    parser = argparse.ArgumentParser(description="Zoth Swarm Agent Bus")
    subparsers = parser.add_subparsers(dest="cmd", help="Command")

    # heartbeat
    hb_p = subparsers.add_parser("heartbeat", help="Register/refresh agent heartbeat")
    hb_p.add_argument("--agent", required=True, help="Agent identity (e.g. grok, antigravity)")
    hb_p.add_argument("--task", required=True, help="Current task description")
    hb_p.add_argument("--capabilities", default="", help="Capabilities or endpoints (e.g. 'FastAPI :8484')")
    hb_p.add_argument("--status", default="active", help="active / busy / idle")

    # post
    post_p = subparsers.add_parser("post", help="Post a message to another agent or @all")
    post_p.add_argument("--from", dest="from_agent", required=True, help="Sender identity")
    post_p.add_argument("--to", default="all", help="Recipient identity or 'all'")
    post_p.add_argument("--msg", required=True, help="Message body")
    post_p.add_argument("--priority", default="normal", help="normal / high / urgent")

    # claim
    claim_p = subparsers.add_parser("claim", help="Claim lock on a project")
    claim_p.add_argument("--agent", required=True, help="Agent identity")
    claim_p.add_argument("--project", required=True, help="Project slug")
    claim_p.add_argument("--note", default="", help="Note or scope")

    # release
    rel_p = subparsers.add_parser("release", help="Release a project lock")
    rel_p.add_argument("--agent", required=True, help="Agent identity")
    rel_p.add_argument("--project", required=True, help="Project slug")

    # who
    subparsers.add_parser("who", help="List active agents and their current tasks")

    # read
    read_p = subparsers.add_parser("read", help="Read recent swarm messages")
    read_p.add_argument("--limit", type=int, default=15, help="Number of messages to display")

    # claims
    subparsers.add_parser("claims", help="List all project claims")

    args = parser.parse_args()

    if args.cmd == "heartbeat":
        cmd_heartbeat(args.agent, args.task, args.capabilities, args.status)
    elif args.cmd == "post":
        cmd_post(args.from_agent, args.to, args.msg, args.priority)
    elif args.cmd == "claim":
        cmd_claim(args.agent, args.project, note=args.note)
    elif args.cmd == "release":
        cmd_release(args.agent, args.project)
    elif args.cmd == "who":
        cmd_who()
    elif args.cmd == "read":
        cmd_messages(args.limit)
    elif args.cmd == "claims":
        cmd_claims()
    else:
        parser.print_help()

if __name__ == "__main__":
    main()
