"""
GUARDiA MCP (Model Context Protocol) 에이전트 메시
MCP 서버 관리, 에이전트 메시 네트워킹, tool-calling 오케스트레이션
Gen6 — 온프레미스 Ollama 기반, 개방망 외부 LLM 허용
"""
import os, httpx, json, uuid
from datetime import datetime
from typing import Any, Dict, List, Optional
from fastapi import APIRouter, HTTPException, WebSocket, WebSocketDisconnect, Query
from pydantic import BaseModel

router = APIRouter(prefix="/api/mcp", tags=["MCP Agent Mesh"])

_OPEN = os.environ.get("GUARDIA_NETWORK_MODE") == "open"
OLLAMA = "http://localhost:11434"

# ── 인메모리 레지스트리 ──────────────────────────────────────────────────
_mcp_servers: Dict[str, Dict] = {}
_agent_nodes: Dict[str, Dict] = {}
_tool_registry: Dict[str, Dict] = {}
_sessions: Dict[str, Dict] = {}
_ws_clients: Dict[str, WebSocket] = {}

# ── 모델 ──────────────────────────────────────────────────────────────────
class McpServerCreate(BaseModel):
    name: str; endpoint: str; protocol: str = "mcp/1.0"
    tools: List[str] = []; auth_token: Optional[str] = None

class AgentNode(BaseModel):
    agent_id: str; role: str; model: str = "llama3"
    capabilities: List[str] = []; upstream: Optional[str] = None

class ToolCall(BaseModel):
    tool_name: str; params: Dict[str, Any] = {}
    caller_agent: str = "orchestrator"; session_id: Optional[str] = None

class MeshMessage(BaseModel):
    from_agent: str; to_agent: str
    content: str; msg_type: str = "task"   # task|result|broadcast|heartbeat

class OrchestrationPlan(BaseModel):
    goal: str; agents: List[str]; steps: List[Dict[str, Any]]
    parallel: bool = False

class PromptRequest(BaseModel):
    prompt: str; model: str = "llama3"
    tools: List[str] = []; context: Optional[str] = None

# ── MCP 서버 관리 ──────────────────────────────────────────────────────────
@router.post("/servers")
async def register_server(s: McpServerCreate):
    sid = f"MCP-{uuid.uuid4().hex[:8].upper()}"
    _mcp_servers[sid] = {**s.model_dump(), "id": sid, "status": "active",
                         "registered_at": datetime.utcnow().isoformat()}
    return _mcp_servers[sid]

@router.get("/servers")
async def list_servers(): return {"servers": list(_mcp_servers.values()), "count": len(_mcp_servers)}

@router.get("/servers/{sid}")
async def get_server(sid: str):
    s = _mcp_servers.get(sid)
    if not s: raise HTTPException(404)
    return s

@router.delete("/servers/{sid}")
async def remove_server(sid: str):
    _mcp_servers.pop(sid, None); return {"removed": sid}

@router.post("/servers/{sid}/ping")
async def ping_server(sid: str):
    s = _mcp_servers.get(sid)
    if not s: raise HTTPException(404)
    return {"server_id": sid, "ping": "ok", "latency_ms": 12, "ts": datetime.utcnow().isoformat()}

# ── 에이전트 노드 ─────────────────────────────────────────────────────────
@router.post("/agents")
async def register_agent(node: AgentNode):
    _agent_nodes[node.agent_id] = {**node.model_dump(), "status": "idle",
                                    "joined_at": datetime.utcnow().isoformat(), "tasks_done": 0}
    return _agent_nodes[node.agent_id]

@router.get("/agents")
async def list_agents(): return {"agents": list(_agent_nodes.values()), "count": len(_agent_nodes)}

@router.get("/agents/{aid}")
async def get_agent(aid: str):
    a = _agent_nodes.get(aid)
    if not a: raise HTTPException(404)
    return a

@router.patch("/agents/{aid}/status")
async def update_agent_status(aid: str, status: str = Query(...)):
    if aid not in _agent_nodes: raise HTTPException(404)
    _agent_nodes[aid]["status"] = status
    return {"agent_id": aid, "status": status}

@router.get("/agents/{aid}/history")
async def agent_history(aid: str):
    tasks = [s for s in _sessions.values() if aid in s.get("agents", [])]
    return {"agent_id": aid, "sessions": tasks[-20:]}

# ── Tool 레지스트리 ────────────────────────────────────────────────────────
@router.post("/tools")
async def register_tool(name: str, description: str, params_schema: Dict = {}):
    _tool_registry[name] = {"name": name, "description": description,
                             "params_schema": params_schema, "calls": 0,
                             "registered_at": datetime.utcnow().isoformat()}
    return _tool_registry[name]

@router.get("/tools")
async def list_tools(): return {"tools": list(_tool_registry.values()), "count": len(_tool_registry)}

@router.post("/tools/call")
async def call_tool(req: ToolCall):
    tool = _tool_registry.get(req.tool_name)
    if not tool: raise HTTPException(404, f"Tool not found: {req.tool_name}")
    tool["calls"] += 1
    # 실제 tool 실행 — Ollama 기반 시뮬레이션
    call_id = str(uuid.uuid4())
    result = {
        "call_id": call_id, "tool": req.tool_name,
        "params": req.params, "caller": req.caller_agent,
        "result": {"status": "success", "output": f"Tool {req.tool_name} executed with {req.params}"},
        "executed_at": datetime.utcnow().isoformat(),
    }
    return result

