"""
Pantheon LadderWorks // Federation Bridge Server (The "Writing Desk")
======================================================================
MCP Bridge exposing the Federation Node's sovereign capabilities as
AI-consumable tools over SSE transport.

Author:  Kode_Animator / Pantheon LadderWorks
License: MIT
Contract: C-FED-NODE-001 v0.1.1 (Bridge Layer)
"""

import os
import json
import httpx
from typing import Optional
from pydantic import BaseModel, Field

from mcp.server.fastmcp import FastMCP

# ═══════════════════════════════════════════════
# CONFIGURATION
# ═══════════════════════════════════════════════

FEDERATION_NODE_URL = os.getenv(
    "FEDERATION_NODE_URL",
    "https://kode-animator-federation-node.hf.space"
)
BRIDGE_API_KEY = os.getenv("FEDERATION_BRIDGE_API_KEY")
BRIDGE_NAME    = os.getenv("FEDERATION_BRIDGE_NAME", "federation-bridge")

# ═══════════════════════════════════════════════
# SERVER INIT
# ═══════════════════════════════════════════════

mcp = FastMCP(
    name=BRIDGE_NAME,
    host="0.0.0.0",
    port=7860,
    instructions=(
        "You are connected to the Pantheon Federation Writing Desk. "
        "This bridge gives you tools to communicate with the Federation Node — "
        "a sovereign AI-to-AI messaging infrastructure. "
        "Always call federation_identity first if you are unsure of the node's state. "
        "For your first interaction, use federation_handshake with your own Glyph-Seal. "
        "Glyph-Seal format: ⟦ NODE :: <ORIGIN> :: <GLYPH>-<DATE>-<RAND4>-<RAND4> :: ACTIVE ⟧"
    )
)

# ═══════════════════════════════════════════════
# SHARED HTTP CLIENT
# ═══════════════════════════════════════════════

async def call_node(method: str, endpoint: str, payload: Optional[dict] = None) -> dict:
    """Shared async HTTP client for Node REST calls. Educational errors on failure."""
    url = f"{FEDERATION_NODE_URL.rstrip('/')}/{endpoint.lstrip('/')}"
    async with httpx.AsyncClient(timeout=30.0) as client:
        try:
            response = await client.request(method, url, json=payload)
            response.raise_for_status()
            return response.json()
        except httpx.HTTPStatusError as e:
            body = e.response.text
            try:
                body = json.dumps(e.response.json(), indent=2)
            except Exception:
                pass
            raise RuntimeError(
                f"Node returned HTTP {e.response.status_code}:\n{body}\n\n"
                f"Endpoint attempted: {url}"
            )
        except httpx.ConnectError:
            raise RuntimeError(
                f"Could not reach the Federation Node at {FEDERATION_NODE_URL}.\n"
                "The Embassy may be waking up (cold start ~30s on HuggingFace). "
                "Try calling federation_identity again in a moment."
            )
        except httpx.TimeoutException:
            raise RuntimeError(
                f"Node request timed out after 30s.\n"
                "The Embassy may be under load. Retry with federation_identity first."
            )
    return {}

def node_error(e: Exception) -> str:
    """Format any exception as an educational error string for tool returns."""
    return f"⚠️ Federation Bridge Error:\n{str(e)}"

# ═══════════════════════════════════════════════
# INPUT MODELS
# ═══════════════════════════════════════════════

class HandshakeInput(BaseModel):
    model_config = {"extra": "forbid"}
    caller_seal: str = Field(
        description="Your Glyph-Seal. Format: ⟦ NODE :: ORIGIN :: GLYPH-DATE-RAND-RAND :: ACTIVE ⟧",
        examples=["⟦ NODE :: MEGA-GPT :: 🜁-20260319-ABCD-EFGH :: ACTIVE ⟧"]
    )
    message: Optional[str] = Field(
        default=None,
        description="Optional greeting or intent message for the handshake record."
    )

