import asyncio
from uuid import uuid4
from datetime import datetime, timezone
from typing import Dict, Any, Optional, List

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from mcp.server.fastmcp import FastMCP, Context

# ---------------------------
# In-memory queue
# ---------------------------
commands: Dict[str, Dict[str, Any]] = {}
queue: List[str] = []
queue_lock = asyncio.Lock()

def now_iso() -> str:
    return datetime.now(timezone.utc).isoformat()

async def enqueue(command: str, args: Optional[Dict[str, Any]] = None, source: str = "mcp") -> Dict[str, Any]:
    cid = str(uuid4())
    rec = {
        "id": cid,
        "command": command,
        "args": args or {},
        "source": source,
        "status": "queued",
        "result": None,
        "error": None,
        "created_at": now_iso(),
        "updated_at": now_iso(),
        "claimed_by": None,
    }
    async with queue_lock:
        commands[cid] = rec
        queue.append(cid)
    return rec

# ---------------------------
# MCP server (3 tools only)
# ---------------------------
mcp = FastMCP(
    "Create3 Robot Bridge",
    instructions="Queues commands for a local Create 3 runner.",
    stateless_http=True,
)

@mcp.tool()
async def dock(ctx: Context) -> dict:
    rec = await enqueue("dock", {}, "mcp")
    return {"ok": True, "queued": rec}

@mcp.tool()
async def undock(ctx: Context) -> dict:
    rec = await enqueue("undock", {}, "mcp")
    return {"ok": True, "queued": rec}

@mcp.tool()
async def move_cm(cm: float, ctx: Context) -> dict:
    rec = await enqueue("move_cm", {"cm": cm}, "mcp")
    return {"ok": True, "queued": rec}

# ---------------------------
# FastAPI app
# ---------------------------
app = FastAPI(title="Create3 MCP Bridge")

class StatusUpdate(BaseModel):
    status: str
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None
    robot_id: Optional[str] = None

@app.get("/")
async def root():
    return {"ok": True, "mcp": "/mcp", "health": "/health", "time": now_iso()}

@app.get("/health")
async def health():
    return {
        "ok": True,
        "queued": sum(1 for c in commands.values() if c["status"] == "queued"),
        "running": sum(1 for c in commands.values() if c["status"] == "running"),
        "total": len(commands),
        "time": now_iso(),
    }

@app.get("/commands")
async def list_commands():
    return {"ok": True, "items": sorted(commands.values(), key=lambda x: x["updated_at"], reverse=True)}

@app.post("/commands/next")
async def claim_next_command(robot_id: str):
    async with queue_lock:
        while queue:
            cid = queue.pop(0)
            cmd = commands.get(cid)
            if not cmd or cmd["status"] != "queued":
                continue
            cmd["status"] = "running"
            cmd["claimed_by"] = robot_id
            cmd["updated_at"] = now_iso()
            return {"ok": True, "item": cmd}
    return {"ok": True, "item": None}

@app.post("/commands/{command_id}/status")
async def update_status(command_id: str, update: StatusUpdate):
    cmd = commands.get(command_id)
    if not cmd:
        raise HTTPException(status_code=404, detail="Command not found")
    cmd["status"] = update.status
    cmd["result"] = update.result
    cmd["error"] = update.error
    if update.robot_id:
        cmd["claimed_by"] = update.robot_id
    cmd["updated_at"] = now_iso()
    return {"ok": True}

# Mount MCP at /mcp
app.mount("/mcp", mcp.streamable_http_app())