#!/usr/bin/env python3
"""Google Trends MCP server (JSON-RPC 2.0 over stdio) using trendspyg.

Capabilities (callable as `mcp__google-trends__<tool>`):
  status                 — self-check: confirms trendspyg is importable
  get_trending           — current trending topics per country (fast, RSS, no browser)
  get_interest_over_time — keyword interest history (browser-based, ~10-25s)
  explore                — full keyword analysis: interest over time + related
                           queries + interest by region (browser-based, ~10-25s)

Data source: trendspyg (MIT) — Google Trends scraping library.
The RSS path is lightweight (HTTP only); the Explore paths drive Chrome/Selenium.

Run with:
  python3 google_trends_mcp.py
"""

from __future__ import annotations

import json
import sys
from typing import Any

# ── Logging ─────────────────────────────────────────────────────────────────────

def log(msg: str) -> None:
    print(f"[google_trends_mcp] {msg}", file=sys.stderr, flush=True)


# ── trendspyg import ────────────────────────────────────────────────────────────
# Imported at module load, but a failure must not kill the server: the MCP
# handshake still has to succeed so skald can surface a readable tool error
# instead of a process that dies on startup.

try:
    import trendspyg
    _IMPORT_ERROR: str | None = None
except Exception as e:  # pragma: no cover - depends on the install
    trendspyg = None  # type: ignore[assignment]
    _IMPORT_ERROR = str(e)
    log(f"trendspyg is not importable: {e}")


# Google's Explore pages soft-throttle; trendspyg retries past it. Its defaults
# (10 attempts × 8s) allow a ~100s call, far beyond an agent's patience — cap
# the worst case at roughly 25s instead.
MAX_RETRIES = 3
RETRY_WAIT = 6.0


def _describe_error(e: Exception) -> str:
    """Turn a trendspyg exception into a message an agent can act on."""
    name = type(e).__name__
    if name == "RateLimitError":
        return f"Error: Google Trends is throttling this IP ({e}). Retry in a few minutes."
    if name == "BrowserError":
        return f"Error: Chrome could not be started ({e}). This tool needs Chrome/Selenium installed."
    if name == "InvalidParameterError":
        return f"Error: Invalid parameter: {e}"
    return f"Error: Google Trends request failed ({name}): {e}"


# ── Argument helpers ────────────────────────────────────────────────────────────

def _str_arg(args: dict[str, Any], key: str, default: str = "") -> str:
    value = args.get(key, default)
    return value.strip() if isinstance(value, str) else default


def _int_arg(args: dict[str, Any], key: str, default: int, lo: int, hi: int) -> int:
    try:
        value = int(args.get(key, default))
    except (TypeError, ValueError):
        return default
    return max(lo, min(hi, value))


def _bool_arg(args: dict[str, Any], key: str, default: bool) -> bool:
    value = args.get(key, default)
    return value if isinstance(value, bool) else default


# ── Tool implementations ────────────────────────────────────────────────────────

def _status(args: dict[str, Any]) -> str:
    """Self-check: confirm trendspyg is importable."""
    if trendspyg is None:
        return f"Error: trendspyg is not installed: {_IMPORT_ERROR}"
    version = getattr(trendspyg, "__version__", "unknown")
    return (
        f"OK: Google Trends integration is operational (trendspyg v{version}).\n"
        "  - get_trending: ready (RSS, no browser needed)\n"
        "  - get_interest_over_time / explore: require Chrome/Selenium"
    )


def _get_trending(args: dict[str, Any]) -> str:
    """Current trending search topics for a country (RSS — fast, no browser)."""
    if trendspyg is None:
        return f"Error: trendspyg is not installed: {_IMPORT_ERROR}"

    geo = _str_arg(args, "geo", "US") or "US"
    max_trends = _int_arg(args, "max_trends", 10, 1, 20)
    include_articles = _bool_arg(args, "include_articles", False)

    try:
        raw = trendspyg.download_google_trends_rss(
            geo=geo,
            output_format="dict",
            include_images=False,
            include_articles=include_articles,
            max_articles_per_trend=3,
        )
    except Exception as e:
        return _describe_error(e)

    trends = []
    for t in (raw or [])[:max_trends]:
        item = {
            "title": t.get("trend", ""),
            "traffic": t.get("traffic", ""),
            "traffic_min": t.get("traffic_min"),
            "published": t.get("published", ""),
            "explore_link": t.get("explore_link", ""),
        }
        if include_articles:
            item["articles"] = [
                {
                    "headline": a.get("headline", ""),
                    "url": a.get("url", ""),
                    "source": a.get("source", ""),
                }
                for a in (t.get("news_articles") or [])[:3]
            ]
        trends.append(item)

    return _json({"geo": geo, "count": len(trends), "trends": trends})


