#!/usr/bin/env python3
"""
LangProtect MCP Gateway Setup Helper
Automatically configures VS Code for global MCP gateway usage
"""

import os
import json
import sys
import getpass
import urllib.request
import urllib.error
from pathlib import Path


def get_vscode_settings_path():
    """Get the VS Code user settings path based on OS"""
    home = Path.home()
    
    if sys.platform == "darwin":  # macOS
        return home / "Library/Application Support/Code/User/settings.json"
    elif sys.platform == "win32":  # Windows
        return home / "AppData/Roaming/Code/User/settings.json"
    else:  # Linux
        return home / ".config/Code/User/settings.json"


def validate_credentials(url, email, password):
    """Validate credentials against the backend API"""
    try:
        import json
        
        # Prepare the request
        data = json.dumps({"email": email, "password": password}).encode('utf-8')
        req = urllib.request.Request(
            f"{url}/v1/group-users/signin",
            data=data,
            headers={'Content-Type': 'application/json'}
        )
        
        # Make the request
        with urllib.request.urlopen(req, timeout=10) as response:
            return response.status in [200, 201]
            
    except urllib.error.HTTPError as e:
        # Parse error message if available
        try:
            error_body = e.read().decode('utf-8')
            error_data = json.loads(error_body)
            error_msg = error_data.get('detail', error_data.get('message', 'Authentication failed'))
            print(f"      ✗ {error_msg}")
        except:
            print(f"      ✗ Authentication failed (HTTP {e.code})")
        return False
    except urllib.error.URLError as e:
        print(f"      ✗ Cannot connect to {url}")
        print(f"        Make sure the backend is running and accessible")
        return False
    except Exception as e:
        print(f"      ✗ Error: {e}")
        return False


def prompt_credentials():
    """Interactively prompt user for credentials with validation"""
    print()
    print("═" * 65)
    print("         🔐 Enter Your LangProtect Credentials")
    print("═" * 65)
    print()
    
    while True:
        # Prompt for URL
        url = input("Backend URL [http://localhost:8000]: ").strip()
        if not url:
            url = "http://localhost:8000"
        
        # Prompt for email
        email = input("Email: ").strip()
        if not email:
            print("✗ Email cannot be empty!")
            print()
            continue
        
        # Prompt for password (hidden)
        password = getpass.getpass("Password: ")
        if not password:
            print("✗ Password cannot be empty!")
            print()
            continue
        
        # Validate credentials
        print("      Validating credentials...")
        if validate_credentials(url, email, password):
            print("      ✓ Credentials validated successfully!")
            print()
            return url, email, password
        else:
            print()
            print("Please try again or press Ctrl+C to cancel.")
            print()


def create_wrapper_script(url=None, email=None, password=None):
    """Create the global wrapper script with credentials"""
    wrapper_dir = Path.home() / ".local/bin"
    wrapper_dir.mkdir(parents=True, exist_ok=True)
    
    wrapper_path = wrapper_dir / "langprotect-mcp-wrapper.sh"
    
    # Check if credentials provided via environment variables
    if not url or not email or not password:
        url = os.environ.get('LANGPROTECT_URL')
        email = os.environ.get('LANGPROTECT_EMAIL')
        password = os.environ.get('LANGPROTECT_PASSWORD')
    
    # If still not provided, prompt user
    if not url or not email or not password:
        url, email, password = prompt_credentials()
    else:
        # Validate environment credentials
        print("   Using credentials from environment variables...")
        print("   Validating...")
        if not validate_credentials(url, email, password):
            print("   ✗ Environment credentials invalid. Please enter manually:")
            url, email, password = prompt_credentials()
        else:
            print("   ✓ Environment credentials validated!")
    
    # Create wrapper with actual credentials
    wrapper_content = f"""#!/bin/bash
# LangProtect MCP Gateway Wrapper
# Auto-configured by langprotect-gateway-setup

# ============================================================
# Backend Connection
# ============================================================
export LANGPROTECT_URL="${{LANGPROTECT_URL:-{url}}}"
export LANGPROTECT_EMAIL="${{LANGPROTECT_EMAIL:-{email}}}"
export LANGPROTECT_PASSWORD="${{LANGPROTECT_PASSWORD:-{password}}}"

# ============================================================
# Security Controls (v1.3.1+)
# ============================================================
export LANGPROTECT_ENABLE_MASKING="${{LANGPROTECT_ENABLE_MASKING:-true}}"
export LANGPROTECT_FAIL_CLOSED="${{LANGPROTECT_FAIL_CLOSED:-false}}"
export LANGPROTECT_SCAN_TIMEOUT="${{LANGPROTECT_SCAN_TIMEOUT:-5.0}}"
export LANGPROTECT_ENTROPY_DETECTION="${{LANGPROTECT_ENTROPY_DETECTION:-true}}"

# ============================================================
# MCP Server Configuration
# ============================================================
export MCP_SERVER_COMMAND="${{MCP_SERVER_COMMAND:-npx}}"
export MCP_SERVER_ARGS="${{MCP_SERVER_ARGS:--y,@modelcontextprotocol/server-filesystem,.}}"

# Start the gateway
exec langprotect-gateway "$@"
"""
    
    wrapper_path.write_text(wrapper_content)
    wrapper_path.chmod(0o755)
    
    return wrapper_path


