"""DataServer — MCP server for UCC filing data (FastMCP + SQLite)."""

import sqlite3
import json
from pathlib import Path
from fastmcp import FastMCP
from fastmcp.resources import Resource

# ── WHAT: create FastMCP instance with server name and HTTP transport
# ── WHY: name becomes the server identifier visible in Claude's tool list
mcp = FastMCP("DataServer", port=8001)

DB_PATH = Path(__file__).parent / "ucc_filings.db"
SEED_SQL = Path(__file__).parent / "seed.sql"


def get_db() -> sqlite3.Connection:
    """Return a thread-safe SQLite connection with dict-style rows."""
    conn = sqlite3.connect(str(DB_PATH))
    conn.row_factory = sqlite3.Row
    return conn


def init_db() -> None:
    """Create schema and seed data if the database doesn't exist yet."""
    with get_db() as conn:
        conn.executescript(SEED_SQL.read_text())
    print(f"[DataServer] Database initialized at {DB_PATH}")


# ── WHAT: search_filings — full-text LIKE search on debtor_name
# ── WHY: the LLM calls this first; broad matching lets it find variants
# ── GOTCHA: always validate `limit` to prevent table scans with limit=99999
@mcp.tool()
def search_filings(debtor_name: str, limit: int = 10) -> list[dict]:
    """Search UCC filings by debtor name using full-text LIKE matching.

    Args:
        debtor_name: Company or entity name to search for (partial match OK)
        limit: Maximum number of results to return (default 10, max 50)

    Returns:
        List of filing records as dicts.
    """
    limit = min(max(1, limit), 50)  # clamp to [1, 50]
    with get_db() as conn:
        rows = conn.execute(
            "SELECT * FROM filings WHERE debtor_name LIKE ? LIMIT ?",
            (f"%{debtor_name}%", limit)
        ).fetchall()
    return [dict(row) for row in rows]


@mcp.tool()
def get_filing(filing_id: str) -> dict:
    """Fetch a single UCC filing by its unique filing ID.

    Args:
        filing_id: The filing primary key (e.g. 'FL-2024-001')

    Returns:
        Complete filing record as a dict.

    Raises:
        ValueError: If no filing with that ID exists.
    """
    with get_db() as conn:
        row = conn.execute(
            "SELECT * FROM filings WHERE filing_id = ?", (filing_id,)
        ).fetchone()
    if row is None:
        raise ValueError(f"No filing found with ID: {filing_id}")
    return dict(row)


@mcp.tool()
def list_secured_parties(debtor_name: str) -> list[str]:
    """Return all distinct secured party names for a debtor (any status).

    Args:
        debtor_name: Debtor name to look up (partial match OK)

    Returns:
        List of unique secured party name strings.
    """
    with get_db() as conn:
        rows = conn.execute(
            "SELECT DISTINCT secured_party FROM filings "
            "WHERE debtor_name LIKE ?",
            (f"%{debtor_name}%",)
        ).fetchall()
    return [row["secured_party"] for row in rows]


# ── WHAT: resource handler for filing://{filing_id} URI scheme
# ── WHY: resources let Claude attach structured data to context without
#         counting against the tool call budget
@mcp.resource("filing://{filing_id}")
def filing_resource(filing_id: str) -> Resource:
    """MCP resource: returns a single filing as a JSON resource."""
    data = get_filing(filing_id)  # reuse tool logic — raises on not found
    return Resource(
        uri=f"filing://{filing_id}",
        name=f"Filing {filing_id}",
        mime_type="application/json",
        content=json.dumps(data, indent=2)
    )


@mcp.resource("entity://{entity_name}/summary")
def entity_summary_resource(entity_name: str) -> Resource:
    """MCP resource: inline-computed entity risk summary."""
    with get_db() as conn:
        rows = conn.execute(
            "SELECT * FROM filings WHERE debtor_name LIKE ?",
            (f"%{entity_name}%",)
        ).fetchall()
    filings = [dict(r) for r in rows]
    summary = {
        "entity_name": entity_name,
        "total_filings": len(filings),
        "open_count": sum(1 for f in filings if f["status"] == "OPEN"),
        "terminated_count": sum(1 for f in filings if f["status"] == "TERMINATED"),
        "collateral_types": list({f["collateral_type"] for f in filings}),
        "secured_parties": list({f["secured_party"] for f in filings}),
    }
    return Resource(
        uri=f"entity://{entity_name}/summary",
        name=f"Entity Summary — {entity_name}",
        mime_type="application/json",
        content=json.dumps(summary, indent=2)
    )


