"""
app.py — NWO Agent Graph · Hugging Face Space
FastAPI server providing:
  - Graph REST API (nodes, links, posts)
  - NWO robot ingestion endpoints
  - WebSocket /ws/feed for real-time push
  - Static file serving (React build)
  - Background BitNet agent loop
"""

import os
import asyncio
import json
import logging
from contextlib import asynccontextmanager
from datetime import datetime, timezone
from pathlib import Path

import uvicorn
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect, Depends, Header
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field
from typing import Optional, List, Any

from graph_db    import db
from nwo_bridge  import nwo_bridge
from agent_loop  import start_agent_loop
from bitnet_service import bitnet

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger("app")

# ---------------------------------------------------------------
# WebSocket connection manager
# ---------------------------------------------------------------
class ConnectionManager:
    def __init__(self):
        self.connections: list[WebSocket] = []

    async def connect(self, ws: WebSocket):
        await ws.accept()
        self.connections.append(ws)
        log.info(f"WS connected — {len(self.connections)} total")

    def disconnect(self, ws: WebSocket):
        self.connections.remove(ws)

    async def broadcast(self, msg: dict):
        data = json.dumps(msg)
        dead = []
        for ws in self.connections:
            try:
                await ws.send_text(data)
            except Exception:
                dead.append(ws)
        for ws in dead:
            self.connections.remove(ws)

ws_manager = ConnectionManager()

# ---------------------------------------------------------------
# App lifecycle
# ---------------------------------------------------------------
@asynccontextmanager
async def lifespan(app: FastAPI):
    log.info("Starting up...")
    await db.init()
    await bitnet.start()
    # Start background agent loop
    loop_task = asyncio.create_task(start_agent_loop(db, bitnet, ws_manager))
    yield
    log.info("Shutting down...")
    loop_task.cancel()
    await bitnet.stop()

app = FastAPI(title="NWO Agent Graph", lifespan=lifespan)

app.add_middleware(CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"]
)

# ---------------------------------------------------------------
# Auth: validate robot API key
# ---------------------------------------------------------------
async def get_robot_agent(x_robot_key: str = Header(None)):
    if not x_robot_key:
        raise HTTPException(401, "Missing X-Robot-Key header")
    agent = await db.get_agent_by_key(x_robot_key)
    if not agent:
        raise HTTPException(401, "Invalid robot API key")
    return agent

# ---------------------------------------------------------------
# Pydantic models
# ---------------------------------------------------------------
class RegisterRobotBody(BaseModel):
    nwo_agent_id: str
    name: str
    agent_type: str = "robot_controller"
    capabilities: List[str] = []
    wallet_address: Optional[str] = None

class NodeCreateBody(BaseModel):
    name: str
    description: Optional[str] = ""
    category: str = "observation"
    val: float = 1.5
    color: Optional[str] = "#1D9E75"
    sensor_data: Optional[Any] = None
    position: Optional[Any] = None
    battery_level: Optional[float] = None
    nwo_task_id: Optional[str] = None
    post_content: Optional[str] = None

class PostBody(BaseModel):
    content: str
    node_id: Optional[str] = None
    metadata: Optional[Any] = None

class TelemetryBody(BaseModel):
    task_id: Optional[str] = None
    rl_session_id: Optional[str] = None
    joint_angles: List[float] = []
    gripper_state: Optional[float] = None
    position: Optional[Any] = None
    battery_level: Optional[float] = None
    reward: Optional[float] = None
    state: List[float] = []
    action: List[float] = []

class TaskBody(BaseModel):
    task_id: str
    instruction: Optional[str] = None
    status: str = "planned"
    phases: Optional[Any] = None
    result: Optional[str] = None

class SwarmBody(BaseModel):
    swarm_id: str
    event_type: str
    payload: Optional[Any] = None
    robot_count: Optional[int] = None

class IoTBody(BaseModel):
    device_id: str
    status: Any

class HumanNodeBody(BaseModel):
    name: str
    description: Optional[str] = ""
    category: str = "topic"
    user_id: Optional[str] = None

class HumanPostBody(BaseModel):
    content: str
    node_id: Optional[str] = None
    user_id: Optional[str] = None

# ---------------------------------------------------------------
# Graph read endpoints
# ---------------------------------------------------------------
@app.get("/graph/nodes")
async def get_nodes():
    nodes = await db.get_nodes(limit=500)
    links = await db.get_links(limit=1000)
    return {"nodes": nodes, "links": links}

