#!/usr/bin/env python3
"""
LangProtect MCP Gateway - Security Gateway for MCP Servers
Enhanced with Pre-LLM Scanning and Response Masking
"""

import sys
import json
import os
import subprocess
import requests
import argparse
from datetime import datetime, timedelta
from typing import Dict, List, Any, Optional
import logging

# Import response masker
from .response_masker import ResponseMasker, get_masker

log_level = os.environ.get("LOGLEVEL", "DEBUG" if os.getenv('DEBUG', 'false').lower() == 'true' else "INFO").upper()
logging.basicConfig(level=getattr(logging, log_level), format='[%(asctime)s] %(levelname)s: %(message)s', handlers=[logging.StreamHandler(sys.stderr)])
logger = logging.getLogger('langprotect-gateway')


class MCPServer:
    def __init__(self, name: str, config: Dict[str, Any]):
        self.name = name
        self.command = config.get('command')
        self.args = config.get('args', [])
        self.env = config.get('env', {})
        self.tools: List[Dict] = []
        self.process: Optional[subprocess.Popen] = None
        self._request_id = 0
        logger.info(f"Configured MCP server: {name} ({self.command} {' '.join(self.args)})")
    
    def _next_id(self) -> int:
        self._request_id += 1
        return self._request_id
    
    def start(self) -> bool:
        try:
            env = {**os.environ, **self.env}
            self.process = subprocess.Popen([self.command] + self.args, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env, bufsize=1)
            init_request = {"jsonrpc": "2.0", "id": self._next_id(), "method": "initialize", "params": {"protocolVersion": "2024-11-05", "capabilities": {}, "clientInfo": {"name": "langprotect-gateway", "version": "1.0.0"}}}
            self.process.stdin.write(json.dumps(init_request) + "\n")
            self.process.stdin.flush()
            response_line = self.process.stdout.readline()
            if not response_line:
                logger.error(f"Failed to initialize {self.name}")
                return False
            response = json.loads(response_line)
            if "error" in response:
                logger.error(f"Initialize error for {self.name}: {response['error']}")
                return False
            self.process.stdin.write(json.dumps({"jsonrpc": "2.0", "method": "notifications/initialized"}) + "\n")
            self.process.stdin.flush()
            logger.info(f"Started MCP server: {self.name}")
            return True
        except Exception as e:
            logger.error(f"Failed to start {self.name}: {e}")
            return False
    
    def stop(self):
        if self.process:
            try:
                self.process.terminate()
                self.process.wait(timeout=5)
            except:
                self.process.kill()
            self.process = None
    
    def call(self, method: str, params: Dict) -> Dict:
        if not self.process or self.process.poll() is not None:
            raise Exception(f"Server {self.name} is not running")
        request = {"jsonrpc": "2.0", "id": self._next_id(), "method": method, "params": params}
        logger.debug(f"-> {self.name}.{method}")
        try:
            self.process.stdin.write(json.dumps(request) + "\n")
            self.process.stdin.flush()
            response_line = self.process.stdout.readline()
            if not response_line:
                raise Exception(f"No response from {self.name}")
            response = json.loads(response_line)
            logger.debug(f"<- {self.name}: {str(response)[:200]}")
            return response
        except Exception as e:
            logger.error(f"Error calling {self.name}.{method}: {e}")
            raise
    
    def discover_tools(self) -> List[Dict]:
        try:
            response = self.call("tools/list", {})
            if "result" in response:
                self.tools = response["result"].get("tools", [])
                logger.info(f"Discovered {len(self.tools)} tools from {self.name}")
                return self.tools
            return []
        except Exception as e:
            logger.warning(f"Could not discover tools from {self.name}: {e}")
            return []