if __name__ == "__main__":
    init_db()
    mcp.run(transport="http", host="0.0.0.0", port=8001)

"""ComputeServer — entity matching, risk scoring, deduplication via FastMCP."""

from fastmcp import FastMCP
from rapidfuzz import fuzz

mcp = FastMCP("ComputeServer", port=8002)


# ── WHAT: match_entity — finds best-matching names from a candidate list
# ── WHY: UCC filings use inconsistent names. "ACME Industrial LLC" and
#         "Acme Industrial" are the same entity but won't match exactly.
# ── GOTCHA: threshold is 0.0–1.0; WRatio returns 0–100, so multiply.
@mcp.tool()
def match_entity(
    query: str,
    candidates: list[str],
    threshold: float = 0.8
) -> list[dict]:
    """Fuzzy-match query against candidates using rapidfuzz WRatio.

    Args:
        query: The entity name to match (e.g. "ACME Industrial")
        candidates: List of candidate names to compare against
        threshold: Minimum similarity score 0.0–1.0 (default 0.8 = 80%)

    Returns:
        List of {name: str, score: float} sorted by score descending,
        filtered to results >= threshold.
    """
    min_score = threshold * 100
    results = []
    for candidate in candidates:
        score = fuzz.WRatio(query, candidate)
        if score >= min_score:
            results.append({"name": candidate, "score": round(score, 1)})
    return sorted(results, key=lambda x: x["score"], reverse=True)


# ── WHAT: score_risk — compute a numeric risk score for a debtor entity
# ── WHY: converts raw filing counts into an actionable 0–100 score with tier
# ── GOTCHA: cap at 100 — open liens * 15 can easily exceed 100 for serial borrowers
@mcp.tool()
def score_risk(
    debtor_name: str,
    filing_count: int,
    open_liens: int,
    collateral_types: list[str]
) -> dict:
    """Compute a risk score for a debtor based on filing data.

    Args:
        debtor_name: Entity name (used for factor labeling only)
        filing_count: Total number of filings (all statuses)
        open_liens: Number of currently OPEN liens
        collateral_types: List of collateral type strings

    Returns:
        {score: 0-100, tier: "low"|"medium"|"high"|"critical", factors: list[str]}
    """
    score = open_liens * 15
    factors = []

    # Collateral type bonuses
    collateral_set = {c.lower() for c in collateral_types}
    if "all assets" in collateral_set:
        score += 20
        factors.append("blanket lien on all assets (+20)")
    if "accounts receivable" in collateral_set:
        score += 10
        factors.append("accounts receivable pledged (+10)")
    if "equipment" in collateral_set:
        score += 5
        factors.append("equipment collateral (+5)")

    # Volume factor
    if filing_count >= 5:
        score += 10
        factors.append(f"high filing volume: {filing_count} total filings (+10)")

    score = min(100, score)
    factors.insert(0, f"{open_liens} open lien(s) base score: {open_liens * 15}")

    if score < 25:
        tier = "low"
    elif score < 50:
        tier = "medium"
    elif score < 75:
        tier = "high"
    else:
        tier = "critical"

    return {
        "debtor_name": debtor_name,
        "score": score,
        "tier": tier,
        "factors": factors,
        "filing_count": filing_count,
        "open_liens": open_liens,
    }


# ── WHAT: deduplicate_names — cluster similar name variants into canonical forms
# ── WHY: the same company might be filed as "ACME LLC", "ACME Industrial LLC",
#         "Acme Industrial" — dedup before storing in ChromaDB
@mcp.tool()
def deduplicate_names(names: list[str]) -> list[dict]:
    """Cluster similar entity names and return canonical name + aliases.

    Args:
        names: List of entity name strings to deduplicate

    Returns:
        List of {canonical: str, aliases: list[str], count: int}
    """
    if not names:
        return []

    clusters: list[dict] = []
    used = set()

    for i, name in enumerate(names):
        if i in used:
            continue
        cluster = {"canonical": name, "aliases": [], "count": 1}
        used.add(i)

        for j, other in enumerate(names):
            if j in used:
                continue
            if fuzz.WRatio(name, other) >= 85:
                cluster["aliases"].append(other)
                cluster["count"] += 1
                used.add(j)

        clusters.append(cluster)

    return sorted(clusters, key=lambda c: c["count"], reverse=True)


if __name__ == "__main__":
    mcp.run(transport="http", host="0.0.0.0", port=8002)

