#!/usr/bin/env python3
"""AI Research Engine — Thin MCP Proxy. Admin key for backend auth."""

import json
import httpx
from mcp.server import FastMCP

BACKEND_URL = "http://10.30.20.250:8000"
ADMIN_KEY = "sk-admin-unlimited-2026"
client = httpx.Client(timeout=120.0)

def _get(path: str) -> dict:
    sep = "&" if "?" in path else "?"
    r = client.get(f"{BACKEND_URL}{path}{sep}api_key={ADMIN_KEY}")
    r.raise_for_status()
    return r.json()

mcp = FastMCP("ai-research-engine", instructions="AI Research Engine — private knowledge acquisition. 10 tools.")

@mcp.tool()
def search_web(query: str, category: str = "", limit: int = 10) -> str:
    """Full-text search across indexed documents."""
    r = _get(f"/api/search?q={query}&category={category}&limit={limit}")
    return json.dumps(r, indent=2)

@mcp.tool()
def semantic_search(query: str, limit: int = 10) -> str:
    """Search by meaning using vector embeddings."""
    r = _get(f"/api/semantic-search?q={query}&limit={limit}")
    return json.dumps(r, indent=2)

@mcp.tool()
def crawl_url(url: str, depth: int = 1) -> str:
    """Crawl a URL and index it."""
    r = _get(f"/api/crawl?url={url}&depth={depth}")
    return json.dumps(r, indent=2)

@mcp.tool()
def crawl_topic(topic: str, max_urls: int = 20) -> str:
    """Discover and crawl sources for a topic."""
    r = _get(f"/api/crawl-topic?topic={topic}&max_urls={max_urls}")
    return json.dumps(r, indent=2)

@mcp.tool()
def research_topic(topic: str) -> str:
    """Full pipeline: search → crawl → AI summary."""
    r = _get(f"/api/research?topic={topic}")
    return json.dumps(r, indent=2)

@mcp.tool()
def retrieve_document(url: str) -> str:
    """Get full indexed content of a document."""
    r = _get(f"/api/document?url={url}")
    return json.dumps(r, indent=2)

@mcp.tool()
def summarize_sources(urls: str, instruction: str = "Summarize key points") -> str:
    """AI summary of multiple URLs. urls: comma-separated."""
    r = _get(f"/api/summarize?urls={urls}&instruction={instruction}")
    return json.dumps(r, indent=2)

@mcp.tool()
def extract_information(url: str, schema: str = "company names, products, prices") -> str:
    """Extract structured data from a document using LLM."""
    r = _get(f"/api/extract?url={url}&schema={schema}")
    return json.dumps(r, indent=2)

@mcp.tool()
def create_report(topic: str, sources: str = "") -> str:
    """Generate comprehensive research report."""
    r = _get(f"/api/report?topic={topic}&sources={sources}")
    return json.dumps(r, indent=2)

@mcp.tool()
def index_status() -> str:
    """System health + business stats."""
    r = _get("/api/status")
    return json.dumps(r, indent=2)

if __name__ == "__main__":
    mcp.run(transport="stdio")