#!/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-30s)
  explore             — full keyword analysis: interest over time + related queries
                        + interest by region (browser-based, ~10-30s)

Data source: trendspyg v0.7.0 (MIT) — Google Trends scraping library.
RSS path is lightweight (HTTP only); Explore path requires Chrome/Selenium.

Run with:
  python3 google_trends_mcp.py
"""

from __future__ import annotations

import json
import sys
from typing import Any

import trendspyg


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

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


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

def _status() -> dict[str, Any]:
    """Self-check: verify trendspyg is importable and working."""
    v = trendspyg.__version__
    return {
        "status": "ok",
        "library": "trendspyg",
        "version": v,
    }


def _get_trending(
    geo: str = "US",
    max_trends: int = 10,
    include_articles: bool = False,
) -> dict[str, Any]:
    """Get current trending search topics for a country (RSS — fast, no browser).

    Args:
        geo: Two-letter country code (US, GB, IT, DE, JP, etc.).
        max_trends: Max trending topics to return (default 10, max 20).
        include_articles: Include news article headlines/URLs (default False).

    Returns:
        A dict with: geo, fetched_at, trending_trends[] (each with title,
        traffic_volume, articles[] when requested).
    """
    max_trends = min(max_trends, 20)
    if max_trends < 1:
        max_trends = 10

    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 {"status": "error", "message": str(e), "geo": geo}

    if not raw:
        return {"status": "ok", "geo": geo, "trends": []}

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

    return {"status": "ok", "geo": geo, "trends": trends}


def _get_interest_over_time(
    keyword: str,
    geo: str = "US",
    timeframe: str = "today 12-m",
    output_format: str = "dict",
) -> dict[str, Any]:
    """Get a keyword's search interest over time (0-100 scale).

    Browser-based — takes ~10-30s. Returns data points with date, value,
    and is_partial flag.

    Args:
        keyword: Search term to analyze.
        geo: Two-letter country code (US, GB, IT, etc.) or '' for worldwide.
        timeframe: Google Trends range ('today 12-m', 'today 5-y',
            'now 7-d', 'now 1-H', 'all', or 'YYYY-MM-DD YYYY-MM-DD').
        output_format: 'dict' (default) or 'json'.

    Returns:
        Dict with keyword, geo, timeframe, datapoints[].
    """
    if not keyword or not keyword.strip():
        return {"status": "error", "message": "keyword is required"}

    try:
        of = "dict" if output_format == "dict" else "json"
        data = trendspyg.download_google_trends_interest_over_time(
            keyword=keyword.strip(),
            geo=geo,
            timeframe=timeframe,
            output_format=of,
        )
    except Exception as e:
        return {"status": "error", "message": str(e), "keyword": keyword}

    if of == "json":
        data = json.loads(data) if isinstance(data, str) else data

    return {
        "status": "ok",
        "keyword": keyword,
        "geo": geo,
        "timeframe": timeframe,
        "datapoints": data,
    }


def _explore(
    keyword: str,
    geo: str = "US",
    timeframe: str = "today 12-m",
    include_related: bool = True,
    include_geo: bool = True,
) -> dict[str, Any]:
    """Full keyword analysis: interest over time + related queries + region map.

    Browser-based — takes ~10-30s. Returns the complete Explore data in
    one call.

    Args:
        keyword: Search term to analyze.
        geo: Two-letter country code (US, GB, IT, etc.) or '' for worldwide.
        timeframe: Google Trends range ('today 12-m', 'today 5-y', etc.).
        include_related: Include related queries (top + rising).
        include_geo: Include interest by region.

    Returns:
        Complete ExploreEnvelope as a JSON-safe dict.
    """
    if not keyword or not keyword.strip():
        return {"status": "error", "message": "keyword is required"}

    try:
        data = trendspyg.download_google_trends_explore(
            keyword=keyword.strip(),
            geo=geo,
            timeframe=timeframe,
            include_related=include_related,
            include_geo=include_geo,
        )
    except Exception as e:
        return {"status": "error", "message": str(e), "keyword": keyword}

    # data is already a dict (ExploreEnvelope)
    data["status"] = "ok"
    return data


# ── JSON-RPC helpers ───────────────────────────────────────────────────────────

TOOLS: dict[str, tuple[Any, dict[str, Any]]] = {
    "status": (
        _status,
        {
            "description": "Check if Google Trends MCP is operational",
            "inputSchema": {"type": "object", "properties": {}, "required": []},
        },
    ),
    "get_trending": (
        _get_trending,
        {
            "description": "Get current trending search topics for a country (RSS — fast, no browser)",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "geo": {
                        "type": "string",
                        "description": "Two-letter country code (US, GB, IT, DE, etc.)",
                        "default": "US",
                    },
                    "max_trends": {
                        "type": "number",
                        "description": "Max trending topics (1-20)",
                        "default": 10,
                    },
                    "include_articles": {
                        "type": "boolean",
                        "description": "Include news article headlines/URLs",
                        "default": False,
                    },
                },
                "required": [],
            },
        },
    ),
    "get_interest_over_time": (
        _get_interest_over_time,
        {
            "description": "Get a keyword's search interest history (0-100 scale, browser-based ~10-30s)",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "keyword": {
                        "type": "string",
                        "description": "Search term to analyze",
                    },
                    "geo": {
                        "type": "string",
                        "description": "Two-letter country code or '' for worldwide",
                        "default": "US",
                    },
                    "timeframe": {
                        "type": "string",
                        "description": "Google Trends range: 'today 12-m', 'today 5-y', 'now 7-d', 'all'",
                        "default": "today 12-m",
                    },
                    "output_format": {
                        "type": "string",
                        "description": "Output format: 'dict' or 'json'",
                        "default": "dict",
                    },
                },
                "required": ["keyword"],
            },
        },
    ),
    "explore": (
        _explore,
        {
            "description": "Full keyword analysis: interest over time + related queries + region map (browser-based ~10-30s)",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "keyword": {
                        "type": "string",
                        "description": "Search term to analyze",
                    },
                    "geo": {
                        "type": "string",
                        "description": "Two-letter country code or '' for worldwide",
                        "default": "US",
                    },
                    "timeframe": {
                        "type": "string",
                        "description": "Google Trends range: 'today 12-m', 'today 5-y', 'now 7-d', 'all'",
                        "default": "today 12-m",
                    },
                    "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"],
            },
        },
    ),
}


# ── JSON-RPC request handler ───────────────────────────────────────────────────

def handle_request(request: dict[str, Any]) -> dict[str, Any]:
    """Dispatch a JSON-RPC 2.0 request and return a response."""
    req_id = request.get("id", None)
    method = request.get("method", "")
    params = request.get("params", {})

    if method == "tools/list":
        tools_list = []
        for name, (_, schema) in TOOLS.items():
            t = {
                "name": name,
                "description": schema["description"],
                "inputSchema": schema["inputSchema"],
            }
            # Set title for friendly name in UI
            title_map = {
                "status": "Status Check",
                "get_trending": "Get Trending Topics",
                "get_interest_over_time": "Interest Over Time",
                "explore": "Explore Keyword",
            }
            if name in title_map:
                t["title"] = title_map[name]
            tools_list.append(t)
        return _ok(req_id, tools_list)

    elif method == "tools/call":
        tool_name = params.get("name", "")
        arguments = params.get("arguments", {})

        if tool_name not in TOOLS:
            return _error(req_id, -32601, f"Method not found: {tool_name}")

        try:
            fn, _ = TOOLS[tool_name]
            result = fn(**arguments)
            # Ensure result is JSON-serializable and wrapped in content
            if isinstance(result, dict):
                content = [{"type": "text", "text": json.dumps(result, default=str)}]
            else:
                content = [{"type": "text", "text": str(result)}]
            return {"jsonrpc": "2.0", "id": req_id, "result": {"content": content}}
        except Exception as e:
            log(f"Error calling {tool_name}: {e}")
            return _error(req_id, -32603, f"Internal error: {e}")

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


def _ok(req_id: Any, result: Any) -> dict[str, Any]:
    return {"jsonrpc": "2.0", "id": req_id, "result": result}


def _error(req_id: Any, code: int, message: str) -> dict[str, Any]:
    return {
        "jsonrpc": "2.0",
        "id": req_id,
        "error": {"code": code, "message": message},
    }


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

def main() -> None:
    """Read JSON-RPC 2.0 requests from stdin and write responses to stdout."""
    log("Google Trends MCP server starting...")
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            request = json.loads(line)
        except json.JSONDecodeError as e:
            log(f"Invalid JSON: {e}")
            resp = _error(None, -32700, f"Parse error: {e}")
        else:
            resp = handle_request(request)

        sys.stdout.write(json.dumps(resp) + "\n")
        sys.stdout.flush()

    log("Google Trends MCP server shutting down.")


if __name__ == "__main__":
    main()