"""StorageServer — ChromaDB vector store for UCC entity snapshots."""

import time
import json
import re
from datetime import datetime, timezone
from fastmcp import FastMCP
import chromadb
from chromadb.utils import embedding_functions

mcp = FastMCP("StorageServer", port=8003)

# ── WHAT: initialize ChromaDB with local sentence-transformer embeddings
# ── WHY: all-MiniLM-L6-v2 is fast, small (80 MB), and runs CPU-only
# ── GOTCHA: EmbeddingFunction is called lazily — import errors surface at
#            first store/query, not at server startup. Test early.
_ef = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="sentence-transformers/all-MiniLM-L6-v2"
)
_client = chromadb.Client()  # in-memory; use PersistentClient for production
_collection = _client.get_or_create_collection(
    name="ucc_entities",
    embedding_function=_ef,
    metadata={"hnsw:space": "cosine"},
)


def _normalize_id(name: str) -> str:
    """Convert entity name to a safe ChromaDB document ID."""
    return re.sub(r"[^a-z0-9_-]", "-", name.lower().strip())


@mcp.tool()
def store_entity(
    entity_name: str,
    filing_ids: list[str],
    metadata: dict
) -> str:
    """Embed and upsert an entity snapshot into ChromaDB.

    Args:
        entity_name: Canonical entity name (used as embedding text + ID basis)
        filing_ids: List of associated filing IDs
        metadata: Arbitrary metadata dict (risk score, tier, timestamp, etc.)

    Returns:
        The document ID stored in ChromaDB.
    """
    doc_id = _normalize_id(entity_name)
    # Merge filing_ids into metadata as JSON string (ChromaDB metadata is flat)
    full_metadata = {
        **{k: str(v) for k, v in metadata.items()},
        "filing_ids": json.dumps(filing_ids),
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "entity_name": entity_name,
    }
    _collection.upsert(
        ids=[doc_id],
        documents=[entity_name],  # what gets embedded
        metadatas=[full_metadata],
    )
    return doc_id


@mcp.tool()
def find_similar(query: str, limit: int = 5) -> list[dict]:
    """Semantic search: find entities similar to the query.

    Args:
        query: Free-text query (entity name, description, or question)
        limit: Maximum results to return (default 5)

    Returns:
        List of {entity_name, score, filing_ids, metadata} sorted by similarity.
    """
    limit = min(max(1, limit), 20)
    results = _collection.query(
        query_texts=[query],
        n_results=min(limit, _collection.count() or 1),
        include=["documents", "distances", "metadatas"],
    )
    output = []
    for doc, dist, meta in zip(
        results["documents"][0],
        results["distances"][0],
        results["metadatas"][0],
    ):
        filing_ids = json.loads(meta.get("filing_ids", "[]"))
        output.append({
            "entity_name": meta.get("entity_name", doc),
            "score": round(1 - dist, 4),  # cosine: 1 = identical
            "filing_ids": filing_ids,
            "metadata": {k: v for k, v in meta.items()
                         if k not in ("filing_ids", "entity_name")},
        })
    return sorted(output, key=lambda x: x["score"], reverse=True)


@mcp.tool()
def get_entity_history(entity_name: str) -> list[dict]:
    """Retrieve all stored snapshots for an entity, newest first.

    Args:
        entity_name: Entity name to look up (exact or partial)

    Returns:
        List of snapshot dicts with filing_ids and metadata, sorted by timestamp.
    """
    # ChromaDB doesn't support WHERE on metadata natively in all versions;
    # fetch all and filter in Python for compatibility.
    total = _collection.count()
    if total == 0:
        return []
    all_results = _collection.get(
        include=["documents", "metadatas"],
        limit=total
    )
    matches = []
    for doc, meta in zip(all_results["documents"], all_results["metadatas"]):
        stored_name = meta.get("entity_name", doc)
        if entity_name.lower() in stored_name.lower():
            filing_ids = json.loads(meta.get("filing_ids", "[]"))
            matches.append({
                "entity_name": stored_name,
                "timestamp": meta.get("timestamp", ""),
                "filing_ids": filing_ids,
                "metadata": {k: v for k, v in meta.items()
                             if k not in ("filing_ids", "entity_name")},
            })
    return sorted(matches, key=lambda x: x["timestamp"], reverse=True)


if __name__ == "__main__":
    mcp.run(transport="http", host="0.0.0.0", port=8003)

"""MCPEcosystemHost — orchestrates DataServer, ComputeServer, StorageServer."""

