#!/usr/bin/env python3
"""DeepWiki extension - MCP HTTP client for GitHub repo documentation"""

import json
import urllib.request
import time

# Global session cache
_session_id = None

def get_session_id():
    """Initialize MCP session and get session ID"""
    global _session_id
    if _session_id:
        return _session_id

    url = "https://mcp.deepwiki.com/mcp"
    data = {
        "jsonrpc": "2.0",
        "id": "init",
        "method": "initialize",
        "params": {
            "protocolVersion": "2024-11-05",
            "capabilities": {},
            "clientInfo": {"name": "nanocode", "version": "1.0"}
        }
    }

    try:
        req = urllib.request.Request(
            url,
            json.dumps(data).encode(),
            {
                "Content-Type": "application/json",
                "Accept": "application/json, text/event-stream"
            }
        )
        resp = urllib.request.urlopen(req, timeout=30)
        _session_id = resp.headers.get("mcp-session-id")
        return _session_id
    except Exception as e:
        return None

def mcp_call(tool_name, arguments):
    """Make MCP JSON-RPC call to DeepWiki"""
    session_id = get_session_id()
    if not session_id:
        return "error: Failed to initialize MCP session"

    url = "https://mcp.deepwiki.com/mcp"
    request_id = str(int(time.time() * 1000))

    data = {
        "jsonrpc": "2.0",
        "id": request_id,
        "method": "tools/call",
        "params": {
            "name": tool_name,
            "arguments": arguments
        }
    }

    try:
        req = urllib.request.Request(
            url,
            json.dumps(data).encode(),
            {
                "Content-Type": "application/json",
                "Accept": "application/json, text/event-stream",
                "mcp-session-id": session_id
            }
        )
        resp = urllib.request.urlopen(req, timeout=30)

        # Handle SSE response
        response_text = resp.read().decode()

        # Parse SSE format: extract JSON from "data: " lines
        for line in response_text.split('\n'):
            if line.startswith('data: '):
                data = line[6:].strip()  # Skip "data: " prefix

                # Skip ping events
                if data == "ping":
                    continue

                try:
                    result = json.loads(data)

                    if "error" in result:
                        return f"error: {result['error']['message']}"

                    # Extract text from MCP response
                    content = result.get("result", {}).get("content", [])
                    if content:
                        return "\n".join(item["text"] for item in content if item.get("type") == "text")
                except json.JSONDecodeError:
                    continue

        return "error: No data in response"
    except Exception as e:
        return f"error: {e}"

@register_tool(
    "deepwiki_ask",
    "Ask questions about GitHub repo",
    {"repo": "string", "question": "string"}
)
def deepwiki_ask(args):
    """Ask a question about a GitHub repository"""
    return mcp_call("ask_question", {
        "repoName": args["repo"],
        "question": args["question"]
    })

@register_tool(
    "deepwiki_structure",
    "Get repo documentation structure",
    {"repo": "string"}
)
def deepwiki_structure(args):
    """Get documentation structure/topics for a GitHub repository"""
    return mcp_call("read_wiki_structure", {
        "repoName": args["repo"]
    })

@register_tool(
    "deepwiki_contents",
    "Get repo documentation contents",
    {"repo": "string"}
)
def deepwiki_contents(args):
    """Get full documentation contents for a GitHub repository"""
    return mcp_call("read_wiki_contents", {
        "repoName": args["repo"]
    })