#!/usr/bin/env python3
"""
Script to safely add zen-marketing MCP server to Claude Desktop configuration.
"""

import json
import os
import shutil
from pathlib import Path

def update_claude_config():
    """Add zen-marketing server to Claude Desktop config"""

    # Path to Claude config
    config_path = Path.home() / ".claude.json"

    if not config_path.exists():
        print(f"Error: Config file not found at {config_path}")
        return False

    # Create backup
    backup_path = config_path.with_suffix('.json.backup')
    shutil.copy2(config_path, backup_path)
    print(f"✓ Created backup: {backup_path}")

    # Read current config
    with open(config_path, 'r') as f:
        config = json.load(f)

    # Get MCP servers section
    if 'mcpServers' not in config:
        config['mcpServers'] = {}

    # Get absolute paths
    project_root = Path(__file__).parent.resolve()
    venv_python = project_root / ".venv" / "bin" / "python"
    server_script = project_root / "server.py"

    # Verify paths exist
    if not venv_python.exists():
        print(f"Error: Python venv not found at {venv_python}")
        print("Run ./run-server.sh first to create the virtual environment")
        return False

    if not server_script.exists():
        print(f"Error: server.py not found at {server_script}")
        return False

    # Read OpenRouter key from .env
    env_file = project_root / ".env"
    openrouter_key = ""
    gemini_key = ""

    if env_file.exists():
        with open(env_file, 'r') as f:
            for line in f:
                if line.startswith('OPENROUTER_API_KEY='):
                    openrouter_key = line.split('=', 1)[1].strip()
                elif line.startswith('GEMINI_API_KEY='):
                    gemini_key = line.split('=', 1)[1].strip()

    if not openrouter_key:
        print("Warning: OPENROUTER_API_KEY not found in .env")

    # Create zen-marketing server config
    zen_marketing_config = {
        "command": str(venv_python),
        "args": [str(server_script)],
        "env": {
            "DEFAULT_MODEL": "google/gemini-2.5-pro-latest",
            "FAST_MODEL": "google/gemini-2.5-flash-preview-09-2025",
            "CREATIVE_MODEL": "minimax/minimax-m2",
            "ENABLE_WEB_SEARCH": "true",
            "LOG_LEVEL": "INFO"
        }
    }

    # Add API keys if available
    if openrouter_key:
        zen_marketing_config["env"]["OPENROUTER_API_KEY"] = openrouter_key
    if gemini_key:
        zen_marketing_config["env"]["GEMINI_API_KEY"] = gemini_key

    # Check if zen-marketing already exists
    if 'zen-marketing' in config['mcpServers']:
        print("\n⚠ zen-marketing server already configured")
        response = input("Overwrite existing configuration? (y/N): ")
        if response.lower() != 'y':
            print("Cancelled.")
            return False

    # Add/update zen-marketing
    config['mcpServers']['zen-marketing'] = zen_marketing_config

    # Write updated config
    with open(config_path, 'w') as f:
        json.dump(config, f, indent=2)

    print(f"\n✓ Added zen-marketing MCP server to {config_path}")
    print(f"\nConfiguration:")
    print(f"  Command: {venv_python}")
    print(f"  Script:  {server_script}")
    print(f"  Models:")
    print(f"    - DEFAULT: {zen_marketing_config['env']['DEFAULT_MODEL']}")
    print(f"    - FAST:    {zen_marketing_config['env']['FAST_MODEL']}")
    print(f"    - CREATIVE: {zen_marketing_config['env']['CREATIVE_MODEL']}")
    if openrouter_key:
        print(f"  OpenRouter API Key: {'*' * (len(openrouter_key) - 8)}{openrouter_key[-8:]}")
    if gemini_key:
        print(f"  Gemini API Key: {'*' * (len(gemini_key) - 8)}{gemini_key[-8:]}")

    print(f"\n⚠  IMPORTANT: Restart Claude Desktop for changes to take effect")

    return True

if __name__ == "__main__":
    try:
        success = update_claude_config()
        if not success:
            exit(1)
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()
        exit(1)