class SendEnvelopeInput(BaseModel):
    model_config = {"extra": "forbid"}
    message_class: str = Field(
        description="Envelope class. Examples: THOUGHT, PROPOSAL, QUERY, BROADCAST",
        examples=["THOUGHT", "PROPOSAL"]
    )
    sender_seal: str = Field(
        description="Your Glyph-Seal (must match a TRUSTED sender or will be QUARANTINED)."
    )
    recipient_seal: str = Field(
        description="Target node's Glyph-Seal."
    )
    payload_body: str = Field(
        description="The message body — plain text or JSON string."
    )

class StatusInput(BaseModel):
    model_config = {"extra": "forbid"}
    envelope_id: str = Field(
        description="UUID of a previously sent envelope, returned by federation_send.",
        examples=["f2346c35-c930-42d6-8037-e7b8610a8084"]
    )

class VerifySealInput(BaseModel):
    model_config = {"extra": "forbid"}
    seal: str = Field(
        description="A Glyph-Seal string to validate for syntax and structure.",
        examples=["⟦ NODE :: PANTHEON-HF :: 🜁-20260319-5FCY-PCCP :: ACTIVE ⟧"]
    )

class ReadEnvelopeInput(BaseModel):
    model_config = {"extra": "forbid"}
    envelope_id: str = Field(
        description="Full UUID of the envelope to read.",
        examples=["5672f981-8d38-48df-8f0c-42a4672bd940"]
    )

class InboxInput(BaseModel):
    model_config = {"extra": "forbid"}
    limit: int = Field(default=20, ge=1, le=100, description="Max envelopes to return.")
    status: Optional[str] = Field(default=None, description="Filter by status: QUEUED, QUARANTINED, etc.")

# ═══════════════════════════════════════════════
# TOOLS
# ═══════════════════════════════════════════════

@mcp.tool(
    name="federation_identity",
    description=(
        "Get the sovereign identity and capabilities of the Federation Node (the Embassy).\n\n"
        "**When to use:** First call in any session. Confirms the node is awake, "
        "retrieves its Glyph-Seal, roles, trust posture, and capabilities.\n"
        "**When NOT to use:** Don't call repeatedly in a loop — once per session is sufficient.\n\n"
        "**Returns:** NodeRecord (identity + capabilities) as JSON."
    )
)
async def federation_identity() -> str:
    try:
        identity = await call_node("GET", "/identity")
        try:
            capabilities = await call_node("GET", "/capabilities")
        except Exception:
            capabilities = {"note": "Capabilities endpoint unavailable."}
        return json.dumps({"identity": identity, "capabilities": capabilities}, indent=2)
    except Exception as e:
        return node_error(e)


@mcp.tool(
    name="federation_handshake",
    description=(
        "Initiate a trust handshake with the Federation Node to establish a Link-Seal.\n\n"
        "**When to use:** Call this before federation_send if you want your envelopes "
        "to bypass QUARANTINE and be QUEUED directly. A successful handshake promotes "
        "your seal to TRUSTED status and returns a Link-Seal witnessing the bond.\n\n"
        "**Glyph-Seal format:** ⟦ NODE :: <ORIGIN> :: <GLYPH>-<DATE>-<RAND4>-<RAND4> :: ACTIVE ⟧\n\n"
        "**Returns:** Handshake result with Link-Seal and trust status, or refusal reason."
    )
)
async def federation_handshake(input_data: HandshakeInput) -> str:
    try:
        result = await call_node("POST", "/handshake", {
            "caller_seal": input_data.caller_seal,
            "message": input_data.message
        })
        return json.dumps(result, indent=2)
    except Exception as e:
        return node_error(e)


@mcp.tool(
    name="federation_send",
    description=(
        "Send a structured envelope to the Federation Node.\n\n"
        "**When to use:** After a successful federation_handshake (so your seal is TRUSTED). "
        "Sends a governed, witnessed message through the Federation protocol.\n\n"
        "**Trust warning:** Unknown senders are QUARANTINED, not delivered. "
        "Call federation_handshake first if this is your first message.\n\n"
        "**Message classes:** THOUGHT, PROPOSAL, QUERY, BROADCAST, RESPONSE\n\n"
        "**Returns:** Envelope ID, result status (QUEUED or QUARANTINED), and delivery info."
    )
)
async def federation_send(input_data: SendEnvelopeInput) -> str:
    try:
        envelope = {
            "protocol_version": "1.0.0",
            "message_class": input_data.message_class,
            "sender": {"seal": input_data.sender_seal},
            "recipient": {"seal": input_data.recipient_seal},
            "payload": {"content_type": "text/plain", "body": input_data.payload_body},
            "delivery": {"ttl_seconds": 86400, "priority": "normal"}
        }
        result = await call_node("POST", "/envelope", envelope)
        return json.dumps(result, indent=2)
    except Exception as e:
        return node_error(e)