def _get_interest_over_time(args: dict[str, Any]) -> str:
    """A keyword's search interest over time (0-100 scale). Browser-based."""
    if trendspyg is None:
        return f"Error: trendspyg is not installed: {_IMPORT_ERROR}"

    keyword = _str_arg(args, "keyword")
    if not keyword:
        return "Error: Missing required parameter 'keyword'."

    geo = _str_arg(args, "geo", "US")
    timeframe = _str_arg(args, "timeframe", "today 12-m") or "today 12-m"

    try:
        points = trendspyg.download_google_trends_interest_over_time(
            keyword=keyword,
            geo=geo,
            timeframe=timeframe,
            output_format="dict",
            max_retries=MAX_RETRIES,
            retry_wait=RETRY_WAIT,
        )
    except Exception as e:
        return _describe_error(e)

    points = points or []
    return _json({
        "keyword": keyword,
        "geo": geo,
        "timeframe": timeframe,
        "count": len(points),
        "datapoints": points,
    })


def _explore(args: dict[str, Any]) -> str:
    """Full keyword analysis: interest over time + related queries + region map."""
    if trendspyg is None:
        return f"Error: trendspyg is not installed: {_IMPORT_ERROR}"

    keyword = _str_arg(args, "keyword")
    if not keyword:
        return "Error: Missing required parameter 'keyword'."

    try:
        envelope = trendspyg.download_google_trends_explore(
            keyword=keyword,
            geo=_str_arg(args, "geo", "US"),
            timeframe=_str_arg(args, "timeframe", "today 12-m") or "today 12-m",
            include_related=_bool_arg(args, "include_related", True),
            include_geo=_bool_arg(args, "include_geo", True),
            max_retries=MAX_RETRIES,
            retry_wait=RETRY_WAIT,
        )
    except Exception as e:
        return _describe_error(e)

    # trendspyg returns an ExploreEnvelope: already JSON-safe, every field present.
    return _json(envelope)


def _json(payload: Any) -> str:
    return json.dumps(payload, ensure_ascii=False, default=str)


# ── Tool manifest ────────────────────────────────────────────────────────────────

TOOLS = [
    {
        "name": "status",
        "title": "Status Check",
        "description": (
            "Self-check that the Google Trends integration is operational: confirms "
            "the trendspyg library is installed and reports its version. "
            "Call this first if any Google Trends tool fails."
        ),
        "inputSchema": {"type": "object", "properties": {}},
    },
    {
        "name": "get_trending",
        "title": "Get Trending Topics",
        "description": (
            "Get the search topics trending RIGHT NOW in a country, with their traffic "
            "volume and (optionally) the news articles behind them. Fast (~1s, no browser). "
            "Always current — it cannot look at past dates or filter by keyword; "
            "use get_interest_over_time for a specific term."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "geo": {
                    "type": "string",
                    "description": "Country or region code (e.g. 'US', 'GB', 'IT', 'US-CA'). Default: 'US'.",
                },
                "max_trends": {
                    "type": "integer",
                    "description": "Maximum number of trending topics to return (1-20). Default: 10.",
                },
                "include_articles": {
                    "type": "boolean",
                    "description": "Include up to 3 news article headlines/URLs per trend. Default: false.",
                },
            },
            "required": [],
        },
    },
    {
        "name": "get_interest_over_time",
        "title": "Interest Over Time",
        "description": (
            "Get a keyword's search interest history as Google's 0-100 relative index, "
            "oldest point first. Use it to see how a term's popularity moved over a period. "
            "Browser-based: takes ~10-25s and needs Chrome. A keyword Google has no data "
            "for comes back as a series of zeros."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "keyword": {
                    "type": "string",
                    "description": "Search term to analyze (e.g. 'bitcoin').",
                },
                "geo": {
                    "type": "string",
                    "description": "Country or region code (e.g. 'US', 'IT', 'US-CA'), or '' for worldwide. Default: 'US'.",
                },
                "timeframe": {
                    "type": "string",
                    "description": (
                        "Google Trends date range: 'today 12-m' (default), 'today 5-y', "
                        "'today 3-m', 'now 7-d', 'now 1-H', 'all', or a custom "
                        "'YYYY-MM-DD YYYY-MM-DD'."
                    ),
                },
            },
            "required": ["keyword"],
        },
    },
    {
        "name": "explore",
        "title": "Explore Keyword",
        "description": (
            "Full Google Trends Explore picture for one keyword in a single call: interest "
            "over time, related queries (top + rising), and interest by region. Prefer this "
            "over get_interest_over_time when you also want related terms or the regional "
            "breakdown. Browser-based: takes ~10-25s and needs Chrome."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "keyword": {
                    "type": "string",
                    "description": "Search term to analyze (e.g. 'bitcoin').",
                },
                "geo": {
                    "type": "string",
                    "description": "Country or region code (e.g. 'US', 'IT', 'US-CA'), or '' for worldwide. Default: 'US'.",
                },
                "timeframe": {
                    "type": "string",
                    "description": (
                        "Google Trends date range: 'today 12-m' (default), 'today 5-y', "
                        "'now 7-d', 'all', or a custom 'YYYY-MM-DD YYYY-MM-DD'."
                    ),
                },
                "include_related": {
                    "type": "boolean",
                    "description": "Include top + rising related queries. Default: true.",
                },
                "include_geo": {
                    "type": "boolean",
                    "description": "Include interest by region. Default: true.",
                },
            },
            "required": ["keyword"],
        },
    },
]

