#!/usr/bin/env python3
"""
MCP Search Server — provides web search via opencode infrastructure.

Two search methods (automatic fallback):
  1. npx opencode run (rate-limited but returns real web results)
  2. opencode-go API + model training data (no rate limit, less fresh)

Usage:
  python3 scripts/mcp_search_server.py          # MCP server (stdio)
  python3 scripts/mcp_search_server.py --query Q # one-shot search
  python3 scripts/mcp_search_server.py --url  U  # one-shot webfetch
"""
import json, os, subprocess, sys, time
from pathlib import Path
from typing import Any, Dict, List, Optional

PROJECT_ROOT = Path(__file__).resolve().parent.parent
CACHE_FILE = PROJECT_ROOT / "automation" / "data" / "mcp_search_cache.json"
SESSION_FILE = PROJECT_ROOT / "automation" / "data" / "mcp_session.txt"
CACHE_TTL = 3600
SESSION_TITLE = "opencode搜索"

API_BASE = "https://opencode.ai/zen/go/v1"
API_KEY = os.environ.get("OPENCODE_API_KEY", "")
if not API_KEY:
    try:
        from dotenv import load_dotenv
        env_path = PROJECT_ROOT / "platform" / "backend" / ".env"
        load_dotenv(env_path)
        API_KEY = os.environ.get("OPENCODE_API_KEY", "")
    except Exception:
        pass


# ── session (reuse same session for all MCP searches) ─────────────
def _load_session() -> Optional[str]:
    if SESSION_FILE.exists():
        try:
            return SESSION_FILE.read_text().strip() or None
        except Exception:
            pass
    return None

def _save_session_from_output(stdout: str):
    for line in stdout.strip().split("\n"):
        try:
            ev = json.loads(line)
            sid = ev.get("sessionID") or ev.get("part", {}).get("sessionID")
            if sid:
                SESSION_FILE.parent.mkdir(parents=True, exist_ok=True)
                SESSION_FILE.write_text(sid)
                return
        except Exception:
            pass


# ── cache ─────────────────────────────────────────────────────────
def _check_cache(query: str) -> Optional[List[Dict]]:
    if not CACHE_FILE.exists():
        return None
    try:
        data = json.loads(CACHE_FILE.read_text())
        entry = data.get(query)
        if entry and time.time() - entry.get("ts", 0) < CACHE_TTL:
            return entry.get("results")
    except Exception:
        pass
    return None

def _write_cache(query: str, results: List[Dict]):
    CACHE_FILE.parent.mkdir(parents=True, exist_ok=True)
    data = {}
    if CACHE_FILE.exists():
        try:
            data = json.loads(CACHE_FILE.read_text())
        except Exception:
            pass
    data[query] = {"ts": time.time(), "results": results}
    keys = sorted(data.keys(), key=lambda k: data[k].get("ts", 0), reverse=True)[:200]
    CACHE_FILE.write_text(json.dumps({k: data[k] for k in keys}, ensure_ascii=False))


# ── method 1: npx opencode run ────────────────────────────────────
def _search_via_opencode_cli(query: str, max_results: int) -> Optional[List[Dict]]:
    """Use npx opencode run to execute websearch tool (short timeout)."""
    sid = _load_session()
    args = ["npx", "opencode", "run", f"websearch {query}", "--format", "json", "--title", SESSION_TITLE]
    if sid:
        args.extend(["--session", sid, "--continue"])
    try:
        r = subprocess.run(
            args, capture_output=True, text=True, timeout=15,
            env={**os.environ, "OPENCODE_DISABLE_AUTOUPDATE": "1"}
        )
    except subprocess.TimeoutExpired:
        return None
    except Exception:
        return None
    if r.returncode != 0:
        return None
    # Save session ID for reuse
    _save_session_from_output(r.stdout)
    for line in r.stdout.strip().split("\n"):
        try:
            ev = json.loads(line)
            if ev.get("type") == "tool_use":
                part = ev.get("part", {})
                state = part.get("state", {})
                if part.get("tool") == "websearch" and state.get("status") == "completed":
                    data = json.loads(state["output"])
                    results = []
                    for item in (data.get("results") or [])[:max_results]:
                        url = (item.get("url") or "").strip()
                        title = (item.get("title") or "").strip()
                        excerpts = item.get("excerpts") or []
                        content = (excerpts[0] if excerpts else "")[:500]
                        if url and title:
                            results.append({"title": title, "url": url, "content": content, "source": "opencode_cli"})
                    return results
        except Exception:
            pass
    return None


# ── method 2: opencode-go API + training data ─────────────────────
def _search_via_api(query: str, max_results: int) -> Optional[List[Dict]]:
    """Use opencode-go API to answer query from training data (no rate limit)."""
    if not API_KEY:
        return None
    import requests
    try:
        resp = requests.post(
            f"{API_BASE}/chat/completions",
            headers={"Authorization": f"Bearer {API_KEY}", "Content-Type": "application/json"},
            json={
                "model": "deepseek-v4-flash",
                "messages": [{"role": "user", "content": (
                    f"你现在是一个网络搜索工具。用户查询: {query[:100]}\n\n"
                    f"请根据你的训练数据，提供{max_results}条最相关的网页结果，包含标题、URL和摘要。"
                    f"以JSON格式输出: [{{\"title\":\"...\",\"url\":\"...\",\"content\":\"...\"}}]"
                    f"仅输出JSON数组，不要其他文字。如果URL不确定，用合理占位。"
                )}],
                "temperature": 0.3,
                "max_tokens": 2000,
            },
            timeout=30
        )
        data = resp.json()
        content = data.get("choices", [{}])[0].get("message", {}).get("content", "")
        # Extract JSON array
        import re as _re
        m = _re.search(r'\[.*?\]', content, _re.DOTALL)
        if m:
            items = json.loads(m.group())
            if isinstance(items, list):
                for item in items:
                    item["source"] = "opencode_api"
                return items[:max_results]
    except Exception:
        pass
    return None