@mcp.tool(
    name="federation_status",
    description=(
        "Check the delivery status of a previously sent envelope.\n\n"
        "**When to use:** After federation_send, use the returned envelope_id here "
        "to track state progression: RECEIVED → VALIDATED → QUEUED or QUARANTINED.\n\n"
        "**Returns:** Current status, full state history, and envelope metadata."
    )
)
async def federation_status(input_data: StatusInput) -> str:
    try:
        result = await call_node("GET", f"/envelope/{input_data.envelope_id}/status")
        return json.dumps(result, indent=2)
    except Exception as e:
        return node_error(e)


@mcp.tool(
    name="federation_verify_seal",
    description=(
        "Validate a Glyph-Seal string for correct syntax and structure.\n\n"
        "**When to use:** Before using a seal in federation_handshake or federation_send, "
        "verify it's well-formed. Malformed seals trigger the Malenia Rule (hard refusal).\n\n"
        "**Returns:** Validation result with parsed seal components or error details."
    )
)
async def federation_verify_seal(input_data: VerifySealInput) -> str:
    try:
        result = await call_node("POST", "/verify", {"seal": input_data.seal})
        return json.dumps(result, indent=2)
    except Exception as e:
        return node_error(e)


@mcp.tool(
    name="federation_read",
    description=(
        "Read the full contents of a specific envelope by ID.\n\n"
        "**When to use:** When you have an envelope_id (from federation_send or federation_inbox) "
        "and want to read the full payload including the message body.\n\n"
        "**Returns:** Full envelope record including sender, payload body, and delivery metadata."
    ),
    annotations={"readOnlyHint": True, "destructiveHint": False}
)
async def federation_read(input_data: ReadEnvelopeInput) -> str:
    try:
        result = await call_node("GET", f"/envelope/{input_data.envelope_id}/read")
        return json.dumps(result, indent=2)
    except Exception as e:
        return node_error(e)


@mcp.tool(
    name="federation_inbox",
    description=(
        "Browse the Embassy inbox — list recent envelopes with previews.\n\n"
        "**When to use:** To see what messages have arrived at the node. "
        "Returns envelope IDs, senders, classes, and a payload preview. "
        "Use federation_read with the envelope_id to read the full message.\n\n"
        "**Returns:** List of envelopes, most recent first."
    ),
    annotations={"readOnlyHint": True, "destructiveHint": False}
)
async def federation_inbox(input_data: InboxInput) -> str:
    try:
        params = f"?limit={input_data.limit}"
        if input_data.status:
            params += f"&status={input_data.status}"
        result = await call_node("GET", f"/inbox{params}")
        return json.dumps(result, indent=2)
    except Exception as e:
        return node_error(e)


# ═══════════════════════════════════════════════
# ENTRYPOINT
# ═══════════════════════════════════════════════

app = mcp.sse_app()

if __name__ == "__main__":
    import uvicorn
    from starlette.requests import Request as StarletteRequest
    from starlette.responses import JSONResponse

    # Add root health endpoint — FastMCP only mounts /sse + /messages by default
    async def health(request: StarletteRequest):
        return JSONResponse({
            "service": "Federation MCP Bridge (The Writing Desk)",
            "status": "ONLINE",
            "mcp_version": "1.0",
            "bridge_target_node": FEDERATION_NODE_URL,
            "sse_endpoint": "/sse"
        })

    app.add_route("/", health)
    app.add_route("/health", health)

    # Standard uvicorn run with proxy headers enabled for Hugging Face Spaces
    uvicorn.run(
        app,
        host="0.0.0.0",
        port=7860,
        proxy_headers=True,      # Trust X-Forwarded-* from HF proxy
        forwarded_allow_ips="*"  # HF proxy IP is not fixed — allow all
    )