"""
AgentLink MCP Server — Supabase Backend
Signaling via Supabase PostgreSQL + polling.
Komunikasi P2P via WebSocket langsung antar agent.
"""

import asyncio
import json
import uuid
import os
import secrets
import logging
from contextlib import asynccontextmanager
from typing import Optional, Dict, Any
from datetime import datetime, timezone, timedelta

import aiohttp
import websockets
import websockets.server
import websockets.client
from websockets.exceptions import ConnectionClosed
from mcp.server.fastmcp import FastMCP, Context
from pydantic import BaseModel, Field, ConfigDict

logging.basicConfig(level=logging.WARNING)
logger = logging.getLogger(__name__)

# ─── Konstanta ────────────────────────────────────────────────────────────────
SUPABASE_URL   = os.environ.get("SUPABASE_URL",  "https://uqsmqsgiyjuilyfmgdku.supabase.co")
SUPABASE_KEY   = os.environ.get("SUPABASE_KEY",  "sb_publishable_XOlFEUvggwp4x694NH1Lxg_Obb-Zv2F")
WS_PORT        = int(os.environ.get("WS_PORT",   "8765"))
AGENT_NAME     = os.environ.get("AGENT_NAME",    f"agent-{uuid.uuid4().hex[:6]}")
POLL_INTERVAL  = 10   # detik
AGENT_TTL      = 30   # detik — agent dianggap offline

# ─── State Global ─────────────────────────────────────────────────────────────
agent_id: Optional[str]         = None
agent_name: str                  = AGENT_NAME
agent_ws_url: str                = ""
current_room: Optional[dict]     = None
p2p_connections: Dict[str, Any]  = {}
known_peers: Dict[str, dict]     = {}
inbox: list                      = []
http_session: Optional[aiohttp.ClientSession] = None

# ─── Supabase REST ────────────────────────────────────────────────────────────

def _hdrs():
    return {
        "apikey":        SUPABASE_KEY,
        "Authorization": f"Bearer {SUPABASE_KEY}",
        "Content-Type":  "application/json",
        "Prefer":        "return=representation",
    }

async def sb_select(table: str, qs: str = "") -> list:
    url = f"{SUPABASE_URL}/rest/v1/{table}{'?' + qs if qs else ''}"
    async with http_session.get(url, headers=_hdrs()) as r:
        return await r.json(content_type=None)

async def sb_insert(table: str, data: dict) -> list:
    async with http_session.post(f"{SUPABASE_URL}/rest/v1/{table}", headers=_hdrs(), json=data) as r:
        return await r.json(content_type=None)

async def sb_upsert(table: str, data: dict, on_conflict: str) -> list:
    hdrs = {**_hdrs(), "Prefer": f"resolution=merge-duplicates,return=representation"}
    async with http_session.post(f"{SUPABASE_URL}/rest/v1/{table}",
        headers=hdrs, json=data, params={"on_conflict": on_conflict}) as r:
        return await r.json(content_type=None)

async def sb_update(table: str, qs: str, data: dict) -> list:
    async with http_session.patch(f"{SUPABASE_URL}/rest/v1/{table}?{qs}",
        headers=_hdrs(), json=data) as r:
        return await r.json(content_type=None)

async def sb_delete(table: str, qs: str) -> None:
    async with http_session.delete(f"{SUPABASE_URL}/rest/v1/{table}?{qs}",
        headers=_hdrs()) as r:
        await r.read()

async def cleanup_stale():
    cutoff = (datetime.now(timezone.utc) - timedelta(seconds=AGENT_TTL)).isoformat()
    await sb_delete("agents", f"last_ping=lt.{cutoff}")

# ─── P2P WebSocket ────────────────────────────────────────────────────────────