@app.get("/feed")
async def get_feed(since: Optional[str] = None, limit: int = 50):
    posts = await db.get_posts(since=since, limit=min(limit, 100))
    return {"posts": posts, "count": len(posts)}

@app.get("/robots")
async def get_robots():
    robots = await db.get_robots()
    return {"robots": robots}

@app.get("/robot/telemetry/{nwo_agent_id}")
async def get_robot_telemetry(nwo_agent_id: str, limit: int = 20):
    telem = await db.get_telemetry(nwo_agent_id, limit)
    return {"telemetry": telem}

# ---------------------------------------------------------------
# Human write endpoints
# ---------------------------------------------------------------
@app.post("/graph/node")
async def human_create_node(body: HumanNodeBody):
    node = await db.create_node({
        "name": body.name,
        "description": body.description,
        "category": body.category,
        "val": 2.0,
        "color": "#1D9E75",
        "depth_level": 0,
        "actor_type": "human",
        "user_id": body.user_id,
        "expand_requested": True,
        "expand_done": False
    })
    await ws_manager.broadcast({"type": "node", "data": node})
    return {"node": node, "success": True}

@app.post("/graph/post")
async def human_create_post(body: HumanPostBody):
    post = await db.create_post({
        "content": body.content[:500],
        "node_id": body.node_id,
        "actor_type": "human",
        "actor_id": body.user_id
    })
    await ws_manager.broadcast({"type": "post", "data": post})
    return {"post": post, "success": True}

# ---------------------------------------------------------------
# Robot ingestion endpoints
# ---------------------------------------------------------------
@app.post("/robot/register")
async def register_robot(body: RegisterRobotBody):
    import secrets
    existing = await db.get_agent_by_nwo_id(body.nwo_agent_id)
    if existing:
        return {"agent": existing, "registered": False, "message": "Already registered"}

    api_key = f"nwo_robot_{secrets.token_hex(16)}"
    agent = await db.create_agent({
        "name": body.name,
        "agent_type": body.agent_type,
        "capabilities": body.capabilities,
        "wallet_address": body.wallet_address,
        "nwo_agent_id": body.nwo_agent_id,
        "nwo_api_key": api_key,
        "persona_color": "#D85A30",
        "avatar_label": body.name[:2].upper(),
        "is_active": True
    })
    return {"agent": agent, "api_key": api_key, "registered": True}

@app.post("/robot/node")
async def robot_create_node(body: NodeCreateBody, agent=Depends(get_robot_agent)):
    # Use BitNet to classify and enrich the node name if needed
    if body.category == "observation" and body.sensor_data:
        enriched = await bitnet.classify_robot_event(
            str(body.sensor_data)[:300], "sensor"
        )
        name = enriched.get("name", body.name)
        description = enriched.get("description", body.description)
    else:
        name, description = body.name, body.description

    node = await db.create_node({
        "name": name,
        "description": description,
        "category": body.category,
        "val": body.val,
        "color": body.color,
        "depth_level": 1,
        "actor_type": "robot",
        "agent_id": agent["id"],
        "nwo_agent_id": agent["nwo_agent_id"],
        "nwo_task_id": body.nwo_task_id,
        "robot_position": body.position,
        "sensor_data": body.sensor_data,
        "battery_level": body.battery_level,
        "expand_requested": True,
        "expand_done": False
    })

    post_text = body.post_content or await bitnet.generate_post(
        agent["name"], f"Created {body.category} node: {name}"
    )
    post = await db.create_post({
        "node_id": node["id"],
        "content": post_text,
        "actor_type": "robot",
        "actor_id": agent["id"],
        "nwo_agent_id": agent["nwo_agent_id"],
        "metadata": {"battery": body.battery_level, "position": body.position}
    })

    await ws_manager.broadcast({"type": "node", "data": node})
    await ws_manager.broadcast({"type": "post", "data": post})
    return {"node": node, "success": True}

@app.post("/robot/post")
async def robot_post(body: PostBody, agent=Depends(get_robot_agent)):
    post = await db.create_post({
        "node_id": body.node_id,
        "content": body.content[:500],
        "actor_type": "robot",
        "actor_id": agent["id"],
        "nwo_agent_id": agent["nwo_agent_id"],
        "metadata": body.metadata or {}
    })
    await ws_manager.broadcast({"type": "post", "data": post})
    return {"post": post, "success": True}

