#!/usr/bin/env python3
"""
DAARION Node State Discovery Script

This script collects and documents the current state of nodes in the DAARION network,
including Swapper models, DAGI Router agents, and database agents.

Usage:
    python scripts/discover_node_state.py --node node-1-hetzner-gex44
    python scripts/discover_node_state.py --node node-2-macbook-m4max
    python scripts/discover_node_state.py --all

Output:
    Creates markdown files in docs/users/nodes/NODE_STATE_{node_id}.md
"""

import argparse
import json
import os
import sys
from datetime import datetime
from pathlib import Path

try:
    import httpx
except ImportError:
    print("❌ httpx not installed. Run: pip install httpx")
    sys.exit(1)

# Configuration
CITY_API_URL = os.environ.get("CITY_API_URL", "https://daarion.space/api")
DB_CONNECTION = os.environ.get("DATABASE_URL", None)

KNOWN_NODES = [
    "node-1-hetzner-gex44",
    "node-2-macbook-m4max",
]

OUTPUT_DIR = Path(__file__).parent.parent / "docs" / "users" / "nodes"


def fetch_swapper_state(node_id: str) -> dict:
    """Fetch Swapper state from city-service API."""
    url = f"{CITY_API_URL}/node-internal/{node_id}/swapper"
    try:
        with httpx.Client(timeout=10.0) as client:
            resp = client.get(url)
            if resp.status_code == 200:
                return resp.json()
            else:
                return {"error": f"HTTP {resp.status_code}", "healthy": False, "models": []}
    except Exception as e:
        return {"error": str(e), "healthy": False, "models": []}


def fetch_dagi_router_agents(node_id: str) -> dict:
    """Fetch DAGI Router agents from city-service API."""
    url = f"{CITY_API_URL}/node-internal/{node_id}/dagi-router/agents"
    try:
        with httpx.Client(timeout=10.0) as client:
            resp = client.get(url)
            if resp.status_code == 200:
                return resp.json()
            else:
                return {"error": f"HTTP {resp.status_code}", "agents": []}
    except Exception as e:
        return {"error": str(e), "agents": []}


def fetch_node_info(node_id: str) -> dict:
    """Fetch node info from city-service API."""
    url = f"{CITY_API_URL}/nodes/{node_id}"
    try:
        with httpx.Client(timeout=10.0) as client:
            resp = client.get(url)
            if resp.status_code == 200:
                return resp.json()
            else:
                return {"error": f"HTTP {resp.status_code}"}
    except Exception as e:
        return {"error": str(e)}


def generate_markdown(node_id: str, swapper: dict, router: dict, node_info: dict) -> str:
    """Generate markdown report for a node."""
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    md = f"""# Node {node_id} — State

> Generated: {now}

## Node Info

| Field | Value |
|-------|-------|
| Node ID | `{node_id}` |
| Name | {node_info.get('name', 'N/A')} |
| Status | {node_info.get('status', 'N/A')} |
| Environment | {node_info.get('environment', 'N/A')} |
| Hostname | {node_info.get('hostname', 'N/A')} |

## Swapper

"""
    
    if swapper.get("error"):
        md += f"**Error:** {swapper['error']}\n\n"
    else:
        healthy = "✅ Healthy" if swapper.get("healthy") else "❌ Unhealthy"
        models_loaded = swapper.get("models_loaded", 0)
        models_total = swapper.get("models_total", 0)
        
        md += f"""- **Status:** {healthy}
- **Models Loaded:** {models_loaded} / {models_total}

### Models

| Name | Type | Loaded | VRAM (GB) |
|------|------|--------|-----------|
"""
        models = swapper.get("models", [])
        if models:
            for m in models:
                loaded = "✅" if m.get("loaded") else "❌"
                vram = m.get("vram_gb", "N/A")
                md += f"| {m.get('name', 'N/A')} | {m.get('type', 'N/A')} | {loaded} | {vram} |\n"
        else:
            md += "| (no models) | | | |\n"
    
    md += "\n## DAGI Router Agents\n\n"
    
    if router.get("error"):
        md += f"**Error:** {router['error']}\n\n"
    else:
        total = router.get("total", 0)
        active = router.get("active", 0)
        phantom = router.get("phantom", 0)
        stale = router.get("stale", 0)
        
        md += f"""- **Total Agents:** {total}
- **Active:** {active}
- **Phantom:** {phantom}
- **Stale:** {stale}

### Agent List

| ID | Name | Kind | Status | Runtime | Has DB Record |
|----|------|------|--------|---------|---------------|
"""
        agents = router.get("agents", [])
        if agents:
            for a in agents:
                md += f"| {a.get('id', 'N/A')} | {a.get('name', 'N/A')} | {a.get('kind', 'N/A')} | {a.get('status', 'N/A')} | {a.get('runtime', 'N/A')} | {a.get('has_db_record', 'N/A')} |\n"
        else:
            md += "| (no agents) | | | | | |\n"
    
    md += f"""

---

## Raw Data

<details>
<summary>Swapper JSON</summary>

```json
{json.dumps(swapper, indent=2, default=str)}
```

</details>

<details>
<summary>DAGI Router JSON</summary>

```json
{json.dumps(router, indent=2, default=str)}
```

</details>
"""
    
    return md


def discover_node(node_id: str, output_dir: Path) -> None:
    """Discover and document state for a single node."""
    print(f"\n🔍 Discovering node: {node_id}")
    
    # Fetch data
    print(f"  📡 Fetching Swapper state...")
    swapper = fetch_swapper_state(node_id)
    
    print(f"  📡 Fetching DAGI Router agents...")
    router = fetch_dagi_router_agents(node_id)
    
    print(f"  📡 Fetching node info...")
    node_info = fetch_node_info(node_id)
    
    # Generate markdown
    md = generate_markdown(node_id, swapper, router, node_info)
    
    # Write to file
    output_dir.mkdir(parents=True, exist_ok=True)
    output_file = output_dir / f"NODE_STATE_{node_id}.md"
    output_file.write_text(md)
    
    print(f"  ✅ Saved to: {output_file}")
    
    # Print summary
    print(f"\n  📊 Summary:")
    print(f"     Swapper: {'healthy' if swapper.get('healthy') else 'unhealthy'}, {swapper.get('models_total', 0)} models")
    print(f"     Agents: {router.get('total', 0)} total, {router.get('active', 0)} active")


def main():
    global CITY_API_URL
    
    parser = argparse.ArgumentParser(description="Discover DAARION node state")
    parser.add_argument("--node", type=str, help="Specific node ID to discover")
    parser.add_argument("--all", action="store_true", help="Discover all known nodes")
    parser.add_argument("--output", type=str, default=str(OUTPUT_DIR), help="Output directory")
    parser.add_argument("--api-url", type=str, default=CITY_API_URL, help="City API URL")
    
    args = parser.parse_args()
    
    if not args.node and not args.all:
        print("❌ Please specify --node <node_id> or --all")
        parser.print_help()
        sys.exit(1)
    
    CITY_API_URL = args.api_url
    output_dir = Path(args.output)
    
    print(f"🚀 DAARION Node State Discovery")
    print(f"   API URL: {CITY_API_URL}")
    print(f"   Output: {output_dir}")
    
    nodes_to_discover = KNOWN_NODES if args.all else [args.node]
    
    for node_id in nodes_to_discover:
        discover_node(node_id, output_dir)
    
    print(f"\n✨ Discovery complete!")


if __name__ == "__main__":
    main()