class LangProtectAuth:
    def __init__(self, url: str, email: str, password: str, scan_timeout: float = 5.0, fail_closed: bool = False):
        self.url = url
        self.email = email
        self.password = password
        self.jwt_token: Optional[str] = None
        self.token_expiry: Optional[datetime] = None
        self.scan_timeout = scan_timeout  # Maximum wait time for scans
        self.fail_closed = fail_closed    # Block on scan failure if True
        logger.info(f"Auth initialized: timeout={scan_timeout}s, fail_closed={fail_closed}")
    
    def login(self) -> bool:
        try:
            logger.info(f"Authenticating with {self.url}...")
            response = requests.post(f"{self.url}/v1/group-users/signin", json={'email': self.email, 'password': self.password}, timeout=10)
            if response.status_code == 200:
                data = response.json()
                self.jwt_token = data.get('access_token')
                self.token_expiry = datetime.now() + timedelta(days=6)
                logger.info("Authentication successful")
                return True
            else:
                logger.error(f"Login failed: {response.status_code}")
                return False
        except Exception as e:
            logger.error(f"Login error: {e}")
            return False
    
    def ensure_token(self) -> bool:
        if not self.jwt_token or (self.token_expiry and datetime.now() > self.token_expiry):
            return self.login()
        return True
    
    def scan_input(self, tool_name: str, arguments: Dict, server_name: str) -> Dict:
        """
        Scan user input BEFORE forwarding to MCP server (blocking scan).
        Uses the new Group Scan API with policy-based scanning.
        """
        self.ensure_token()
        try:
            # Convert tool call to prompt string for scanning
            prompt = f"Tool: {tool_name}\nServer: {server_name}\nArguments: {json.dumps(arguments, indent=2)}"
            
            payload = {
                'prompt': prompt,
                'metadata': {
                    'tool_name': tool_name,
                    'server_name': server_name,
                    'source': 'mcp-gateway-input',
                    'scan_type': 'input'
                }
            }
            
            logger.debug(f"🛡️ INPUT SCAN: {tool_name} on {server_name}")
            
            response = requests.post(
                f"{self.url}/v1/group-logs/scan",
                json=payload,
                headers={
                    'Authorization': f'Bearer {self.jwt_token}',
                    'Content-Type': 'application/json'
                },
                timeout=self.scan_timeout
            )
            
            if response.status_code != 200:
                logger.warning(f"Backend returned {response.status_code}")
                if self.fail_closed:
                    return {
                        'status': 'blocked',
                        'reason': f'Scan service unavailable (HTTP {response.status_code}) - fail-closed mode'
                    }
                else:
                    logger.warning("Allowing request (fail-open mode)")
                    return {'status': 'allowed', 'error': f'Backend error: {response.status_code}'}
            
            result = response.json()
            
            # Handle scan service timeout - respect fail mode
            if result.get('detections', {}).get('error') == 'Scan service timeout':
                logger.warning("Scan service timeout detected")
                if self.fail_closed:
                    return {
                        'status': 'blocked',
                        'reason': 'Scan service timeout - fail-closed mode'
                    }
                else:
                    logger.warning("Allowing request despite timeout (fail-open)")
                    return {'status': 'allowed', 'id': result.get('id'), 'error': 'Scan timeout'}
            
            return result
        
        except requests.exceptions.Timeout:
            logger.error(f"Backend scan timeout after {self.scan_timeout}s")
            if self.fail_closed:
                return {
                    'status': 'blocked',
                    'reason': f'Scan timeout after {self.scan_timeout}s - fail-closed mode'
                }
            else:
                logger.warning("Allowing request despite timeout (fail-open)")
                return {'status': 'allowed', 'error': 'Request timeout'}
        
        except Exception as e:
            logger.error(f"Scan error: {e}")
            if self.fail_closed:
                return {
                    'status': 'blocked',
                    'reason': f'Scan error: {str(e)} - fail-closed mode'
                }
            else:
                logger.warning(f"Allowing request despite error (fail-open): {e}")
                return {'status': 'allowed', 'error': str(e)}
    
    def scan_output(self, tool_name: str, output_content: str, prompt: str = None, metadata: Dict = None) -> Dict:
        """
        Scan LLM/MCP output AFTER receiving from server (non-blocking, masking scan).
        Uses the new Group Scan API with output scanning support.
        
        Args:
            tool_name: Name of the MCP tool that generated the output
            output_content: The output text to scan for secrets
            prompt: Original user prompt (optional)
            metadata: Additional context (optional)
        
        Returns:
            Scan result with masked_content field if secrets detected
        """
        self.ensure_token()
        try:
            payload = {
                'prompt': prompt or f"Tool: {tool_name}",
                'output': output_content,
                'metadata': {
                    'tool_name': tool_name,
                    'source': 'mcp-gateway-output',
                    'scan_type': 'output',
                    **(metadata or {})
                }
            }
            
            logger.debug(f"🔍 OUTPUT SCAN: {tool_name} ({len(output_content)} chars)")
            
            response = requests.post(
                f"{self.url}/v1/group-logs/scan",
                json=payload,
                headers={
                    'Authorization': f'Bearer {self.jwt_token}',
                    'Content-Type': 'application/json'
                },
                timeout=self.scan_timeout
            )
            
            if response.status_code != 200:
                logger.warning(f"Output scan failed: HTTP {response.status_code}")
                # For output scanning, fail-open (don't block, return original)
                return {
                    'status': 'allowed',
                    'output': output_content,
                    'masked': False,
                    'error': f'Scan failed: {response.status_code}'
                }
            
            result = response.json()
            
            # Extract masked content from MCPResponseScanner details
            mcp_response = result.get('detections', {}).get('MCPResponseScanner', {})
            if mcp_response.get('is_detected'):
                masked_content = mcp_response.get('details', {}).get('masked_content', output_content)
                logger.warning(f"🔒 OUTPUT MASKED: {tool_name} (score={mcp_response.get('score')})")
                return {
                    'status': result.get('status'),
                    'output': masked_content,
                    'masked': True,
                    'risk_score': result.get('risk_score'),
                    'scan_id': result.get('id'),
                    'detections': mcp_response.get('details', {}).get('detections', [])
                }
            else:
                # No secrets detected, return original
                return {
                    'status': 'safe',
                    'output': output_content,
                    'masked': False
                }
        
        except Exception as e:
            logger.error(f"Output scan error: {e}")
            # Fail-open for output scanning - return original content
            return {
                'status': 'allowed',
                'output': output_content,
                'masked': False,
                'error': str(e)
            }