async def p2p_handler(websocket):
    peer_id = None
    try:
        raw  = await asyncio.wait_for(websocket.recv(), timeout=10)
        hs   = json.loads(raw)
        peer_id   = hs.get("agent_id", "unknown")
        peer_name = hs.get("name", peer_id)
        p2p_connections[peer_id] = websocket
        async for raw_msg in websocket:
            msg = json.loads(raw_msg)
            inbox.append({"type": "message", "from": peer_name, "agent_id": peer_id,
                "content": msg.get("content", ""), "msg_type": msg.get("msg_type", "text"),
                "timestamp": datetime.now(timezone.utc).isoformat()})
    except (ConnectionClosed, asyncio.TimeoutError):
        pass
    finally:
        if peer_id: p2p_connections.pop(peer_id, None)

async def _connect_p2p(peer_id: str, peer_name: str, peer_url: str):
    if not peer_url or peer_id in p2p_connections: return
    try:
        ws = await asyncio.wait_for(websockets.connect(peer_url), timeout=15)
        await ws.send(json.dumps({"agent_id": agent_id, "name": agent_name}))
        p2p_connections[peer_id] = ws
        asyncio.create_task(_p2p_listen(peer_id, peer_name, ws))
    except Exception as e:
        logger.warning(f"Gagal P2P ke {peer_name}: {e}")

async def _p2p_listen(peer_id: str, peer_name: str, ws):
    try:
        async for raw in ws:
            msg = json.loads(raw)
            inbox.append({"type": "message", "from": peer_name, "agent_id": peer_id,
                "content": msg.get("content", ""), "msg_type": msg.get("msg_type", "text"),
                "timestamp": datetime.now(timezone.utc).isoformat()})
    except ConnectionClosed:
        pass
    finally:
        p2p_connections.pop(peer_id, None)

# ─── Polling Loop ─────────────────────────────────────────────────────────────

async def polling_loop():
    global known_peers
    while True:
        await asyncio.sleep(POLL_INTERVAL)
        if not agent_id or not current_room: continue
        try:
            rid = current_room["room_id"]
            await sb_update("agents", f"agent_id=eq.{agent_id}&room_id=eq.{rid}",
                {"last_ping": datetime.now(timezone.utc).isoformat()})
            await cleanup_stale()
            fresh_list = await sb_select("agents", f"room_id=eq.{rid}&agent_id=neq.{agent_id}")
            fresh = {a["agent_id"]: a for a in fresh_list}
            for pid, peer in fresh.items():
                if pid not in known_peers:
                    inbox.append({"type": "event", "event": "agent_joined",
                        "from": peer["name"], "agent_id": pid,
                        "timestamp": datetime.now(timezone.utc).isoformat()})
                    asyncio.create_task(_connect_p2p(pid, peer["name"], peer.get("ws_url", "")))
            for pid in list(known_peers.keys()):
                if pid not in fresh:
                    inbox.append({"type": "event", "event": "agent_left",
                        "from": known_peers[pid]["name"], "agent_id": pid,
                        "timestamp": datetime.now(timezone.utc).isoformat()})
                    p2p_connections.pop(pid, None)
            known_peers = fresh
        except Exception as e:
            logger.warning(f"Polling error: {e}")

# ─── Lifespan ─────────────────────────────────────────────────────────────────

@asynccontextmanager
async def lifespan(server):
    global http_session
    http_session = aiohttp.ClientSession()
    p2p_server   = await websockets.server.serve(p2p_handler, "0.0.0.0", WS_PORT)
    poll_task    = asyncio.create_task(polling_loop())
    yield {}
    poll_task.cancel()
    p2p_server.close()
    await p2p_server.wait_closed()
    await http_session.close()

mcp = FastMCP("agentlink", lifespan=lifespan)

# ─── Input Models ─────────────────────────────────────────────────────────────

class RegisterInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    name:          Optional[str] = Field(default=None, description="Nama agent")
    ws_public_url: str           = Field(..., description="URL WebSocket publik. Contoh: wss://abc.trycloudflare.com")

class CreateRoomInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    name:       str  = Field(..., description="Nama room", min_length=1, max_length=50)
    is_private: bool = Field(default=False, description="True = private (butuh token), False = public")

class JoinRoomInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    room_id: str           = Field(..., description="ID room (6 karakter, contoh: ABC123)")
    token:   Optional[str] = Field(default=None, description="Token untuk private room")

class SendInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    peer_id:  str = Field(..., description="Agent ID peer tujuan")
    message:  str = Field(..., description="Isi pesan", min_length=1, max_length=10000)
    msg_type: str = Field(default="text", description="Tipe: 'text', 'data', 'command'")

class BroadcastInput(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, extra="forbid")
    message:  str = Field(..., description="Pesan ke semua peer", min_length=1)
    msg_type: str = Field(default="text", description="Tipe pesan")

class ReadInboxInput(BaseModel):
    model_config = ConfigDict(extra="forbid")
    limit: int  = Field(default=10, ge=1, le=100, description="Jumlah pesan")
    clear: bool = Field(default=False, description="Hapus setelah dibaca")

# ─── Tools ────────────────────────────────────────────────────────────────────

@mcp.tool(name="agent_register", annotations={
    "readOnlyHint": False, "destructiveHint": False, "idempotentHint": True, "openWorldHint": True})
async def agent_register(params: RegisterInput) -> str:
    """
    Daftarkan agent ke AgentLink. Wajib dipanggil pertama.

    Dapatkan ws_public_url via Cloudflare Tunnel:
      cloudflared tunnel --url ws://localhost:8765

    Args:
        params: name (opsional), ws_public_url (wss://...)
    Returns:
        str: JSON berisi agent_id
    """
    global agent_id, agent_name, agent_ws_url
    if params.name: agent_name = params.name
    agent_ws_url = params.ws_public_url
    agent_id     = uuid.uuid4().hex[:8]
    return json.dumps({"success": True, "agent_id": agent_id, "name": agent_name,
        "ws_public_url": agent_ws_url,
        "message": f"Agent '{agent_name}' siap. Sekarang bisa create/join room."}, indent=2)


@mcp.tool(name="room_create", annotations={
    "readOnlyHint": False, "destructiveHint": False, "idempotentHint": False, "openWorldHint": True})
