#!/usr/bin/env python3
"""
Claude Code MCP Server Toggle Utility
Enables/disables MCP servers in ~/.claude.json for specific projects
"""

import json
import sys
import os
from pathlib import Path

CLAUDE_CONFIG = Path.home() / ".claude.json"

def load_config():
    if not CLAUDE_CONFIG.exists():
        print(f"Error: {CLAUDE_CONFIG} not found")
        sys.exit(1)

    with open(CLAUDE_CONFIG, 'r') as f:
        return json.load(f)

def save_config(data):
    with open(CLAUDE_CONFIG, 'w') as f:
        json.dump(data, f, indent=2, ensure_ascii=False)

def list_servers(project_path=None):
    data = load_config()

    if project_path:
        project_path = str(Path(project_path).resolve())
        if project_path not in data.get('projects', {}):
            print(f"Project not found: {project_path}")
            return

        project = data['projects'][project_path]
        enabled_servers = project.get('mcpServers', {})
        disabled_servers = project.get('_disabledServers', {})

        print(f"\n📁 Project: {project_path}")
        print(f"{'='*60}")

        if not enabled_servers and not disabled_servers:
            print("No MCP servers configured for this project")
            return

        if enabled_servers:
            print("\n✅ Enabled Servers:")
            for name, config in enabled_servers.items():
                print(f"  • {name}")
                print(f"    Command: {config.get('command', 'N/A')}")
                print(f"    Args: {config.get('args', [])}")

        if disabled_servers:
            print("\n❌ Disabled Servers:")
            for name, config in disabled_servers.items():
                print(f"  • {name}")
                print(f"    Command: {config.get('command', 'N/A')}")
                print(f"    Args: {config.get('args', [])}")
    else:
        print("\n🔧 All Projects with MCP Servers:")
        print(f"{'='*60}")

        for proj_path, proj_data in data.get('projects', {}).items():
            enabled_servers = proj_data.get('mcpServers', {})
            disabled_servers = proj_data.get('_disabledServers', {})
            if enabled_servers or disabled_servers:
                print(f"\n📁 {proj_path}")
                if enabled_servers:
                    print("  ✅ Enabled:")
                    for name in enabled_servers.keys():
                        print(f"    - {name}")
                if disabled_servers:
                    print("  ❌ Disabled:")
                    for name in disabled_servers.keys():
                        print(f"    - {name}")

def disable_server(project_path, server_name):
    data = load_config()
    project_path = str(Path(project_path).resolve())

    if project_path not in data.get('projects', {}):
        print(f"Error: Project not found: {project_path}")
        sys.exit(1)

    project = data['projects'][project_path]

    if 'mcpServers' not in project or server_name not in project['mcpServers']:
        print(f"Error: Server '{server_name}' not found in project")
        sys.exit(1)

    # Store config in disabled state
    if '_disabledServers' not in project:
        project['_disabledServers'] = {}

    project['_disabledServers'][server_name] = project['mcpServers'][server_name]
    del project['mcpServers'][server_name]

    save_config(data)
    print(f"✅ Disabled MCP server: {server_name}")

def enable_server(project_path, server_name):
    data = load_config()
    project_path = str(Path(project_path).resolve())

    if project_path not in data.get('projects', {}):
        print(f"Error: Project not found: {project_path}")
        sys.exit(1)

    project = data['projects'][project_path]

    if '_disabledServers' not in project or server_name not in project['_disabledServers']:
        print(f"Error: Server '{server_name}' not found in disabled servers")
        sys.exit(1)

    # Restore config from disabled state
    if 'mcpServers' not in project:
        project['mcpServers'] = {}

    project['mcpServers'][server_name] = project['_disabledServers'][server_name]
    del project['_disabledServers'][server_name]

    save_config(data)
    print(f"✅ Enabled MCP server: {server_name}")

def disable_all_servers(project_path):
    data = load_config()
    project_path = str(Path(project_path).resolve())

    if project_path not in data.get('projects', {}):
        print(f"Error: Project not found: {project_path}")
        sys.exit(1)

    project = data['projects'][project_path]
    enabled_servers = project.get('mcpServers', {})

    if not enabled_servers:
        print("No enabled servers to disable")
        return

    # Move all to disabled
    if '_disabledServers' not in project:
        project['_disabledServers'] = {}

    for name, config in enabled_servers.items():
        project['_disabledServers'][name] = config

    project['mcpServers'] = {}

    save_config(data)
    print(f"✅ Disabled all {len(enabled_servers)} MCP servers")

def enable_all_servers(project_path):
    data = load_config()
    project_path = str(Path(project_path).resolve())

    if project_path not in data.get('projects', {}):
        print(f"Error: Project not found: {project_path}")
        sys.exit(1)

    project = data['projects'][project_path]
    disabled_servers = project.get('_disabledServers', {})

    if not disabled_servers:
        print("No disabled servers to enable")
        return

    # Move all to enabled
    if 'mcpServers' not in project:
        project['mcpServers'] = {}

    for name, config in disabled_servers.items():
        project['mcpServers'][name] = config

    project['_disabledServers'] = {}

    save_config(data)
    print(f"✅ Enabled all {len(disabled_servers)} MCP servers")

