"""
afs.livenet.node — Spec 22 P1

Python compute node for Living Networks. Connects to the WS broker as a
persistent process, receives Sim/Gate module assignments, runs them using
the full Python AFS engine, and sends results back via WebSocket.

Usage:
    cd ~/datasphere/afs-engine
    python3 -m afs.livenet.node --server ws://localhost:3000 --network <networkId> --lobby
    python3 -m afs.livenet.node --server ws://localhost:3000 --session <sessionId>
"""

import asyncio
import json
import argparse
import sys
import time
import platform
import os
import urllib.request
from urllib.parse import urlparse

try:
    import websockets
except ImportError:
    print("[node] ERROR: websockets not installed.  Run: pip install websockets")
    sys.exit(1)

from ..core import AfsEngine, AfsSimulationConfig
from ..compiler import FieldCompiler
from ..analysis.cluster_detection import detect_clusters, compute_concept_weights

# Prevents FieldCompiler from spinning up a DomainAnalyzer (needs OpenAI key)
_DUMMY_ANALYZER = type("NoOpAnalyzer", (), {})()

# Lite sim parameters — same as Gate mode in the browser
NUM_CELLS       = 30
NUM_GENERATIONS = 28
TICKS_PER_GEN   = 100


# ─── Utilities ────────────────────────────────────────────────────────────────

def _get_capabilities():
    caps = {
        "node_type": "python",
        "gpu": False,
        "gpu_backend": None,
        "afs_version": "0.3.0",
        "cpu_cores": os.cpu_count() or 1,
        "platform": platform.system(),
    }
    try:
        import wgpu  # noqa: F401
        caps["gpu"] = True
        caps["gpu_backend"] = "wgpu"
    except Exception:
        pass
    return caps


def _http_base(ws_url: str) -> str:
    """ws://host → http://host  |  wss://host → https://host"""
    parsed = urlparse(ws_url)
    scheme = "https" if parsed.scheme == "wss" else "http"
    return f"{scheme}://{parsed.netloc}"


def _http_get(base_url: str, path: str) -> dict:
    url = f"{base_url}{path}"
    req = urllib.request.Request(url, headers={"Accept": "application/json"})
    with urllib.request.urlopen(req, timeout=15) as resp:
        return json.loads(resp.read().decode())


# ─── Simulation ───────────────────────────────────────────────────────────────

def _normalize_world_model(world_model: dict) -> dict:
    """
    Livenet stores world models with 'factors' (flat list of names).
    compile_from_world_model expects 'concepts' with dataDictionary.
    This bridges the two formats without mutating the original.
    """
    if "concepts" in world_model:
        return world_model  # already in engine format
    wm = dict(world_model)
    wm["concepts"] = [
        {
            "name": f.get("name", f) if isinstance(f, dict) else f,
            "description": f.get("description", "") if isinstance(f, dict) else "",
            "dataDictionary": {"validRange": [0, 100]},
        }
        for f in world_model.get("factors", [])
    ]
    return wm


def _run_sim(world_model: dict, seed_centroid=None, seed_spread: float = 0.15) -> list:
    """Compile world model and run a lite AFS simulation. Returns list of strategy dicts."""
    compiler = FieldCompiler(analyzer=_DUMMY_ANALYZER)
    compile_result = compiler.compile_from_world_model(_normalize_world_model(world_model))
    field    = compile_result["field"]
    concepts = compile_result.get("concepts", [])

    config = AfsSimulationConfig()
    config.engine.initial_cell_count = NUM_CELLS
    config.engine.max_population     = NUM_CELLS * 3

    engine = AfsEngine(field, config, use_gpu=False)
    engine.total_ticks = NUM_GENERATIONS * TICKS_PER_GEN

    if seed_centroid:
        engine.initialize(seed_centroid=seed_centroid, seed_spread=seed_spread)
    else:
        engine.initialize()

    for _ in range(NUM_GENERATIONS):
        engine.run_generation(ticks=TICKS_PER_GEN)

    cell_positions    = engine.get_cell_positions()
    concept_positions = engine.get_concept_positions()
    concept_names     = engine.get_concept_names()

    clusters_raw = detect_clusters(cell_positions)
    total_cells  = sum(cl.cell_count for cl in clusters_raw) or 1

    strategies = []
    for cl in sorted(clusters_raw, key=lambda c: -c.cell_count):
        weights     = compute_concept_weights(cl, concept_positions, concept_names, engine.n_dims)
        top_concept = max(weights, key=weights.get) if weights else "Strategy"
        strategies.append({
            "name":            top_concept,
            "density":         cl.cell_count / total_cells,
            "centroid":        cl.centroid.tolist(),
            "cellCount":       cl.cell_count,
            "concept_weights": weights,
        })

    return strategies