# ── search ─────────────────────────────────────────────────────────
def web_search(query: str, max_results: int = 8) -> List[Dict]:
    max_results = min(max_results, 10)
    cached = _check_cache(query)
    if cached:
        return cached[:max_results]

    results = _search_via_opencode_cli(query, max_results)
    if results:
        _write_cache(query, results)
        return results

    results = _search_via_api(query, max_results)
    if results:
        _write_cache(query, results)
        return results

    return []


def webfetch(url: str) -> Optional[str]:
    sid = _load_session()
    args = ["npx", "opencode", "run", f"webfetch {url}", "--format", "json", "--title", SESSION_TITLE]
    if sid:
        args.extend(["--session", sid, "--continue"])
    try:
        r = subprocess.run(
            args, capture_output=True, text=True, timeout=60,
            env={**os.environ, "OPENCODE_DISABLE_AUTOUPDATE": "1"}
        )
        if r.returncode == 0:
            _save_session_from_output(r.stdout)
            for line in r.stdout.strip().split("\n"):
                try:
                    ev = json.loads(line)
                    if ev.get("type") == "tool_use":
                        p = ev.get("part", {})
                        s = p.get("state", {})
                        if p.get("tool") == "webfetch" and s.get("status") == "completed":
                            return s.get("output", "")[:10000]
                except Exception:
                    pass
    except Exception:
        pass
    return None


# ── MCP protocol (JSON-RPC 2.0 over stdio) ────────────────────────
def _read_msg() -> Optional[Dict]:
    line = sys.stdin.readline()
    if not line:
        return None
    try:
        return json.loads(line)
    except json.JSONDecodeError:
        return None

def _send_msg(msg: Dict):
    sys.stdout.write(json.dumps(msg, ensure_ascii=False) + "\n")
    sys.stdout.flush()

def _send_error(req_id: Any, code: int, message: str):
    _send_msg({"jsonrpc": "2.0", "id": req_id, "error": {"code": code, "message": message}})

def _send_result(req_id: Any, result: Any):
    _send_msg({"jsonrpc": "2.0", "id": req_id, "result": result})


def serve():
    sys.stdin.reconfigure(encoding="utf-8")
    sys.stdout.reconfigure(encoding="utf-8")
    while True:
        msg = _read_msg()
        if msg is None:
            break
        req_id = msg.get("id")
        method = msg.get("method", "")
        params = msg.get("params", {})
        if method == "initialize":
            _send_result(req_id, {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {"listChanged": False}},
                "serverInfo": {"name": "opencode-search-mcp", "version": "1.0.0"}
            })
        elif method == "notifications/initialized":
            pass
        elif method == "tools/list":
            _send_result(req_id, {"tools": [
                {"name": "web_search", "description": "Search the web. Returns up to 10 results with title, url, content.", "inputSchema": {
                    "type": "object", "properties": {
                        "query": {"type": "string", "description": "Search query"},
                        "max_results": {"type": "number", "description": "Max results (1-10)", "default": 8}
                    }, "required": ["query"]
                }},
                {"name": "webfetch", "description": "Fetch and extract content from a URL.", "inputSchema": {
                    "type": "object", "properties": {"url": {"type": "string", "description": "URL to fetch"}},
                    "required": ["url"]
                }}
            ]})
        elif method == "tools/call":
            name = params.get("name", "")
            args = params.get("arguments", {})
            try:
                if name == "web_search":
                    results = web_search(args.get("query", ""), int(args.get("max_results", 8)))
                    _send_result(req_id, {"content": [{"type": "text", "text": json.dumps(results, ensure_ascii=False)}]})
                elif name == "webfetch":
                    content = webfetch(args.get("url", ""))
                    _send_result(req_id, {"content": [{"type": "text", "text": content or "Failed to fetch URL"}]})
                else:
                    _send_error(req_id, -32601, f"Unknown tool: {name}")
            except Exception as e:
                _send_error(req_id, -32603, str(e))
        elif method == "shutdown":
            _send_result(req_id, {})
            break
        else:
            _send_error(req_id, -32601, f"Unknown method: {method}")


def main():
    if "--query" in sys.argv:
        idx = sys.argv.index("--query")
        q = sys.argv[idx + 1] if idx + 1 < len(sys.argv) else ""
        print(json.dumps(web_search(q), ensure_ascii=False, indent=2))
        return
    if "--url" in sys.argv:
        idx = sys.argv.index("--url")
        u = sys.argv[idx + 1] if idx + 1 < len(sys.argv) else ""
        print(webfetch(u) or "Failed")
        return
    serve()


if __name__ == "__main__":
    main()