#!/usr/bin/env python3
"""Context7 extension - Query up-to-date documentation for any programming library"""

import json
import urllib.request
import urllib.parse
import os

CONTEXT7_API = "https://context7.com/api"
API_KEY = os.environ.get("CONTEXT7_API_KEY")  # Optional, for higher rate limits

def context7_get(endpoint, params=None):
    """Make GET request to Context7 API"""
    try:
        url = f"{CONTEXT7_API}{endpoint}"
        if params:
            url += "?" + urllib.parse.urlencode(params)

        headers = {
            "X-Context7-Source": "mcp",
            "Accept": "application/json"
        }

        # Add API key if available
        if API_KEY:
            headers["Authorization"] = f"Bearer {API_KEY}"

        req = urllib.request.Request(url, headers=headers)
        resp = urllib.request.urlopen(req, timeout=30)

        content_type = resp.headers.get("Content-Type", "")
        if "application/json" in content_type:
            return json.loads(resp.read().decode())
        else:
            return resp.read().decode()
    except Exception as e:
        return {"error": str(e)}

@register_tool(
    "context7_search",
    "Search for libraries by name (e.g., 'react', 'express', 'django')",
    {"query": "string"}
)
def context7_search(args):
    """Search for libraries in Context7"""
    query = args.get("query", "")
    if not query:
        return "error: query required"

    result = context7_get("/v1/search", {"query": query})

    if isinstance(result, dict) and "error" in result:
        return f"error: {result['error']}"

    if isinstance(result, dict):
        results = result.get("results", [])
    else:
        return f"error: Unexpected response format"

    if not results:
        return f"No libraries found for: {query}"

    # Format results
    output = f"Found {len(results)} libraries:\n\n"
    for lib in results[:5]:  # Top 5
        output += f"• {lib.get('title', 'Unknown')}\n"
        output += f"  ID: {lib.get('id', 'N/A')}\n"
        if lib.get('description'):
            output += f"  {lib['description']}\n"
        output += f"  Snippets: {lib.get('totalSnippets', 0)}\n"
        output += f"  Trust Score: {lib.get('trustScore', 0)}\n"
        if lib.get('versions'):
            output += f"  Versions: {', '.join(lib['versions'][:3])}\n"
        output += "\n"

    # Suggest using the top result
    if results:
        top_lib = results[0]
        lib_id = top_lib.get('id', '')
        output += f"To get documentation:\ncontext7_docs({{\"library_id\": \"{lib_id}\", \"topic\": \"your topic\"}})"

    return output

@register_tool(
    "context7_docs",
    "Get documentation for a library using its Context7 ID",
    {"library_id": "string", "topic": "string?", "tokens": "number?"}
)
def context7_docs(args):
    """Get library documentation from Context7"""
    library_id = args.get("library_id", "")
    topic = args.get("topic", "")
    tokens = args.get("tokens", 5000)

    if not library_id:
        return "error: library_id required (use context7_search first)"

    # Ensure minimum tokens
    if tokens < 1000:
        tokens = 1000

    params = {
        "tokens": tokens,
        "type": "txt"
    }
    if topic:
        params["topic"] = topic

    result = context7_get(f"/v1{library_id}", params)

    if isinstance(result, dict) and "error" in result:
        return f"error: {result['error']}"

    if not result or result == "null":
        return f"No documentation found for {library_id}. Try using context7_search first."

    return result