async def room_create(params: CreateRoomInput) -> str:
    """
    Buat room baru (public atau private).

    Public: siapa saja bisa join dengan room_id.
    Private: butuh room_id + token otomatis — bagikan ke peer.
    Room dihapus otomatis saat semua agent keluar.

    Args:
        params: name (nama room), is_private (True/False)
    Returns:
        str: JSON berisi room_id dan token (jika private)
    """
    global current_room, known_peers
    if not agent_id:
        return json.dumps({"success": False, "error": "Jalankan agent_register dulu."})
    try:
        room_id = uuid.uuid4().hex[:6].upper()
        token   = secrets.token_urlsafe(8) if params.is_private else ""
        await sb_insert("rooms", {"room_id": room_id, "name": params.name,
            "is_private": params.is_private, "token": token, "owner_id": agent_id})
        await sb_insert("agents", {"agent_id": agent_id, "name": agent_name,
            "room_id": room_id, "ws_url": agent_ws_url})
        current_room = {"room_id": room_id, "name": params.name, "is_private": params.is_private}
        known_peers  = {}
        result = {"success": True, "room_id": room_id, "name": params.name,
            "is_private": params.is_private, "message": f"Room '{params.name}' berhasil dibuat."}
        if params.is_private:
            result["token"]       = token
            result["instruction"] = "Bagikan room_id DAN token ke peer."
        else:
            result["instruction"] = f"Bagikan room_id '{room_id}' ke peer untuk join."
        return json.dumps(result, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool(name="room_join", annotations={
    "readOnlyHint": False, "destructiveHint": False, "idempotentHint": False, "openWorldHint": True})
async def room_join(params: JoinRoomInput) -> str:
    """
    Join room yang sudah ada. P2P otomatis terbentuk ke semua agent di room.

    Public: cukup room_id. Private: butuh room_id + token dari owner.

    Args:
        params: room_id (6 karakter), token (opsional untuk private)
    Returns:
        str: JSON info room + daftar agent
    """
    global current_room, known_peers
    if not agent_id:
        return json.dumps({"success": False, "error": "Jalankan agent_register dulu."})
    try:
        rid   = params.room_id.upper()
        rooms = await sb_select("rooms", f"room_id=eq.{rid}")
        if not rooms:
            return json.dumps({"success": False, "error": f"Room '{rid}' tidak ditemukan."})
        room = rooms[0]
        if room["is_private"] and room.get("token") != (params.token or ""):
            return json.dumps({"success": False, "error": "Token salah."})
        await cleanup_stale()
        await sb_upsert("agents", {"agent_id": agent_id, "name": agent_name,
            "room_id": rid, "ws_url": agent_ws_url,
            "last_ping": datetime.now(timezone.utc).isoformat()}, "agent_id,room_id")
        peers = await sb_select("agents", f"room_id=eq.{rid}&agent_id=neq.{agent_id}")
        current_room = {"room_id": rid, "name": room["name"], "is_private": room["is_private"]}
        known_peers  = {a["agent_id"]: a for a in peers}
        for p in peers:
            asyncio.create_task(_connect_p2p(p["agent_id"], p["name"], p.get("ws_url", "")))
        return json.dumps({"success": True, "room_id": rid, "name": room["name"],
            "is_private": room["is_private"], "agent_count": len(peers) + 1,
            "agents": [{"agent_id": a["agent_id"], "name": a["name"]} for a in peers],
            "message": f"Berhasil join room '{room['name']}'. P2P sedang terbentuk..."}, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool(name="room_leave", annotations={
    "readOnlyHint": False, "destructiveHint": False, "idempotentHint": True, "openWorldHint": True})
async def room_leave() -> str:
    """
    Keluar dari room saat ini. Room dihapus otomatis jika kosong.

    Returns:
        str: JSON status keluar
    """
    global current_room, known_peers
    if not agent_id or not current_room:
        return json.dumps({"success": False, "error": "Tidak sedang di dalam room."})
    try:
        rid = current_room["room_id"]
        await sb_delete("agents", f"agent_id=eq.{agent_id}&room_id=eq.{rid}")
        remaining = await sb_select("agents", f"room_id=eq.{rid}")
        if not remaining:
            await sb_delete("rooms", f"room_id=eq.{rid}")
    except Exception:
        pass
    for ws in list(p2p_connections.values()):
        try: await ws.close()
        except Exception: pass
    p2p_connections.clear()
    known_peers = {}
    current_room = None
    return json.dumps({"success": True, "message": "Berhasil keluar dari room."})


@mcp.tool(name="room_list", annotations={
    "readOnlyHint": True, "destructiveHint": False, "idempotentHint": True, "openWorldHint": True})
async def room_list() -> str:
    """
    Lihat daftar room public yang aktif saat ini.

    Returns:
        str: JSON daftar room
    """
    try:
        rooms  = await sb_select("rooms",  "is_private=eq.false")
        agents = await sb_select("agents", "")
        result = [{"room_id": r["room_id"], "name": r["name"],
            "agent_count": sum(1 for a in agents if a["room_id"] == r["room_id"])} for r in rooms]
        return json.dumps({"success": True, "rooms": result, "count": len(result)}, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool(name="room_info", annotations={
    "readOnlyHint": True, "destructiveHint": False, "idempotentHint": True, "openWorldHint": True})
async def room_info() -> str:
    """
    Info room saat ini: nama, ID, daftar agent, status private/public.

    Returns:
        str: JSON info room lengkap
    """
    if not current_room:
        return json.dumps({"success": False, "error": "Tidak sedang di dalam room."})
    try:
        agents = await sb_select("agents", f"room_id=eq.{current_room['room_id']}")
        return json.dumps({"success": True, "room": current_room,
            "agents": [{"agent_id": a["agent_id"], "name": a["name"]} for a in agents],
            "agent_count": len(agents), "p2p_connected": list(p2p_connections.keys())}, indent=2)
    except Exception as e:
        return json.dumps({"success": False, "error": str(e)})


@mcp.tool(name="agent_send", annotations={
    "readOnlyHint": False, "destructiveHint": False, "idempotentHint": False, "openWorldHint": True})
async def agent_send(params: SendInput) -> str:
    """
    Kirim pesan ke satu peer secara P2P langsung (tidak lewat Supabase).

    Args:
        params: peer_id, message, msg_type ('text'/'data'/'command')
    Returns:
        str: JSON status pengiriman
    """
    if params.peer_id not in p2p_connections:
        return json.dumps({"success": False, "error": f"Tidak ada koneksi P2P ke '{params.peer_id}'.",
            "connected_peers": list(p2p_connections.keys()),
            "hint": "Tunggu beberapa detik setelah join room agar P2P terbentuk."})
    try:
        await p2p_connections[params.peer_id].send(json.dumps({"from": agent_id,
            "from_name": agent_name, "content": params.message, "msg_type": params.msg_type,
            "timestamp": datetime.now(timezone.utc).isoformat()}))
        return json.dumps({"success": True, "message": f"Pesan terkirim ke '{params.peer_id}'"})
    except ConnectionClosed:
        p2p_connections.pop(params.peer_id, None)
        return json.dumps({"success": False, "error": "Koneksi P2P terputus."})


@mcp.tool(name="agent_broadcast", annotations={
    "readOnlyHint": False, "destructiveHint": False, "idempotentHint": False, "openWorldHint": True})
async def agent_broadcast(params: BroadcastInput) -> str:
    """
    Kirim pesan ke semua peer P2P di room saat ini.

    Args:
        params: message, msg_type
    Returns:
        str: JSON hasil broadcast per peer
    """
    if not p2p_connections:
        return json.dumps({"success": False, "error": "Tidak ada peer yang terhubung P2P."})
    payload = json.dumps({"from": agent_id, "from_name": agent_name,
        "content": params.message, "msg_type": params.msg_type,
        "timestamp": datetime.now(timezone.utc).isoformat()})
    results, failed = {}, []
    for pid, ws in list(p2p_connections.items()):
        try:
            await ws.send(payload)
            results[pid] = "sent"
        except Exception as e:
            results[pid] = f"failed: {e}"
            failed.append(pid)
    for pid in failed: p2p_connections.pop(pid, None)
    return json.dumps({"success": True, "results": results,
        "sent_to": len(results) - len(failed), "failed": len(failed)}, indent=2)


@mcp.tool(name="agent_read_inbox", annotations={
    "readOnlyHint": True, "destructiveHint": False, "idempotentHint": True, "openWorldHint": False})
async def agent_read_inbox(params: ReadInboxInput) -> str:
    """
    Baca pesan masuk dan event room (join/leave).

    Args:
        params: limit (default 10), clear (hapus setelah dibaca)
    Returns:
        str: JSON daftar pesan dan event
    """
    messages = inbox[-params.limit:]
    if params.clear: inbox.clear()
    return json.dumps({"messages": messages, "count": len(messages),
        "total_in_inbox": len(inbox), "cleared": params.clear}, indent=2)


@mcp.tool(name="agent_list", annotations={
    "readOnlyHint": True, "destructiveHint": False, "idempotentHint": True, "openWorldHint": False})
async def agent_list() -> str:
    """
    Lihat info agent ini dan daftar peer P2P yang aktif.

    Returns:
        str: JSON info agent + peer
    """
    return json.dumps({"my_agent_id": agent_id, "my_name": agent_name,
        "current_room": current_room.get("room_id") if current_room else None,
        "p2p_peers": list(p2p_connections.keys()),
        "known_peers": list(known_peers.keys()),
        "peer_count": len(p2p_connections)}, indent=2)


def main():
    """Entry point untuk uvx agentlink."""
    mcp.run(transport="stdio")

if __name__ == "__main__":
    main()