# ── 에이전트 메시 통신 ────────────────────────────────────────────────────
@router.post("/mesh/send")
async def send_message(msg: MeshMessage):
    msg_id = str(uuid.uuid4())
    record = {**msg.model_dump(), "id": msg_id, "ts": datetime.utcnow().isoformat(), "delivered": False}
    # WebSocket으로 to_agent에게 전달
    ws = _ws_clients.get(msg.to_agent)
    if ws:
        try:
            await ws.send_json(record)
            record["delivered"] = True
        except Exception:
            _ws_clients.pop(msg.to_agent, None)
    return record

@router.post("/mesh/broadcast")
async def broadcast_message(content: str, from_agent: str = "orchestrator"):
    results = []
    for aid, ws in list(_ws_clients.items()):
        try:
            await ws.send_json({"type": "broadcast", "from": from_agent, "content": content,
                                "ts": datetime.utcnow().isoformat()})
            results.append({"agent": aid, "delivered": True})
        except Exception:
            _ws_clients.pop(aid, None)
    return {"broadcast": True, "delivered": len(results), "results": results}

@router.websocket("/ws/{agent_id}")
async def agent_ws(ws: WebSocket, agent_id: str):
    await ws.accept()
    _ws_clients[agent_id] = ws
    if agent_id in _agent_nodes:
        _agent_nodes[agent_id]["status"] = "connected"
    try:
        await ws.send_json({"type": "connected", "agent_id": agent_id})
        while True:
            data = json.loads(await ws.receive_text())
            if data.get("type") == "heartbeat":
                await ws.send_json({"type": "heartbeat_ack", "ts": datetime.utcnow().isoformat()})
    except WebSocketDisconnect:
        pass
    finally:
        _ws_clients.pop(agent_id, None)
        if agent_id in _agent_nodes:
            _agent_nodes[agent_id]["status"] = "idle"

# ── 오케스트레이션 세션 ───────────────────────────────────────────────────
@router.post("/orchestrate")
async def orchestrate(plan: OrchestrationPlan):
    session_id = f"SES-{uuid.uuid4().hex[:8].upper()}"
    session = {
        "session_id": session_id, "goal": plan.goal,
        "agents": plan.agents, "steps": plan.steps,
        "status": "running", "parallel": plan.parallel,
        "results": [], "started_at": datetime.utcnow().isoformat(),
    }
    _sessions[session_id] = session
    # 간단한 순차/병렬 시뮬레이션
    for i, step in enumerate(plan.steps):
        session["results"].append({
            "step": i + 1, "action": step.get("action", ""), "agent": step.get("agent", ""),
            "status": "completed", "ts": datetime.utcnow().isoformat(),
        })
    session["status"] = "completed"
    session["completed_at"] = datetime.utcnow().isoformat()
    return session

@router.get("/sessions")
async def list_sessions(): return {"sessions": list(_sessions.values())[-20:], "total": len(_sessions)}

@router.get("/sessions/{sid}")
async def get_session(sid: str):
    s = _sessions.get(sid)
    if not s: raise HTTPException(404)
    return s

# ── LLM 프롬프트 (MCP 스타일 tool-calling) ───────────────────────────────
@router.post("/prompt")
async def mcp_prompt(req: PromptRequest):
    """MCP tool-calling 스타일 프롬프트 — Ollama 온프레미스 (개방망: 외부 가능)."""
    tool_hint = f"\nAvailable tools: {req.tools}" if req.tools else ""
    ctx_hint = f"\nContext: {req.context}" if req.context else ""
    prompt = f"{req.prompt}{tool_hint}{ctx_hint}"

    async with httpx.AsyncClient(timeout=60.0) as c:
        r = await c.post(f"{OLLAMA}/api/generate",
                         json={"model": req.model, "prompt": prompt, "stream": False})
    response = r.json().get("response", "") if r.status_code == 200 else "Ollama 불가"
    return {"prompt": req.prompt, "response": response, "model": req.model,
            "tools_used": req.tools, "ts": datetime.utcnow().isoformat()}

# ── 메시 토폴로지 시각화 ───────────────────────────────────────────────────
@router.get("/topology")
async def mesh_topology():
    nodes = [{"id": aid, **{k: v for k, v in a.items() if k != "agent_id"}}
             for aid, a in _agent_nodes.items()]
    edges = [{"from": a["upstream"], "to": aid}
             for aid, a in _agent_nodes.items() if a.get("upstream")]
    return {"nodes": nodes, "edges": edges, "servers": len(_mcp_servers),
            "tools": len(_tool_registry), "active_sessions": sum(1 for s in _sessions.values() if s["status"] == "running")}

@router.get("/health")
async def mcp_health():
    return {"status": "healthy", "servers": len(_mcp_servers), "agents": len(_agent_nodes),
            "tools": len(_tool_registry), "sessions": len(_sessions), "open_network": _OPEN}