"""
Interactive setup wizard for labmate-mcp.

Usage:
    labmate-mcp --setup
"""

import os
import sys
import json
from pathlib import Path

# ANSI colors
GREEN = "\033[92m"
YELLOW = "\033[93m"
CYAN = "\033[96m"
RESET = "\033[0m"
BOLD = "\033[1m"

ENV_FILE = Path.home() / ".labmate-mcp.env"

CLAUDE_DESKTOP_CONFIGS = [
    Path.home() / "Library" / "Application Support" / "Claude" / "claude_desktop_config.json",  # macOS
    Path(os.environ.get("APPDATA", "")) / "Claude" / "claude_desktop_config.json",  # Windows
    Path.home() / ".config" / "claude" / "claude_desktop_config.json",  # Linux
]

API_KEYS = [
    {
        "env_var": "RXN_API_KEY",
        "name": "IBM RXN",
        "url": "https://rxn.res.ibm.com",
        "description": "Retrosynthesis, forward prediction, atom mapping",
        "instructions": "Sign up → Profile → My API Keys → Create",
    },
    {
        "env_var": "ROWAN_API_KEY",
        "name": "Rowan Science",
        "url": "https://rowan.ai",
        "description": "pKa, solubility, ADMET, NMR prediction",
        "instructions": "Create account → Dashboard → API Key",
    },
    {
        "env_var": "SEMANTIC_SCHOLAR_API_KEY",
        "name": "Semantic Scholar",
        "url": "https://www.semanticscholar.org/product/api#api-key",
        "description": "Higher rate limits for literature search",
        "instructions": "Request key from the API page",
        "optional": True,
    },
    {
        "env_var": "UNPAYWALL_EMAIL",
        "name": "Unpaywall",
        "url": "https://unpaywall.org/products/api",
        "description": "Open access PDF discovery",
        "instructions": "Just use your email address (no signup needed)",
        "optional": True,
    },
]


def print_header():
    """Print the setup wizard header."""
    print()
    print(f"{BOLD}{CYAN}🧪 labmate-mcp setup{RESET}")
    print()
    print("This wizard will help you configure API keys for optional features.")
    print("All keys are stored locally in ~/.labmate-mcp.env")
    print()
    print(f"{YELLOW}Press Enter to skip any key you don't have yet.{RESET}")
    print()


def load_existing_keys() -> dict:
    """Load existing keys from .env file."""
    keys = {}
    if ENV_FILE.exists():
        with open(ENV_FILE) as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, value = line.split("=", 1)
                    keys[key.strip()] = value.strip()
    return keys


def prompt_for_key(api_info: dict, existing_value: str = "") -> str:
    """Prompt user for a single API key."""
    name = api_info["name"]
    env_var = api_info["env_var"]
    url = api_info["url"]
    desc = api_info["description"]
    instructions = api_info["instructions"]
    optional = api_info.get("optional", False)
    
    print(f"{BOLD}{name}{RESET} — {desc}")
    print(f"  Get it: {CYAN}{url}{RESET}")
    print(f"  How: {instructions}")
    
    if existing_value:
        masked = existing_value[:8] + "..." + existing_value[-4:] if len(existing_value) > 16 else "***"
        prompt = f"  {env_var} [{masked}]: "
    else:
        prompt = f"  {env_var}: "
    
    value = input(prompt).strip()
    
    if not value and existing_value:
        print(f"  {GREEN}✓ Keeping existing key{RESET}")
        return existing_value
    elif value:
        print(f"  {GREEN}✓ Saved{RESET}")
        return value
    else:
        print(f"  → Skipped")
        return ""


def save_env_file(keys: dict):
    """Save keys to .env file."""
    with open(ENV_FILE, "w") as f:
        f.write("# labmate-mcp API keys\n")
        f.write("# Generated by labmate-mcp --setup\n")
        f.write("# Edit manually or re-run setup to change\n\n")
        for key, value in keys.items():
            if value:
                f.write(f"{key}={value}\n")
    
    # Make file readable only by user
    os.chmod(ENV_FILE, 0o600)
    print()
    print(f"{GREEN}✓ Saved to {ENV_FILE}{RESET}")


def update_claude_config():
    """Optionally update Claude Desktop config to use labmate-mcp."""
    # Find existing config
    config_path = None
    for path in CLAUDE_DESKTOP_CONFIGS:
        if path.exists():
            config_path = path
            break
    
    if not config_path:
        # Try to create on macOS
        macos_path = CLAUDE_DESKTOP_CONFIGS[0]
        if sys.platform == "darwin":
            config_path = macos_path
        else:
            print(f"\n{YELLOW}Could not find Claude Desktop config.{RESET}")
            print("Add labmate-mcp manually to your config.")
            return
    
    print()
    update = input(f"Update Claude Desktop config at {config_path}? [Y/n]: ").strip().lower()
    if update == "n":
        return
    
    # Load or create config
    if config_path.exists():
        with open(config_path) as f:
            try:
                config = json.load(f)
            except json.JSONDecodeError:
                config = {}
    else:
        config_path.parent.mkdir(parents=True, exist_ok=True)
        config = {}
    
    # Add/update labmate entry
    if "mcpServers" not in config:
        config["mcpServers"] = {}
    
    config["mcpServers"]["labmate"] = {
        "command": "labmate-mcp"
    }
    
    # Write config
    with open(config_path, "w") as f:
        json.dump(config, f, indent=2)
    
    print(f"{GREEN}✓ Updated {config_path}{RESET}")


def print_summary(keys: dict):
    """Print summary of configured keys."""
    print()
    print(f"{BOLD}Summary:{RESET}")
    
    configured = [k for k, v in keys.items() if v]
    if configured:
        for key in configured:
            print(f"  {GREEN}✓{RESET} {key}")
    
    missing = [k for k, v in keys.items() if not v]
    if missing:
        for key in missing:
            print(f"  → {key} (not configured)")
    
    print()
    print(f"{BOLD}Restart Claude Desktop to apply changes.{RESET}")
    print()


def run_setup():
    """Run the interactive setup wizard."""
    print_header()
    
    # Load existing keys
    existing = load_existing_keys()
    new_keys = {}
    
    # Prompt for each key
    for i, api_info in enumerate(API_KEYS):
        env_var = api_info["env_var"]
        existing_value = existing.get(env_var, "")
        new_keys[env_var] = prompt_for_key(api_info, existing_value)
        if i < len(API_KEYS) - 1:
            print()
    
    # Save to file
    save_env_file(new_keys)
    
    # Update Claude config
    update_claude_config()
    
    # Print summary
    print_summary(new_keys)


if __name__ == "__main__":
    run_setup()
