#!/usr/bin/env python3
import sys
import json
import os
import urllib.request
import urllib.error

# Configuration
RAG_API_URL = os.environ.get("RAG_API_URL", "http://localhost:8010")

def log(msg):
    """Logs a message to stderr (since stdout is reserved for JSON-RPC messages)."""
    sys.stderr.write(f"[MCP-Server] {msg}\n")
    sys.stderr.flush()

def make_api_request(endpoint, data):
    """Sends a POST request to the RAG Engine FastAPI server using built-in urllib."""
    url = f"{RAG_API_URL.rstrip('/')}/{endpoint.lstrip('/')}"
    headers = {"Content-Type": "application/json"}
    req_body = json.dumps(data).encode("utf-8")
    
    log(f"Calling endpoint: {url} with parameters: {data}")
    
    try:
        req = urllib.request.Request(url, data=req_body, headers=headers, method="POST")
        with urllib.request.urlopen(req, timeout=90) as response:
            res_body = response.read().decode("utf-8")
            return json.loads(res_body), None
    except urllib.error.URLError as e:
        err_msg = f"Failed to connect to RAG Engine API at {url}. Error: {e.reason if hasattr(e, 'reason') else e}"
        log(f"API Error: {err_msg}")
        return None, err_msg
    except Exception as e:
        err_msg = f"Unexpected error connecting to RAG Engine API: {str(e)}"
        log(f"API Error: {err_msg}")
        return None, err_msg

def handle_initialize(request_id):
    """Handles the initialize request from the MCP client."""
    response = {
        "jsonrpc": "2.0",
        "id": request_id,
        "result": {
            "protocolVersion": "2024-11-05",
            "capabilities": {
                "tools": {}
            },
            "serverInfo": {
                "name": "juniper-rag",
                "version": "1.0.0"
            }
        }
    }
    return response

def handle_tools_list(request_id):
    """Returns the list of available tools to the MCP client."""
    tools = [
        {
            "name": "search_juniper_kb",
            "description": "Semantic search in Juniper Technical Documents database for relevant information. Returns a list of matching chunks with similarity scores, source paths/URLs, and page numbers.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The search query (e.g. 'BGP down on SRX')"
                    },
                    "top_k": {
                        "type": "integer",
                        "description": "Number of matching chunks to retrieve (default: 5)",
                        "default": 5
                    },
                    "source": {
                        "type": "string",
                        "description": "Filter by source type: 'kb' (Juniper KB articles) or 'book' (Networking books/PDFs)",
                        "enum": ["kb", "book"]
                    }
                },
                "required": ["query"]
            }
        },
        {
            "name": "ask_juniper_rag",
            "description": "Ask a question about Juniper network setups. Retrieves relevant documentation and uses Ollama to synthesize a detailed answer with citations.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "query": {
                        "type": "string",
                        "description": "The question to answer"
                    },
                    "top_k": {
                        "type": "integer",
                        "description": "Number of chunks to retrieve for LLM context (default: 5)",
                        "default": 5
                    },
                    "source": {
                        "type": "string",
                        "description": "Filter by source type: 'kb' (Juniper KB articles) or 'book' (Networking books/PDFs)",
                        "enum": ["kb", "book"]
                    },
                    "llm_model": {
                        "type": "string",
                        "description": "Ollama LLM model name to use for synthesis (default: 'qwen2.5' or fallback model)",
                        "default": "qwen2.5"
                    }
                },
                "required": ["query"]
            }
        }
    ]
    
    response = {
        "jsonrpc": "2.0",
        "id": request_id,
        "result": {
            "tools": tools
        }
    }
    return response