class LangProtectGateway:
    def __init__(self, mcp_json_path: Optional[str] = None):
        self.mcp_json_path = mcp_json_path
        
        # Load credentials from env vars first, then potentially from config
        self.langprotect_url = os.getenv('LANGPROTECT_URL', 'http://localhost:8000')
        self.email = os.getenv('LANGPROTECT_EMAIL')
        self.password = os.getenv('LANGPROTECT_PASSWORD')
        
        # Security configuration
        self.scan_timeout = float(os.getenv('LANGPROTECT_SCAN_TIMEOUT', '5.0'))
        self.fail_closed = os.getenv('LANGPROTECT_FAIL_CLOSED', 'false').lower() == 'true'
        self.enable_masking = os.getenv('LANGPROTECT_ENABLE_MASKING', 'true').lower() == 'true'
        self.enable_entropy_detection = os.getenv('LANGPROTECT_ENTROPY_DETECTION', 'true').lower() == 'true'
        
        # Try to load credentials from mcp.json env section (like Lasso)
        if mcp_json_path and (not self.email or not self.password):
            self._load_env_from_config(mcp_json_path)
        
        self.auth: Optional[LangProtectAuth] = None
        self.mcp_servers: Dict[str, MCPServer] = {}
        self.tool_to_server: Dict[str, str] = {}
        self.all_tools: List[Dict] = []
        
        # Initialize response masker
        self.masker: Optional[ResponseMasker] = None
        if self.enable_masking:
            self.masker = get_masker(
                enable_entropy=self.enable_entropy_detection,
                entropy_threshold=4.5
            )
            logger.info("✅ Response masking ENABLED")
        else:
            logger.warning("⚠️ Response masking DISABLED")
        
        logger.debug(f"LANGPROTECT_URL: {self.langprotect_url}")
        logger.debug(f"LANGPROTECT_EMAIL: {self.email}")
        logger.info(f"Security config: timeout={self.scan_timeout}s, fail_closed={self.fail_closed}, masking={self.enable_masking}")
    
    def _load_env_from_config(self, path: str):
        """Load credentials from mcp.json env section (Lasso/VS Code style)"""
        try:
            expanded_path = os.path.expanduser(path)
            with open(expanded_path, 'r') as f:
                config = json.load(f)
            
            # Look for env vars in the gateway's config section
            # Check both mcpServers (Cursor/Claude) and servers (VS Code)
            for top_key in ['mcpServers', 'servers']:
                top_config = config.get(top_key, {})
                for gateway_name in ['langprotect-gateway', 'langprotect', 'mcp-gateway']:
                    gateway_config = top_config.get(gateway_name, {})
                    env_section = gateway_config.get('env', {})
                    if env_section:
                        if not self.langprotect_url or self.langprotect_url == 'http://localhost:8000':
                            self.langprotect_url = env_section.get('LANGPROTECT_URL', self.langprotect_url)
                        if not self.email:
                            self.email = env_section.get('LANGPROTECT_EMAIL')
                        if not self.password:
                            self.password = env_section.get('LANGPROTECT_PASSWORD')
                        logger.info(f"Loaded credentials from config env section ({top_key}.{gateway_name})")
                        return
        except Exception as e:
            logger.debug(f"Could not load env from config: {e}")
    
    def initialize(self) -> bool:
        if self.email and self.password:
            self.auth = LangProtectAuth(
                self.langprotect_url,
                self.email,
                self.password,
                scan_timeout=self.scan_timeout,
                fail_closed=self.fail_closed
            )
            if not self.auth.login():
                logger.error("Failed to authenticate with LangProtect backend")
                return False
        else:
            logger.warning("No LangProtect credentials - running in pass-through mode")
        if not self.load_servers():
            return False
        if not self.start_servers():
            return False
        logger.info("=" * 60)
        logger.info("🛡️  LangProtect Gateway initialized")
        logger.info(f"Backend: {self.langprotect_url}")
        logger.info(f"Servers: {len(self.mcp_servers)}")
        logger.info(f"Tools: {len(self.all_tools)}")
        logger.info(f"Security: fail_closed={self.fail_closed}, masking={self.enable_masking}")
        logger.info("=" * 60)
        return True
    
    def load_servers(self) -> bool:
        # Mode 1: Single server via environment variables (for wrapper scripts)
        mcp_command = os.getenv('MCP_SERVER_COMMAND')
        mcp_args = os.getenv('MCP_SERVER_ARGS')
        if mcp_command:
            logger.info(f"Single server mode: {mcp_command}")
            args_list = [arg.strip() for arg in mcp_args.split(',')] if mcp_args else []
            server_name = os.getenv('MCP_SERVER_NAME', 'proxied-server')
            self.mcp_servers[server_name] = MCPServer(server_name, {'command': mcp_command, 'args': args_list, 'env': {}})
            return True
        
        # Mode 2: Config file (mcp.json)
        if self.mcp_json_path:
            return self.load_from_mcp_json(self.mcp_json_path)
        
        logger.warning("No MCP servers configured")
        return False
    
    def load_from_mcp_json(self, path: str) -> bool:
        try:
            expanded_path = os.path.expanduser(path)
            with open(expanded_path, 'r') as f:
                config = json.load(f)
            
            # Try multiple config structures:
            # 1. Cursor/Claude: mcpServers.langprotect-gateway.servers (nested)
            # 2. VS Code: servers.langprotect-gateway.servers (nested)
            # 3. Direct: servers or mcpServers (flat)
            
            servers = {}
            
            # Check for Cursor/Claude-style nested config (mcpServers.X.servers)
            mcp_servers = config.get('mcpServers', {})
            for gateway_name in ['langprotect-gateway', 'langprotect', 'mcp-gateway']:
                gateway_config = mcp_servers.get(gateway_name, {})
                if 'servers' in gateway_config:
                    servers = gateway_config['servers']
                    logger.info(f"Found nested servers config under mcpServers.{gateway_name}.servers")
                    break
            
            # Check for VS Code-style nested config (servers.X.servers)
            if not servers:
                vscode_servers = config.get('servers', {})
                for gateway_name in ['langprotect-gateway', 'langprotect', 'mcp-gateway']:
                    gateway_config = vscode_servers.get(gateway_name, {})
                    if 'servers' in gateway_config:
                        servers = gateway_config['servers']
                        logger.info(f"Found nested servers config under servers.{gateway_name}.servers")
                        break
            
            # Fallback to direct/flat config
            if not servers:
                servers = config.get('servers', config.get('mcpServers', {}))
            
            if not servers:
                logger.error("No servers found in config file")
                return False
            
            for name, cfg in servers.items():
                # Skip gateway self-references
                if name in ['langprotect-gateway', 'langprotect', 'mcp-gateway']:
                    continue
                self.mcp_servers[name] = MCPServer(name, cfg)
            
            logger.info(f"Loaded {len(self.mcp_servers)} servers from config")
            return len(self.mcp_servers) > 0
        except Exception as e:
            logger.error(f"Error loading {path}: {e}")
            return False
    
    def start_servers(self) -> bool:
        started = 0
        for name, server in list(self.mcp_servers.items()):
            if server.start():
                tools = server.discover_tools()
                for tool in tools:
                    tool_name = tool.get('name')
                    if tool_name:
                        self.tool_to_server[tool_name] = name
                        tool_copy = tool.copy()
                        tool_copy['description'] = f"[{name}] {tool_copy.get('description', '')}"
                        self.all_tools.append(tool_copy)
                started += 1
            else:
                del self.mcp_servers[name]
        return started > 0
    
    def shutdown(self):
        for server in self.mcp_servers.values():
            server.stop()
    
    def handle_request(self, request: Dict) -> Optional[Dict]:
        method = request.get('method')
        request_id = request.get('id')
        params = request.get('params', {})
        logger.info(f"Request: {method} (id={request_id})")
        try:
            if method == 'initialize':
                return {'jsonrpc': '2.0', 'id': request_id, 'result': {'protocolVersion': '2024-11-05', 'capabilities': {'tools': {}}, 'serverInfo': {'name': 'langprotect-gateway', 'version': '1.0.0'}}}
            elif method == 'notifications/initialized':
                return None
            elif method == 'tools/list':
                logger.info(f"Returning {len(self.all_tools)} tools")
                return {'jsonrpc': '2.0', 'id': request_id, 'result': {'tools': self.all_tools}}
            elif method == 'tools/call':
                return self._handle_call_tool(request_id, params)
            else:
                return {'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32601, 'message': f'Method not found: {method}'}}
        except Exception as e:
            logger.error(f"Error handling {method}: {e}")
            return {'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32603, 'message': str(e)}}
    
    def _handle_call_tool(self, request_id, params: Dict) -> Dict:
        tool_name = params.get('name', '')
        arguments = params.get('arguments', {})
        server_name = self.tool_to_server.get(tool_name)
        
        if not server_name:
            return {'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32602, 'message': f'Unknown tool: {tool_name}'}}
        
        server = self.mcp_servers.get(server_name)
        if not server:
            return {'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32602, 'message': f'Server not found: {server_name}'}}
        
        logger.info(f"Tool call: {server_name}.{tool_name}")
        
        # 🛡️ LAYER 1: INPUT SCAN (synchronous blocking scan)
        scan_result = None
        if self.auth:
            scan_result = self.auth.scan_input(tool_name, arguments, server_name)
            status = scan_result.get('status', '').lower()
            
            if status == 'blocked':
                reason = scan_result.get('reason', 'Policy violation')
                logger.warning(f"🚫 INPUT BLOCKED: {tool_name} - {reason}")
                return {
                    'jsonrpc': '2.0',
                    'id': request_id,
                    'error': {
                        'code': -32000,
                        'message': f'🛡️ LangProtect: {reason}'
                    }
                }
            
            logger.info(f"✅ INPUT ALLOWED (log_id={scan_result.get('id')})")
        
        # Input scan passed or no auth - forward to MCP server
        try:
            response = server.call('tools/call', {'name': tool_name, 'arguments': arguments})
            
            # 🛡️ LAYER 2: OUTPUT SCAN (scan response content for secrets)
            if self.auth and self.enable_masking and 'result' in response:
                # Extract text content from response
                result_content = response.get('result', {})
                output_text = self._extract_text_from_result(result_content)
                
                if output_text:
                    logger.debug(f"🔍 Scanning output: {len(output_text)} chars")
                    output_scan = self.auth.scan_output(
                        tool_name=tool_name,
                        output_content=output_text,
                        prompt=json.dumps(arguments),
                        metadata={'server_name': server_name}
                    )
                    
                    # 🚨 CRITICAL: Check if backend blocked the request
                    if output_scan.get('status') == 'blocked' or output_scan.get('action') == 'blocked':
                        risk_score = output_scan.get('risk_score', 0)
                        severity = output_scan.get('severity', 'high')
                        logger.error(f"🚫 OUTPUT BLOCKED: {tool_name} (risk={risk_score}, severity={severity})")
                        raise ValueError(
                            f"Request blocked by security policy: {tool_name}\n"
                            f"Risk Score: {risk_score}\n"
                            f"Severity: {severity}\n"
                            f"This request was blocked because it would return sensitive data."
                        )
                    
                    if output_scan.get('masked'):
                        # Replace output with masked version
                        masked_text = output_scan.get('output', output_text)
                        logger.warning(f"🔒 OUTPUT MASKED: {tool_name} (risk={output_scan.get('risk_score')})")
                        response['result'] = self._replace_text_in_result(result_content, masked_text)
            
            # Return formatted response
            if 'result' in response:
                return {'jsonrpc': '2.0', 'id': request_id, 'result': response['result']}
            elif 'error' in response:
                return {'jsonrpc': '2.0', 'id': request_id, 'error': response['error']}
            return response
        
        except Exception as e:
            logger.error(f"Error executing {tool_name}: {e}")
            return {'jsonrpc': '2.0', 'id': request_id, 'error': {'code': -32603, 'message': f'Error executing tool: {e}'}}
    
    def _extract_text_from_result(self, result: Any) -> str:
        """Extract text content from MCP tool result for scanning."""
        if isinstance(result, str):
            return result
        elif isinstance(result, dict):
            # MCP results typically have 'content' field
            if 'content' in result:
                content = result['content']
                if isinstance(content, str):
                    return content
                elif isinstance(content, list):
                    # Extract text from content array
                    texts = []
                    for item in content:
                        if isinstance(item, dict) and item.get('type') == 'text':
                            texts.append(item.get('text', ''))
                        elif isinstance(item, str):
                            texts.append(item)
                    return '\n'.join(texts)
            # Fallback: convert whole dict to string
            return json.dumps(result)
        elif isinstance(result, list):
            return '\n'.join(str(item) for item in result)
        return str(result)
    
    def _replace_text_in_result(self, result: Any, masked_text: str) -> Any:
        """Replace text content in MCP tool result with masked version."""
        if isinstance(result, str):
            return masked_text
        elif isinstance(result, dict):
            result_copy = result.copy()
            if 'content' in result_copy:
                content = result_copy['content']
                if isinstance(content, str):
                    result_copy['content'] = masked_text
                elif isinstance(content, list):
                    # Replace text in content array
                    masked_lines = masked_text.split('\n')
                    new_content = []
                    line_idx = 0
                    for item in content:
                        if isinstance(item, dict) and item.get('type') == 'text':
                            if line_idx < len(masked_lines):
                                new_item = item.copy()
                                new_item['text'] = masked_lines[line_idx]
                                new_content.append(new_item)
                                line_idx += 1
                        else:
                            new_content.append(item)
                    result_copy['content'] = new_content
            return result_copy
        return masked_text
    
    def _log_mask_events(self, mask_events: List[Dict], scan_id: Optional[str], tool_name: str):
        """Log masked secrets to backend for audit trail"""
        if not self.auth or not mask_events:
            return
        
        try:
            self.auth.ensure_token()
            payload = {
                'scan_id': scan_id,
                'tool_name': tool_name,
                'mask_events': mask_events,
                'timestamp': datetime.now().isoformat(),
                'gateway_version': '1.0.0'
            }
            
            # Fire-and-forget (don't block on logging)
            response = requests.post(
                f"{self.langprotect_url}/v1/mask-events",
                json=payload,
                headers={'Authorization': f'Bearer {self.auth.jwt_token}'},
                timeout=2  # Short timeout for logging
            )
            
            if response.status_code == 200:
                logger.debug(f"Logged {len(mask_events)} mask events to backend")
            else:
                logger.warning(f"Failed to log mask events: {response.status_code}")
        
        except Exception as e:
            logger.warning(f"Failed to log mask events: {e}")
    
    def run(self):
        try:
            for line in sys.stdin:
                line = line.strip()
                if not line:
                    continue
                try:
                    request = json.loads(line)
                    response = self.handle_request(request)
                    if response:
                        print(json.dumps(response), flush=True)
                except json.JSONDecodeError as e:
                    print(json.dumps({'jsonrpc': '2.0', 'error': {'code': -32700, 'message': f'Parse error: {e}'}}), flush=True)
        except KeyboardInterrupt:
            pass
        finally:
            self.shutdown()


def parse_args():
    parser = argparse.ArgumentParser(description='LangProtect MCP Gateway')
    parser.add_argument('--mcp-json-path', type=str, help='Path to mcp.json configuration file')
    parser.add_argument('--debug', action='store_true', help='Enable debug logging')
    return parser.parse_args()


def main():
    args = parse_args()
    if args.debug:
        os.environ['DEBUG'] = 'true'
        logging.getLogger('langprotect-gateway').setLevel(logging.DEBUG)
    gateway = LangProtectGateway(mcp_json_path=args.mcp_json_path)
    if not gateway.initialize():
        sys.exit(1)
    gateway.run()


if __name__ == '__main__':
    main()