TOOL_DISPATCH = {
    "status":                 _status,
    "get_trending":           _get_trending,
    "get_interest_over_time": _get_interest_over_time,
    "explore":                _explore,
}


# ── JSON-RPC dispatch ────────────────────────────────────────────────────────────

def _ok(req_id: Any, result: Any) -> str:
    return json.dumps({"jsonrpc": "2.0", "id": req_id, "result": result})


def _text_result(req_id: Any, text: str, is_error: bool = False) -> str:
    payload: dict = {
        "jsonrpc": "2.0",
        "id": req_id,
        "result": {"content": [{"type": "text", "text": text}]},
    }
    if is_error:
        payload["result"]["isError"] = True
    return json.dumps(payload)


def _error(req_id: Any, code: int, message: str) -> str:
    return json.dumps({
        "jsonrpc": "2.0",
        "id": req_id,
        "error": {"code": code, "message": message},
    })


def handle_request(msg: dict) -> str | None:
    method = msg.get("method", "")
    req_id = msg.get("id")

    if method == "initialize":
        return _ok(req_id, {
            "protocolVersion": "2024-11-05",
            "capabilities": {"tools": {}},
            "serverInfo": {
                "name": "google-trends",
                "version": "1.1.0",
            },
        })

    if method == "notifications/initialized":
        return None

    if method == "ping":
        return _ok(req_id, {})

    if method == "tools/list":
        return _ok(req_id, {"tools": TOOLS})

    if method == "tools/call":
        params = msg.get("params", {})
        tool_name = params.get("name", "")
        tool_args = params.get("arguments") or {}

        handler = TOOL_DISPATCH.get(tool_name)
        if handler is None:
            return _text_result(req_id, f"Error: Unknown tool: {tool_name}", is_error=True)

        try:
            result = handler(tool_args)
            is_err = result.startswith("Error:")
            return _text_result(req_id, result, is_error=is_err)
        except Exception as e:
            log(f"Unhandled exception in tool '{tool_name}': {e}")
            return _text_result(req_id, f"Error: Internal error in tool '{tool_name}': {e}", is_error=True)

    # A notification we do not handle gets no reply — only requests do.
    if req_id is None:
        return None

    return _error(req_id, -32601, f"Method not found: {method}")


# ── Main loop ────────────────────────────────────────────────────────────────────

def main() -> None:
    log("Starting Google Trends MCP server")
    try:
        for line in sys.stdin:
            line = line.strip()
            if not line:
                continue
            try:
                msg = json.loads(line)
            except json.JSONDecodeError as e:
                log(f"Invalid JSON input: {e}")
                continue

            resp = handle_request(msg)
            if resp is not None:
                sys.stdout.write(resp + "\n")
                sys.stdout.flush()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()