def handle_tools_call(request_id, params):
    """Executes a tool and returns the response."""
    tool_name = params.get("name")
    arguments = params.get("arguments", {})
    
    if tool_name == "search_juniper_kb":
        # Prepare parameters for FastAPI /search endpoint
        query = arguments.get("query")
        top_k = arguments.get("top_k", 5)
        source = arguments.get("source")
        
        payload = {"query": query, "top_k": top_k}
        if source:
            payload["source"] = source
            
        data, err = make_api_request("search", payload)
        if err:
            text_content = f"Error performing search: {err}\nMake sure RAG Engine is running."
        else:
            # Format results in a nice text block
            blocks = []
            for idx, item in enumerate(data, 1):
                block = f"[{idx}] Title: {item.get('title')}\n"
                block += f"    Source Type: {item.get('source_type').upper()}\n"
                if item.get("page_num"):
                    block += f"    Page: {item.get('page_num')}\n"
                if item.get("url"):
                    block += f"    URL: {item.get('url')}\n"
                block += f"    Similarity Score: {item.get('score'):.4f}\n"
                block += f"    Content:\n{item.get('text')}\n"
                blocks.append(block)
            
            if not blocks:
                text_content = "No matching document chunks found."
            else:
                text_content = "\n---\n".join(blocks)
                
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "content": [
                    {"type": "text", "text": text_content}
                ]
            }
        }
        
    elif tool_name == "ask_juniper_rag":
        # Prepare parameters for FastAPI /generate endpoint
        query = arguments.get("query")
        top_k = arguments.get("top_k", 5)
        source = arguments.get("source")
        llm_model = arguments.get("llm_model", "qwen2.5")
        
        payload = {"query": query, "top_k": top_k, "llm_model": llm_model}
        if source:
            payload["source"] = source
            
        data, err = make_api_request("generate", payload)
        if err:
            text_content = f"Error generating answer: {err}\nMake sure RAG Engine and Ollama are running."
        else:
            answer = data.get("answer", "")
            sources = data.get("sources", [])
            
            text_content = f"### Answer:\n{answer}\n\n"
            text_content += "### Sources:\n"
            
            for idx, item in enumerate(sources, 1):
                source_str = f"- [{idx}] {item.get('title')} ({item.get('source_type').upper()})"
                if item.get("page_num"):
                    source_str += f" - Page {item.get('page_num')}"
                if item.get("url"):
                    source_str += f" - {item.get('url')}"
                text_content += f"{source_str}\n"
                
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "result": {
                "content": [
                    {"type": "text", "text": text_content}
                ]
            }
        }
        
    else:
        return {
            "jsonrpc": "2.0",
            "id": request_id,
            "error": {
                "code": -32601,
                "message": f"Method not found: {tool_name}"
            }
        }

def main():
    """Main input loop reading from stdin and writing to stdout."""
    log("Starting Juniper RAG MCP Stdio Server...")
    log(f"Configured RAG Engine API URL: {RAG_API_URL}")
    
    while True:
        try:
            line = sys.stdin.readline()
            if not line:
                break
                
            line = line.strip()
            if not line:
                continue
                
            log(f"Received raw line: {line[:200]}...")
            request = json.loads(line)
            method = request.get("method")
            request_id = request.get("id")
            
            # MCP initialization workflow and routing
            if method == "initialize":
                response = handle_initialize(request_id)
            elif method == "notifications/initialized":
                # client notification, no response required
                continue
            elif method == "tools/list":
                response = handle_tools_list(request_id)
            elif method == "tools/call":
                params = request.get("params", {})
                response = handle_tools_call(request_id, params)
            elif method == "ping":
                response = {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": {}
                }
            else:
                # Unknown method or notification
                if request_id is not None:
                    response = {
                        "jsonrpc": "2.0",
                        "id": request_id,
                        "error": {
                            "code": -32601,
                            "message": f"Method not found: {method}"
                        }
                    }
                else:
                    continue
            
            # Send response back to stdout
            sys.stdout.write(json.dumps(response) + "\n")
            sys.stdout.flush()
            log(f"Sent response for request id: {request_id}")
            
        except json.JSONDecodeError:
            log("Received invalid JSON data.")
        except Exception as e:
            log(f"Error in main loop: {str(e)}")

if __name__ == "__main__":
    main()