def _extract_pipe_data(pipe_type: str, strategies: list) -> dict:
    """Shape strategy results into the typed payload expected by the dest module."""
    if not strategies:
        return {}
    top = strategies[0]
    if pipe_type == "cluster_seed":
        return {"centroid": top["centroid"], "density": top["density"], "basin_label": top["name"]}
    elif pipe_type == "concept_weight":
        weights = top.get("concept_weights", {})
        return [{"concept": k, "weight": v} for k, v in weights.items()]
    elif pipe_type == "quality_gate":
        return {"pass": top["density"] > 0.5, "density": top["density"], "basin_label": top["name"]}
    else:  # raw_data, world_model_delta, outcome_quality, etc.
        return {"strategies": strategies}


# ─── Main node loop ───────────────────────────────────────────────────────────

async def _execute_session(ws, network_id: str, session_id: str, node_id: str,
                           assigned_modules: list, http_base: str):
    """Fetch the network, run assigned modules, send results back via WS."""

    print(f"[node] fetching network design ({network_id})…")
    try:
        network = _http_get(http_base, f"/api/livenet/{network_id}")
    except Exception as e:
        print(f"[node] ERROR fetching network: {e}")
        return

    modules = {m["moduleId"]: m for m in (network.get("modules") or [])}
    pipes   = network.get("pipes") or []

    for module_id in assigned_modules:
        mod = modules.get(module_id)
        if not mod:
            print(f"[node] WARNING: module {module_id} not in network design, skipping")
            continue

        mod_type = mod.get("type")
        if mod_type not in ("Sim", "Gate"):
            print(f"[node] skipping {module_id} (type={mod_type})")
            continue

        world_model = (mod.get("config") or {}).get("worldModel")
        if not world_model:
            print(f"[node] {module_id}: no worldModel, skipping")
            await ws.send(json.dumps({
                "type":      "livenet_module_status",
                "sessionId": session_id,
                "moduleId":  module_id,
                "status":    "error",
                "detail":    "No worldModel configured",
            }))
            continue

        label = mod.get("label") or mod.get("name") or mod_type
        print(f"[node] running {module_id} ({label})…")

        await ws.send(json.dumps({
            "type":      "livenet_module_status",
            "sessionId": session_id,
            "moduleId":  module_id,
            "status":    "running",
            "detail":    "Python node executing",
        }))

        started_at = int(time.time() * 1000)
        try:
            strategies   = _run_sim(world_model)
            completed_at = int(time.time() * 1000)
            duration_ms  = completed_at - started_at

            top = strategies[0] if strategies else {}
            print(f"[node] {module_id}: {duration_ms}ms — "
                  f"top: {top.get('name', '?')} ({top.get('density', 0):.0%})")

            # Send typed pipe data for each outbound pipe
            outbound = [p for p in pipes if p.get("source_module") == module_id]
            for pipe in outbound:
                dest_id   = pipe.get("dest_module")
                pipe_type = pipe.get("pipe_type", "raw_data")
                data      = _extract_pipe_data(pipe_type, strategies)
                await ws.send(json.dumps({
                    "type":         "livenet_pipe_data",
                    "sessionId":    session_id,
                    "pipeId":       pipe.get("pipeId", f"{module_id}->{dest_id}"),
                    "destModuleId": dest_id,
                    "pipeType":     pipe_type,
                    "data":         data,
                }))
                print(f"[node]   → pipe {pipe_type} → {dest_id}")

            await ws.send(json.dumps({
                "type":      "livenet_module_status",
                "sessionId": session_id,
                "moduleId":  module_id,
                "status":    "complete",
                "detail":    f"{len(strategies)} strategies, {duration_ms}ms",
            }))

        except Exception as exc:
            import traceback
            print(f"[node] {module_id}: ERROR — {exc}")
            traceback.print_exc()
            await ws.send(json.dumps({
                "type":      "livenet_module_status",
                "sessionId": session_id,
                "moduleId":  module_id,
                "status":    "error",
                "detail":    str(exc),
            }))

    print(f"[node] session {session_id} complete")