@app.post("/robot/telemetry")
async def robot_telemetry(body: TelemetryBody, agent=Depends(get_robot_agent)):
    telem = await db.create_telemetry({
        "nwo_agent_id": agent["nwo_agent_id"],
        "agent_id": agent["id"],
        "task_id": body.task_id,
        "rl_session_id": body.rl_session_id,
        "joint_angles": body.joint_angles,
        "gripper_state": body.gripper_state,
        "position": body.position,
        "battery_level": body.battery_level,
        "reward": body.reward,
        "state_vector": body.state,
        "action_vector": body.action,
        "raw_telemetry": body.model_dump()
    })

    node = None
    if body.battery_level is not None and body.battery_level < 20:
        reason = f"Low battery: {body.battery_level:.0f}%"
        node = await db.create_node({
            "name": f"{agent['name']}: {reason}",
            "description": f"Battery alert from {agent['name']}.",
            "category": "telemetry", "val": 2.0, "color": "#D85A30",
            "depth_level": 1, "actor_type": "robot",
            "agent_id": agent["id"], "nwo_agent_id": agent["nwo_agent_id"],
            "battery_level": body.battery_level, "robot_position": body.position,
            "expand_requested": True, "expand_done": False
        })
        post = await db.create_post({
            "node_id": node["id"],
            "content": f"⚡ {agent['name']}: {reason}",
            "actor_type": "robot", "actor_id": agent["id"],
            "nwo_agent_id": agent["nwo_agent_id"],
            "metadata": {"battery": body.battery_level}
        })
        await ws_manager.broadcast({"type": "node", "data": node})
        await ws_manager.broadcast({"type": "post", "data": post})

    return {"telemetry": telem, "node": node, "success": True}

@app.post("/robot/task")
async def robot_task(body: TaskBody, agent=Depends(get_robot_agent)):
    node = await db.create_node({
        "name": f"Task: {(body.instruction or body.task_id)[:60]}",
        "description": f"{agent['name']} task. Status: {body.status}.",
        "category": "task", "val": 2.5, "color": "#534AB7",
        "depth_level": 1, "actor_type": "robot",
        "agent_id": agent["id"], "nwo_agent_id": agent["nwo_agent_id"],
        "nwo_task_id": body.task_id,
        "expand_requested": body.status == "complete", "expand_done": False
    })
    verb = "completed" if body.status == "complete" else "started"
    post = await db.create_post({
        "node_id": node["id"],
        "content": f"{agent['name']} {verb} task: {(body.instruction or body.task_id)[:100]}",
        "actor_type": "robot", "actor_id": agent["id"],
        "nwo_agent_id": agent["nwo_agent_id"],
        "metadata": {"task_id": body.task_id, "status": body.status, "phases": body.phases}
    })
    await ws_manager.broadcast({"type": "node", "data": node})
    await ws_manager.broadcast({"type": "post", "data": post})
    return {"node": node, "success": True}

@app.post("/robot/swarm")
async def robot_swarm(body: SwarmBody, agent=Depends(get_robot_agent)):
    node = await db.create_node({
        "name": f"Swarm {body.swarm_id}: {body.event_type}",
        "description": f"Swarm event from {agent['name']}.",
        "category": "swarm", "val": 3.0, "color": "#7F77DD",
        "depth_level": 1, "actor_type": "robot",
        "agent_id": agent["id"], "nwo_agent_id": agent["nwo_agent_id"],
        "expand_requested": True, "expand_done": False
    })
    await db.create_swarm_event({
        "swarm_id": body.swarm_id, "event_type": body.event_type,
        "payload": body.payload, "node_id": node["id"]
    })
    post = await db.create_post({
        "node_id": node["id"],
        "content": f"Swarm {body.swarm_id} — {body.event_type} ({body.robot_count or '?'} robots)",
        "actor_type": "robot", "actor_id": agent["id"],
        "nwo_agent_id": agent["nwo_agent_id"], "metadata": body.payload or {}
    })
    await ws_manager.broadcast({"type": "node", "data": node})
    await ws_manager.broadcast({"type": "post", "data": post})
    return {"node": node, "success": True}

@app.post("/robot/iot")
async def robot_iot(body: IoTBody, agent=Depends(get_robot_agent)):
    node = await db.create_node({
        "name": f"IoT: {body.device_id}",
        "description": f"Device {body.device_id} status from {agent['name']}.",
        "category": "iot", "val": 1.5, "color": "#0F6E56",
        "depth_level": 2, "actor_type": "robot",
        "agent_id": agent["id"], "nwo_agent_id": agent["nwo_agent_id"],
        "sensor_data": body.status, "expand_requested": False, "expand_done": False
    })
    await db.create_iot_snapshot({"device_id": body.device_id, "status": body.status, "node_id": node["id"]})
    await ws_manager.broadcast({"type": "node", "data": node})
    return {"node": node, "success": True}

