#!/usr/bin/env python3
"""
MCP Summary Server (Streamable HTTP transport)

Designed to work with OpenWebUI's MCP (Streamable HTTP) integration.

Summarizes documents by:
1. Checking text length
2. If short, summarizing directly with LLM
3. If long, chunking text, summarizing each chunk, then synthesizing

All processing happens server-side, keeping full text out of the chat context window.

Tools:
- summarize_document: Summarize a document (handles chunking automatically)

Auth:
- If API_KEY is set:
  - Requires: Authorization: Bearer <API_KEY>
- If API_KEY is not set:
  - No auth required (for local/internal use).
"""

import json
import os
import sys
import re
from http.server import HTTPServer, BaseHTTPRequestHandler
from typing import Any, Dict, Optional
import requests

API_KEY = os.environ.get("API_KEY", "").strip()

TOOLS_LIST: Dict[str, Any] = {
    "tools": [
        {
            "name": "summarize_document",
            "description": "Summarize a document. Automatically handles chunking for long text. Returns a concise summary without exposing the full text.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "text": {
                        "type": "string",
                        "description": "The document text to summarize"
                    },
                    "max_length": {
                        "type": "integer",
                        "description": "Maximum length of summary in words (default: 100)"
                    }
                },
                "required": ["text"]
            }
        }
    ]
}


def get_bearer_token(headers: Any) -> Optional[str]:
    auth = (headers.get("Authorization") or "").strip()
    if auth.startswith("Bearer "):
        return auth[len("Bearer "):].strip()
    return None


def require_auth(headers: Any) -> bool:
    # If API_KEY is not set, allow unauthenticated access
    if not API_KEY:
        return True

    token = get_bearer_token(headers)
    if not token or token != API_KEY:
        raise PermissionError("Missing or invalid API key")
    return True


def call_llm(text: str, system_prompt: str, max_tokens: int = 2000) -> str:
    """Make an OpenAPI-compatible LLM call."""
    openapi_url = os.environ.get("OPENAPI_URL", "http://localhost:8080/v1")
    openapi_api_key = os.environ.get("OPENAPI_API_KEY", "")
    model_name = os.environ.get("MODEL_NAME", "gpt-4o")
    timeout = int(os.environ.get("LLM_TIMEOUT", "120"))
    
    url = f"{openapi_url}/chat/completions"
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {openapi_api_key}"
    }
    
    payload = {
        "model": model_name,
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": text}
        ],
        "temperature": 0.3,
        "max_tokens": max_tokens,
        "top_p": 0.9
    }
    
    response = requests.post(url, headers=headers, json=payload, timeout=timeout)
    response.raise_for_status()
    
    data = response.json()
    return data["choices"][0]["message"]["content"]