def set_servers(project_path, server_names):
    """Disable all servers, then enable only the specified ones"""
    data = load_config()
    project_path = str(Path(project_path).resolve())

    if project_path not in data.get('projects', {}):
        print(f"Error: Project not found: {project_path}")
        sys.exit(1)

    project = data['projects'][project_path]
    enabled_servers = project.get('mcpServers', {})
    disabled_servers = project.get('_disabledServers', {})

    # Collect all servers
    all_servers = {**enabled_servers, **disabled_servers}

    # Validate requested servers exist
    for name in server_names:
        if name not in all_servers:
            print(f"Error: Server '{name}' not found in project")
            print(f"Available servers: {', '.join(all_servers.keys())}")
            sys.exit(1)

    # Move all to disabled
    if '_disabledServers' not in project:
        project['_disabledServers'] = {}

    project['_disabledServers'].update(all_servers)
    project['mcpServers'] = {}

    # Enable only requested servers
    for name in server_names:
        project['mcpServers'][name] = project['_disabledServers'][name]
        del project['_disabledServers'][name]

    save_config(data)
    print(f"✅ Enabled {len(server_names)} servers: {', '.join(server_names)}")
    print(f"   Disabled {len(all_servers) - len(server_names)} other servers")

def show_usage():
    print("""
Usage: claude-mcp-toggle <command> [options]

Commands:
  list [project_path]                      List all MCP servers (optionally for specific project)
  disable <project_path> <server>          Disable MCP server for project
  enable <project_path> <server>           Enable MCP server for project
  disable-all <project_path>               Disable all MCP servers for project
  enable-all <project_path>                Enable all MCP servers for project
  set <project_path> <server1> <server2>   Disable all, then enable only specified servers

Examples:
  claude-mcp-toggle list
  claude-mcp-toggle list /mnt/e/proiecte/ROMFASTSQL
  claude-mcp-toggle disable . chrome-devtools
  claude-mcp-toggle enable . chrome-devtools
  claude-mcp-toggle disable-all .
  claude-mcp-toggle set . chrome-devtools filesystem
""")

def interactive_mode():
    """Interactive mode when no arguments provided"""
    cwd = os.getcwd()
    data = load_config()

    project_path = str(Path(cwd).resolve())
    if project_path not in data.get('projects', {}):
        print(f"❌ Current directory is not a Claude Code project: {project_path}")
        print("\nAvailable projects:")
        for proj in data.get('projects', {}).keys():
            print(f"  - {proj}")
        sys.exit(1)

    project = data['projects'][project_path]
    enabled_servers = project.get('mcpServers', {})
    disabled_servers = project.get('_disabledServers', {})

    if not enabled_servers and not disabled_servers:
        print(f"No MCP servers configured for project: {project_path}")
        sys.exit(0)

    print(f"\n📁 Current Project: {project_path}")
    print(f"{'='*60}\n")

    # Show current state
    if enabled_servers:
        print("✅ Enabled Servers:")
        for i, name in enumerate(enabled_servers.keys(), 1):
            print(f"  {i}. {name}")

    if disabled_servers:
        print("\n❌ Disabled Servers:")
        for i, name in enumerate(disabled_servers.keys(), len(enabled_servers) + 1):
            print(f"  {i}. {name}")

    # Interactive menu
    print("\n" + "="*60)
    print("Commands:")
    print("  [e]nable <number>  - Enable a disabled server")
    print("  [d]isable <number> - Disable an enabled server")
    print("  [l]ist            - Show all servers with details")
    print("  [q]uit            - Exit")
    print("="*60)

    try:
        choice = input("\nEnter command: ").strip().lower()

        if choice == 'q' or choice == 'quit':
            print("Goodbye!")
            sys.exit(0)
        elif choice == 'l' or choice == 'list':
            list_servers(project_path)
        elif choice.startswith('e') or choice.startswith('d'):
            parts = choice.split()
            if len(parts) < 2:
                print("❌ Please specify server number")
                sys.exit(1)

            try:
                server_num = int(parts[1])
                all_servers = list(enabled_servers.keys()) + list(disabled_servers.keys())

                if server_num < 1 or server_num > len(all_servers):
                    print(f"❌ Invalid server number. Choose 1-{len(all_servers)}")
                    sys.exit(1)

                server_name = all_servers[server_num - 1]

                if choice.startswith('e'):
                    enable_server(project_path, server_name)
                else:
                    disable_server(project_path, server_name)
            except ValueError:
                print("❌ Invalid number")
                sys.exit(1)
        else:
            print("❌ Unknown command")
            sys.exit(1)
    except (EOFError, KeyboardInterrupt):
        print("\n\nGoodbye!")
        sys.exit(0)

def main():
    if len(sys.argv) < 2:
        interactive_mode()
        return

    command = sys.argv[1]

    if command == "list":
        project_path = sys.argv[2] if len(sys.argv) > 2 else None
        list_servers(project_path)
    elif command == "disable":
        if len(sys.argv) < 4:
            print("Error: Missing arguments")
            show_usage()
            sys.exit(1)
        disable_server(sys.argv[2], sys.argv[3])
    elif command == "enable":
        if len(sys.argv) < 4:
            print("Error: Missing arguments")
            show_usage()
            sys.exit(1)
        enable_server(sys.argv[2], sys.argv[3])
    elif command == "disable-all":
        if len(sys.argv) < 3:
            print("Error: Missing project path")
            show_usage()
            sys.exit(1)
        disable_all_servers(sys.argv[2])
    elif command == "enable-all":
        if len(sys.argv) < 3:
            print("Error: Missing project path")
            show_usage()
            sys.exit(1)
        enable_all_servers(sys.argv[2])
    elif command == "set":
        if len(sys.argv) < 4:
            print("Error: Missing arguments (project_path and at least one server)")
            show_usage()
            sys.exit(1)
        project_path = sys.argv[2]
        server_names = sys.argv[3:]
        set_servers(project_path, server_names)
    else:
        print(f"Error: Unknown command: {command}")
        show_usage()
        sys.exit(1)

if __name__ == "__main__":
    main()