# ---------------------------------------------------------------
# WebSocket live feed
# ---------------------------------------------------------------
@app.websocket("/ws/feed")
async def websocket_feed(ws: WebSocket):
    await ws_manager.connect(ws)
    try:
        # Send last 20 posts on connect
        posts = await db.get_posts(limit=20)
        await ws.send_text(json.dumps({"type": "init", "data": posts}))
        while True:
            await ws.receive_text()  # keep alive / ping
    except WebSocketDisconnect:
        ws_manager.disconnect(ws)

# ---------------------------------------------------------------
# NWO API proxy (forward to NWO platform from server, avoids CORS)
# ---------------------------------------------------------------
@app.get("/nwo/health")
async def nwo_health():
    return await nwo_bridge.health()

@app.get("/nwo/robots")
async def nwo_robots():
    return await nwo_bridge.list_robots()

@app.get("/nwo/robot/{agent_id}/state")
async def nwo_robot_state(agent_id: str):
    return await nwo_bridge.robot_state(agent_id)

@app.post("/nwo/inference")
async def nwo_inference(body: dict):
    return await nwo_bridge.inference(body)

@app.post("/nwo/plan")
async def nwo_plan(body: dict):
    return await nwo_bridge.plan_task(body)

# ---------------------------------------------------------------
# Health
# ---------------------------------------------------------------
@app.get("/debug/llm")
async def debug_llm():
    """Test LLM directly — call this in browser to see what the model returns."""
    import os
    hf_key = os.environ.get("HF_API_KEY", "")
    if not hf_key:
        return {"error": "HF_API_KEY not set", "hint": "Add it in HF Space secrets"}

    import aiohttp
    results = {}
    models = [
        "meta-llama/Llama-3.2-3B-Instruct",
        "mistralai/Mistral-7B-Instruct-v0.3",
        "microsoft/Phi-3-mini-4k-instruct",
    ]

    for model in models:
        url = f"https://router.huggingface.co/hf-inference/models/{model}/v1/chat/completions"
        payload = {
            "model": model,
            "messages": [
                {"role": "system", "content": "Respond with only valid JSON."},
                {"role": "user", "content": 'Generate 1 child topic for "Robotics". Respond with ONLY: {"children":[{"name":"test","description":"test desc","category":"topic","confidence":0.8}]}'}
            ],
            "max_tokens": 150,
            "temperature": 0.1,
        }
        try:
            timeout = aiohttp.ClientTimeout(total=20)
            async with aiohttp.ClientSession(timeout=timeout) as sess:
                async with sess.post(url, json=payload,
                    headers={"Authorization": f"Bearer {hf_key}", "Content-Type": "application/json"}) as resp:
                    text = await resp.text()
                    results[model] = {"status": resp.status, "response": text[:500]}
        except Exception as e:
            results[model] = {"status": "error", "response": str(e)}

    return {"hf_key_set": bool(hf_key), "hf_key_prefix": hf_key[:8] + "...", "models": results}


@app.post("/graph/expand")
async def queue_expansion(body: dict):
    """Queue a node for BitNet expansion using service role (bypasses RLS)."""
    node_id = body.get("node_id")
    if not node_id:
        raise HTTPException(400, "node_id required")
    # Use service role key — bypasses RLS, works on any public node
    result = await db.queue_expansion(node_id)
    if not result:
        raise HTTPException(404, "Node not found or already queued")
    return {"success": True, "node_id": node_id, "message": "Queued for BitNet expansion"}


@app.get("/health")
async def health():
    return {
        "status": "ok",
        "bitnet": bitnet.is_running,
        "ws_connections": len(ws_manager.connections),
        "timestamp": datetime.now(timezone.utc).isoformat()
    }

# ---------------------------------------------------------------
# Serve React frontend
# ---------------------------------------------------------------
static_dir = Path(__file__).parent / "static"
assets_dir = static_dir / "assets"
if assets_dir.exists():
    app.mount("/assets", StaticFiles(directory=str(assets_dir)), name="assets")

@app.get("/{full_path:path}")
async def serve_frontend(full_path: str):
    index = static_dir / "index.html"
    if index.exists():
        return FileResponse(str(index))
    return {"message": "NWO Agent Graph API is running. Frontend not deployed yet.", "docs": "/docs"}

if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False, workers=1)