def chunk_text(text: str) -> list:
    """Split text into chunks with overlap for summarization."""
    chunk_size = int(os.environ.get("CHUNK_SIZE", "4000"))
    overlap = int(os.environ.get("OVERLAP", "200"))
    
    if len(text) <= chunk_size:
        return [text]
    
    chunks = []
    start = 0
    
    while start < len(text):
        end = min(start + chunk_size, len(text))
        
        # Try to break at sentence/paragraph boundary
        break_point = end
        for marker in ["\n\n", "\n", ". ", "! ", "? "]:
            pos = text.rfind(marker, start + chunk_size // 2, end)
            if pos > start:
                break_point = pos
                break
        
        chunk = text[start:break_point]
        if chunk.strip():
            chunks.append(chunk)
        
        start = break_point - overlap if break_point < len(text) else len(text)
        if start >= len(text):
            break
    
    return chunks


def summarize_document(text: str, max_length: int = 100) -> dict:
    """
    Main summarization function.
    
    - If text is short, summarize directly
    - If text is long, chunk and summarize each chunk, then synthesize
    """
    original_length = len(text)
    
    text = text.strip()
    if not text:
        raise ValueError("Empty text provided")
    
    max_direct_length = int(os.environ.get("MAX_DIRECT_TEXT_LENGTH", "8000"))
    intermediate_length = int(os.environ.get("TARGET_INTERMEDIATE_SUMMARY_LENGTH", "150"))
    
    # Direct summarization for shorter texts
    if len(text) <= max_direct_length:
        system_prompt = f"""You are a precise legal assistant creating concise, accurate summaries.

Create a summary that:
- Is approximately {max_length} words
- Captures key points and important details
- Uses clear, professional language
- Preserves names, dates, and specific facts

Format as plain text without bullet points."""
        
        user_prompt = f"""Summarize the following document:

{text}

Summary:"""
        
        summary = call_llm(user_prompt, system_prompt)
        
        return {
            "summary": summary,
            "original_length": original_length,
            "method": "direct",
            "chunks": 1
        }
    
    # Chunked summarization for longer texts
    chunks = chunk_text(text)
    
    chunk_summaries = []
    for i, chunk in enumerate(chunks, 1):
        system_prompt = f"""You are a precise legal assistant creating concise, accurate summaries.

You are processing chunk {i} of {len(chunks)} from a larger document.

Create a focused summary that:
- Captures key points and important details
- Is approximately {intermediate_length} words
- Can be combined with other chunk summaries
- Uses clear, professional language
- Preserves names, dates, and specific facts

Respond as plain text without bullet points."""
        
        user_prompt = f"""Summarize this text (chunk {i} of {len(chunks)}):

{chunk}

Summary:"""
        
        chunk_summary = call_llm(user_prompt, system_prompt)
        chunk_summaries.append(chunk_summary)
    
    # Synthesize into final summary
    combined = "\n\n".join(chunk_summaries)
    
    system_prompt = """You are a precise legal assistant creating executive-level summaries.

Synthesize the provided partial summaries into a single, cohesive summary that:
- Is approximately 100 words
- Captures the complete document picture
- Is clear and professional
- Removes redundancy
- Maintains logical flow
- Preserves all critical information

Format as a single paragraph of plain text."""
    
    user_prompt = f"""Synthesize these partial summaries into one cohesive summary:

{combined}

Final summary:"""
    
    final_summary = call_llm(user_prompt, system_prompt)
    
    return {
        "summary": final_summary,
        "original_length": original_length,
        "method": "chunked",
        "chunks": len(chunks)
    }


class MCPSummaryHandler(BaseHTTPRequestHandler):
    def log_message(self, format, *args):
        # Quiet logs by default
        pass

    def _send_json(self, status: int, payload: Any):
        body = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _auth_or_401(self):
        try:
            return require_auth(self.headers)
        except PermissionError:
            self._send_json(401, {"error": "Missing or invalid API key"})
            return False

    def do_GET(self):
        # Basic info endpoint (not required by MCP, but useful)
        if self.path == "/":
            self._send_json(200, {
                "service": "mcp-summary",
                "transport": "streamable-http",
                "docs": "Use POST / with MCP JSON-RPC (initialize, tools/list, tools/call)."
            })
            return

        self.send_error(404, "Not Found")

    def do_POST(self):
        # Streamable HTTP MCP endpoint
        if self.path not in ("/", "/mcp"):
            self.send_error(404, "Not Found")
            return

        if not self._auth_or_401():
            return

        length = int(self.headers.get("Content-Length", 0))
        if length == 0:
            self._send_json(400, {"error": "Empty body"})
            return

        raw = self.rfile.read(length)
        try:
            req = json.loads(raw)
        except json.JSONDecodeError:
            self._send_json(400, {"error": "Invalid JSON"})
            return

        method = req.get("method")
        params = req.get("params") or {}
        req_id = req.get("id")

        # MCP: initialize
        if method == "initialize":
            self._send_json(200, {
                "jsonrpc": "2.0",
                "id": req_id,
                "result": {
                    "protocolVersion": "2025-11-25",
                    "capabilities": {
                        "tools": {}
                    },
                    "serverInfo": {
                        "name": "mcp-summary",
                        "version": "1.0.0"
                    }
                }
            })
            return

        # MCP: tools/list
        if method == "tools/list":
            self._send_json(200, {
                "jsonrpc": "2.0",
                "id": req_id,
                "result": TOOLS_LIST
            })
            return

        # MCP: tools/call
        if method == "tools/call":
            tool_name = params.get("name")
            tool_args = params.get("arguments") or {}
            try:
                result = self._call_tool(tool_name, tool_args)
                self._send_json(200, {
                    "jsonrpc": "2.0",
                    "id": req_id,
                    "result": {
                        "content": [
                            {"type": "text", "text": json.dumps(result, ensure_ascii=False)}
                        ]
                    }
                })
            except Exception as e:
                self._send_json(200, {
                    "jsonrpc": "2.0",
                    "id": req_id,
                    "error": {
                        "code": -32000,
                        "message": str(e)
                    }
                })
            return

        # Unknown method
        self._send_json(400, {"error": "Unknown method: " + str(method)})

    def _call_tool(self, name: str, args: Dict[str, Any]) -> Any:
        if name == "summarize_document":
            text = args.get("text")
            if not text:
                raise ValueError("Text parameter is required")
            
            max_length = args.get("max_length", 100)
            return summarize_document(text, max_length)

        raise ValueError(f"Unknown tool: {name}")


def main():
    port = int(sys.argv[1]) if len(sys.argv) > 1 else int(os.environ.get("PORT", "8080"))
    server = HTTPServer(("0.0.0.0", port), MCPSummaryHandler)
    mode = "auth enabled (Bearer)" if API_KEY else "no auth (API_KEY not set)"
    print(f"MCP Summary Server listening on 0.0.0.0:{port} [{mode}]")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down...")
        server.server_close()


if __name__ == "__main__":
    main()