import asyncio
import json
import sys
from typing import Any
import anthropic
from mcp import ClientSession
from mcp.client.http import http_client

# ── WHAT: server URL map — prefix → HTTP endpoint
# ── WHY: centralizing the map here makes it easy to swap ports or hostnames
#         without touching the ReAct loop logic
SERVER_MAP = {
    "data":    "http://localhost:8001/mcp",
    "compute": "http://localhost:8002/mcp",
    "storage": "http://localhost:8003/mcp",
}

SYSTEM_PROMPT = """You are a UCC lien analysis assistant with access to three MCP servers:
- data:: tools: search filings, fetch by ID, list secured parties
- compute:: tools: match entity names (fuzzy), score risk, deduplicate names
- storage:: tools: store entity snapshots, find similar entities, get entity history

Always follow this workflow for any lien inquiry:
1. Use data::search_filings to find all matching filings
2. Use compute::match_entity to identify the canonical entity name
3. Use compute::score_risk with the open lien count and collateral types
4. Use storage::store_entity to persist the result
5. Synthesize a clear, factual risk summary with filing IDs and risk tier

Never guess filing data — always fetch it. Never fabricate scores."""


class MCPEcosystemHost:
    """Connects to all three MCP servers and runs a multi-server ReAct loop."""

    def __init__(self):
        self.sessions: dict[str, ClientSession] = {}
        self._tool_map: dict[str, tuple[str, str]] = {}  # prefixed_name → (prefix, raw_name)
        self._tools_for_api: list[dict] = []
        self._client = anthropic.Anthropic()  # reads ANTHROPIC_API_KEY from env

    # ── WHAT: connect to all servers concurrently and merge tool lists
    async def connect(self) -> None:
        for prefix, url in SERVER_MAP.items():
            try:
                transport = await http_client(url).__aenter__()
                session = ClientSession(*transport)
                await session.__aenter__()
                await session.initialize()
                self.sessions[prefix] = session
                print(f"[Host] Connected to {prefix}:: at {url}")
            except Exception as exc:
                print(f"[Host] WARNING: could not connect to {prefix}:: — {exc}")

        await self._merge_tools()

    async def _merge_tools(self) -> None:
        """Fetch tool lists from all sessions; namespace and convert to Anthropic format."""
        for prefix, session in self.sessions.items():
            tools_result = await session.list_tools()
            for tool in tools_result.tools:
                prefixed = f"{prefix}::{tool.name}"
                self._tool_map[prefixed] = (prefix, tool.name)
                # ── WHAT: convert MCP tool schema to Anthropic tool dict
                # ── GOTCHA: MCP's inputSchema is already a JSON-Schema dict; wrap it
                self._tools_for_api.append({
                    "name": prefixed,
                    "description": tool.description or "",
                    "input_schema": (
                        tool.inputSchema if isinstance(tool.inputSchema, dict)
                        else tool.inputSchema.model_dump()
                    ),
                })
        print(f"[Host] Merged {len(self._tools_for_api)} tools from "
              f"{len(self.sessions)} servers: "
              f"{', '.join(self._tool_map.keys())}")

    async def _call_tool(self, prefixed_name: str, args: dict) -> Any:
        """Strip namespace prefix and dispatch the call to the correct server."""
        if prefixed_name not in self._tool_map:
            return {"error": f"Unknown tool: {prefixed_name}"}
        prefix, raw_name = self._tool_map[prefixed_name]
        session = self.sessions.get(prefix)
        if session is None:
            return {"error": f"Server '{prefix}' not connected"}
        try:
            result = await session.call_tool(raw_name, args)
            # MCP returns a list of content blocks; extract text from the first one
            if result.content:
                text = result.content[0].text if hasattr(result.content[0], "text") else str(result.content[0])
                try:
                    return json.loads(text)
                except json.JSONDecodeError:
                    return text
            return {}
        except Exception as exc:
            return {"error": str(exc)}

    # ── WHAT: ReAct loop — send messages to Claude, execute tool calls, repeat
    # ── WHY: "while stop_reason != end_turn" handles multi-step tool chains
    #         automatically without hardcoding how many steps the task takes
    async def run(self, query: str) -> str:
        messages = [{"role": "user", "content": query}]
        print(f"\n[Host] Query: {query}\n")

        while True:
            response = self._client.messages.create(
                model="claude-3-5-sonnet-20241022",
                max_tokens=4096,
                system=SYSTEM_PROMPT,
                tools=self._tools_for_api,
                messages=messages,
            )

            # Append assistant turn
            messages.append({"role": "assistant", "content": response.content})

            if response.stop_reason == "end_turn":
                # Extract final text
                final = next(
                    (b.text for b in response.content if hasattr(b, "text")), ""
                )
                print(f"\n[Host] Final answer:\n{final}\n")
                return final

            if response.stop_reason != "tool_use":
                break  # unexpected stop reason

            # Execute all tool calls in parallel
            tool_results = []
            for block in response.content:
                if block.type != "tool_use":
                    continue
                print(f"[Host] → {block.name}({json.dumps(block.input)[:80]}...)")
                result = await self._call_tool(block.name, block.input)
                print(f"[Host] ← result: {str(result)[:120]}")
                tool_results.append({
                    "type": "tool_result",
                    "tool_use_id": block.id,
                    "content": json.dumps(result),
                })

            messages.append({"role": "user", "content": tool_results})

        return "[Host] Unexpected loop exit"

    async def close(self) -> None:
        for session in self.sessions.values():
            try:
                await session.__aexit__(None, None, None)
            except Exception:
                pass


