# /// script
# requires-python = ">=3.11"
# dependencies = ["mcp", "httpx"]
# ///
"""MCP server for the TAM Graph.

Exposes the TAM Graph (https://api.tamgraph.com) as MCP tools so that
Claude Desktop, Claude Code, or any other MCP client can search company data,
look up individual companies, and resolve email patterns. It is a thin,
self-contained wrapper over the HTTP API — no database, no state, one file.

Requirements
------------
Set your API key in the environment. Sign up at
https://api.tamgraph.com/login — the key is issued instantly, and adding a
card (never charged during the trial) activates 50 trial requests.

    TAMGRAPH_API_KEY   required, e.g. tg_live_...
    TAMGRAPH_BASE_URL  optional, defaults to https://api.tamgraph.com

Running it
----------
The script carries PEP 723 inline metadata, so `uv` installs its dependencies
into a throwaway environment automatically:

    TAMGRAPH_API_KEY=tg_live_... uv run mcp_server.py

It speaks MCP over stdio, so running it by hand just blocks waiting for a
client — normally you let your MCP client launch it.

Registering it with Claude Desktop / Claude Code
------------------------------------------------
Add this to your MCP config (Claude Desktop:
`~/Library/Application Support/Claude/claude_desktop_config.json` on macOS;
Claude Code: `.mcp.json` in your project, or `claude mcp add`):

    {
      "mcpServers": {
        "tamgraph": {
          "command": "uv",
          "args": ["run", "https://api.tamgraph.com/mcp_server.py"],
          "env": {
            "TAMGRAPH_API_KEY": "tg_live_..."
          }
        }
      }
    }

uv runs the script straight from that URL (it is this file, served by the
API). A local absolute path works too. Restart the client afterwards.
"""

import os

import httpx
from mcp.server.fastmcp import FastMCP

BASE_URL = os.environ.get("TAMGRAPH_BASE_URL", "https://api.tamgraph.com").rstrip("/")
API_KEY = os.environ.get("TAMGRAPH_API_KEY")
if not API_KEY:
    raise SystemExit(
        "TAMGRAPH_API_KEY is not set. Get a key at https://api.tamgraph.com/login "
        "and run: TAMGRAPH_API_KEY=tg_live_... uv run mcp_server.py"
    )

TIMEOUT = 30.0

mcp = FastMCP("tamgraph")


def _request(method: str, path: str, json_body: dict | None = None) -> dict:
    """Call the TAM Graph API and return parsed JSON, raising readable errors."""
    url = f"{BASE_URL}{path}"
    try:
        response = httpx.request(
            method,
            url,
            headers={"X-API-Key": API_KEY, "Content-Type": "application/json"},
            json=json_body,
            timeout=TIMEOUT,
            follow_redirects=True,
        )
    except httpx.RequestError as exc:
        raise RuntimeError(f"Could not reach the TAM Graph API at {url}: {exc}") from exc

    if response.is_success:
        return response.json()

    # The API returns {"detail": "..."} on errors; fall back to raw text.
    detail = response.text
    try:
        payload = response.json()
    except ValueError:
        payload = None
    if isinstance(payload, dict):
        detail = payload.get("detail") or payload.get("error") or detail
    raise RuntimeError(
        f"TAM Graph API error {response.status_code} on {method} {path}: {detail}"
    )


@mcp.tool()
def list_datasets() -> dict:
    """Catalog of datasets with sizes, allowed filters, and query hints. Call this first."""
    return _request("GET", "/v1/datasets")


@mcp.tool()
def search(
    dataset: str,
    filters: dict | None = None,
    limit: int = 25,
    cursor: str | None = None,
) -> dict:
    """Search one dataset for companies matching criteria.

    Call this when the user wants a LIST of companies matching criteria
    (industry, country, size, platform, revenue, funding, ...). For a single
    known company use the `company` tool instead.

    Datasets:
      - companies  36M global companies
      - local      2.3M US local businesses
      - places     7.2M Google Maps places
      - ecommerce  2.9M online stores with revenue estimates
      - startups   1.7M companies with funding data
      - saas       254k SaaS companies (b2b/b2c, headcount, ARR)
      - launches   197k Product Hunt launches
      - sponsors   25k newsletter/podcast/YouTube sponsors
      - agencies   378k marketing/dev/design agencies

    Filters support exact values ({"country": "US"}), lists for OR
    ({"country": ["US", "CA"]}), and `_gte` / `_lte` suffixes on numeric and
    date fields ({"monthly_sales_gte": 100000}). Unknown filter fields return a
    400 listing the allowed fields for that dataset — call `list_datasets`
    first if you are unsure what is filterable.

    Paginate by passing the response's `next_cursor` back as `cursor`; a null
    `next_cursor` means there are no more results.
    """
    body = {
        "filters": filters or {},
        "limit": limit,
        "cursor": cursor,
    }
    return _request("POST", f"/v1/search/{dataset}", json_body=body)


@mcp.tool()
def company(domain: str) -> dict:
    """Merged record for one company across all datasets, including emails and email pattern.

    Call this when the user asks about ONE specific company or domain
    (e.g. "what do you know about stripe.com?"). Returns firmographics, the
    dataset-specific extension blocks, known emails, and the email pattern when
    available. For finding many companies by criteria, use `search` instead.
    """
    return _request("GET", f"/v1/company/{domain}")


@mcp.tool()
def email_pattern(domain: str) -> dict:
    """Email naming pattern (e.g. first.last), generic inboxes, and MX records for a domain.

    Call this when constructing or guessing someone's email address at a
    company — it tells you the verified naming convention to apply to a
    person's name, plus known generic inboxes and mail provider.

    Everything this returns comes out of our database, and the email finder
    is included and unlimited under fair use — no credits, no per-address
    pricing.
    """
    return _request("GET", f"/v1/email-pattern/{domain}")


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