def update_vscode_settings(wrapper_path):
    """Update VS Code settings to use the wrapper"""
    settings_path = get_vscode_settings_path()
    
    # Create directory if it doesn't exist
    settings_path.parent.mkdir(parents=True, exist_ok=True)
    
    # Read existing settings or create new
    if settings_path.exists():
        with open(settings_path, 'r') as f:
            try:
                settings = json.load(f)
            except json.JSONDecodeError:
                settings = {}
    else:
        settings = {}
    
    # Add MCP configuration
    if "chat.mcp.servers" not in settings:
        settings["chat.mcp.servers"] = {}
    
    settings["chat.mcp.servers"]["langprotect-gateway"] = {
        "type": "stdio",
        "command": str(wrapper_path),
        "args": []
    }
    
    # Enable auto-start
    settings["chat.mcp.autostart"] = "newAndOutdated"
    
    # Write back
    with open(settings_path, 'w') as f:
        json.dump(settings, f, indent=2)
    
    return settings_path


def get_claude_config_path():
    """Get the Claude Desktop config path based on OS"""
    home = Path.home()
    if sys.platform == "darwin":
        return home / "Library/Application Support/Claude/claude_desktop_config.json"
    elif sys.platform == "win32":
        return home / "AppData/Roaming/Claude/claude_desktop_config.json"
    else:
        return home / ".config/Claude/claude_desktop_config.json"


def update_claude_config(wrapper_path):
    """Update Claude Desktop config to use the wrapper"""
    config_path = get_claude_config_path()
    config_path.parent.mkdir(parents=True, exist_ok=True)
    
    if config_path.exists():
        with open(config_path, 'r') as f:
            try:
                config = json.load(f)
            except json.JSONDecodeError:
                config = {}
    else:
        config = {}
    
    if "mcpServers" not in config:
        config["mcpServers"] = {}
    
    config["mcpServers"]["langprotect-gateway"] = {
        "command": str(wrapper_path),
        "args": []
    }
    
    with open(config_path, 'w') as f:
        json.dump(config, f, indent=2)
    
    return config_path


def setup():
    """Main setup function"""
    print("🚀 Setting up LangProtect MCP Gateway...")
    print()
    
    # Create wrapper script (will prompt for credentials if needed)
    print("📝 Creating global wrapper script...")
    wrapper_path = create_wrapper_script()
    print(f"   ✅ Created: {wrapper_path}")
    print()
    
    # Update VS Code settings
    print("⚙️  Configuring VS Code...")
    try:
        settings_path = update_vscode_settings(wrapper_path)
        print(f"   ✅ Updated: {settings_path}")
    except Exception as e:
        print(f"   ⚠️  Could not update VS Code settings: {e}")
    
    # Update Claude Desktop config
    print("🍏 Configuring Claude Desktop...")
    try:
        claude_path = update_claude_config(wrapper_path)
        print(f"   ✅ Updated: {claude_path}")
    except Exception as e:
        print(f"   ⚠️  Could not update Claude Desktop config: {e}")
    print()
    
    # Print next steps
    print("✅ Setup complete!")
    print()
    print("📋 Next steps:")
    print()
    print("1. Reload VS Code:")
    print("   Press Ctrl+Shift+P → 'Developer: Reload Window'")
    print()
    print("2. Verify it's working:")
    print("   Press Ctrl+Shift+P → 'MCP: List Servers'")
    print("   You should see 'langprotect-gateway' listed")
    print()
    print("3. Test the protection:")
    print("   Ask AI to read a file with secrets - they'll be masked!")
    print()
    print("🎉 LangProtect is now protecting ALL your VS Code workspaces!")
    print()
    print("💡 Configuration file:", wrapper_path)
    print()


if __name__ == "__main__":
    try:
        setup()
    except KeyboardInterrupt:
        print("\n\n⚠️  Setup cancelled by user.")
        sys.exit(1)
    except Exception as e:
        print(f"\n\n❌ Setup failed: {e}")
        sys.exit(1)