async def run_node(server_url: str, network_id: str = None,
                   session_id: str = None, lobby: bool = False):
    http_base = _http_base(server_url)
    caps      = _get_capabilities()

    print(f"[node] connecting → {server_url}")
    print(f"[node] caps: {caps}")

    while True:  # reconnect loop for --lobby persistent mode
        try:
            async with websockets.connect(server_url) as ws:

                if lobby and network_id:
                    # ── Lobby mode: wait for the browser to click Run ──────────
                    await ws.send(json.dumps({
                        "type":        "livenet_lobby_join",
                        "networkId":   network_id,
                        "capabilities": caps,
                    }))
                    print(f"[node] waiting in lobby for network {network_id}…")

                    node_id          = None
                    assigned_session = None
                    assigned_modules = []
                    assigned_net     = network_id

                    async for raw in ws:
                        msg   = json.loads(raw)
                        mtype = msg.get("type")

                        if mtype == "livenet_lobby_joined":
                            node_id = msg.get("nodeId")
                            print(f"[node] lobby confirmed — nodeId={node_id}")

                        elif mtype == "livenet_session_ready":
                            node_id          = msg.get("nodeId", node_id)
                            assigned_session = msg.get("sessionId")
                            assigned_modules = msg.get("assignedModules", [])
                            assigned_net     = msg.get("networkId", network_id)
                            print(f"[node] session ready — {assigned_session}, modules={assigned_modules}")

                            await _execute_session(
                                ws, assigned_net, assigned_session,
                                node_id, assigned_modules, http_base
                            )

                            # Re-join lobby for the next session
                            print(f"[node] re-joining lobby…")
                            await ws.send(json.dumps({
                                "type":        "livenet_lobby_join",
                                "networkId":   network_id,
                                "capabilities": caps,
                            }))

                        elif mtype == "livenet_pipe_data":
                            # Inbound pipe from another node — not used in single-executor mode
                            pass

                elif session_id:
                    # ── Direct join: session already running ──────────────────
                    await ws.send(json.dumps({
                        "type":        "livenet_join",
                        "sessionId":   session_id,
                        "capabilities": caps,
                    }))
                    raw = await ws.recv()
                    msg = json.loads(raw)
                    node_id          = msg.get("nodeId")
                    assigned_modules = msg.get("assignedModules", [])
                    assigned_net     = msg.get("networkId", network_id or "")
                    print(f"[node] joined session {session_id} — modules={assigned_modules}")

                    await _execute_session(
                        ws, assigned_net, session_id,
                        node_id, assigned_modules, http_base
                    )
                    return  # single-shot; exit after one session

                else:
                    print("[node] ERROR: use --network --lobby or --session <id>")
                    return

        except (websockets.ConnectionClosed, ConnectionRefusedError, OSError) as e:
            if lobby:
                print(f"[node] connection lost ({e}), reconnecting in 5s…")
                await asyncio.sleep(5)
            else:
                print(f"[node] connection error: {e}")
                return


# ─── Entry point ──────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="AFS Living Network Python compute node (Spec 22 P1)")
    parser.add_argument("--server",  default="ws://localhost:3000",
                        help="WebSocket server URL (default: ws://localhost:3000)")
    parser.add_argument("--network", help="Network ID — use with --lobby")
    parser.add_argument("--session", help="Session ID — direct join, no lobby wait")
    parser.add_argument("--lobby",   action="store_true",
                        help="Wait in lobby; re-join after each session completes")
    args = parser.parse_args()

    if not args.network and not args.session:
        parser.error("Provide --network <id> --lobby  OR  --session <id>")

    try:
        asyncio.run(run_node(
            server_url = args.server,
            network_id = args.network,
            session_id = args.session,
            lobby      = args.lobby,
        ))
    except KeyboardInterrupt:
        print("\n[node] stopped")


if __name__ == "__main__":
    main()