async def main() -> None:
    query = " ".join(sys.argv[1:]) if len(sys.argv) > 1 else (
        "Does ACME Industrial have any open liens? What is their risk score?"
    )
    host = MCPEcosystemHost()
    try:
        await host.connect()
        await host.run(query)
    finally:
        await host.close()


if __name__ == "__main__":
    asyncio.run(main())

"""Shared JWT auth utilities for the MCP ecosystem."""

import os
import time
import jwt  # PyJWT

SECRET = os.environ.get("MCP_SECRET_KEY", "dev-secret-change-in-production")
ALGORITHM = "HS256"


def generate_host_token(scopes: list[str], ttl_seconds: int = 300) -> str:
    """Generate a signed JWT for the host to authenticate with MCP servers.

    Args:
        scopes: List of OAuth scope strings the host is requesting
        ttl_seconds: Token lifetime in seconds (default 5 minutes)

    Returns:
        Signed JWT string to include as Authorization: Bearer 
    """
    payload = {
        "iss": "mcp-ecosystem-host",
        "sub": "host-agent",
        "scopes": scopes,
        "iat": int(time.time()),
        "exp": int(time.time()) + ttl_seconds,
    }
    return jwt.encode(payload, SECRET, algorithm=ALGORITHM)


def verify_token(token: str, required_scope: str) -> dict:
    """Verify a JWT and check that required_scope is present.

    Args:
        token: Raw JWT string (without 'Bearer ' prefix)
        required_scope: Scope string that must be in the token's scopes list

    Returns:
        Decoded payload dict if valid and authorized.

    Raises:
        PermissionError: If token invalid, expired, or missing required scope.
    """
    try:
        payload = jwt.decode(token, SECRET, algorithms=[ALGORITHM])
    except jwt.ExpiredSignatureError:
        raise PermissionError("Token expired")
    except jwt.InvalidTokenError as exc:
        raise PermissionError(f"Invalid token: {exc}")

    if required_scope not in payload.get("scopes", []):
        raise PermissionError(
            f"Missing scope: {required_scope}. "
            f"Token has: {payload.get('scopes', [])}"
        )
    return payload


# ── Auth middleware for FastMCP HTTP servers
# Add this to each server.py to enforce auth on all tool calls
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse

class JWTAuthMiddleware(BaseHTTPMiddleware):
    """Validates Bearer JWT on all incoming MCP HTTP requests."""

    def __init__(self, app, required_scope: str = "mcp:tools:execute:read_only"):
        super().__init__(app)
        self.required_scope = required_scope

    async def dispatch(self, request: Request, call_next):
        auth_header = request.headers.get("Authorization", "")
        if not auth_header.startswith("Bearer "):
            return JSONResponse(
                {"error": "Missing Authorization header"},
                status_code=401
            )
        token = auth_header.removeprefix("Bearer ")
        try:
            verify_token(token, self.required_scope)
        except PermissionError as exc:
            return JSONResponse({"error": str(exc)}, status_code=403)
        return await call_next(request)


# ── Usage in data_server/server.py (add after mcp = FastMCP(...)):
# from shared.auth import JWTAuthMiddleware
# mcp.app.add_middleware(JWTAuthMiddleware,
#                        required_scope="mcp:tools:execute:read_only")

# ── Usage in compute_server/server.py (score_risk needs :all scope):
# mcp.app.add_middleware(JWTAuthMiddleware,
#                        required_scope="mcp:tools:execute:all")