#!/usr/bin/env python3
"""
Trixy Voice Assistant Setup Script

This script provides automated installation and setup for the Trixy Voice Assistant
system across different deployment modes and platforms.

Usage:
    python setup.py --mode [client|server|standalone|dev|ml] [options]

Examples:
    python setup.py --mode client                    # Minimal client installation
    python setup.py --mode server --with-gpu         # Server with GPU support
    python setup.py --mode dev --install-hooks       # Development environment
    python setup.py --mode ml --cuda-version 11.8    # ML training with CUDA
"""

import os
import sys
import argparse
import subprocess
import platform
import json
from pathlib import Path
from typing import List, Dict, Any, Optional


class TrixySetup:
    """Automated setup for Trixy Voice Assistant."""
    
    def __init__(self):
        self.platform = platform.system().lower()
        self.architecture = platform.machine().lower()
        self.python_version = sys.version_info
        self.cwd = Path.cwd()
        
        # Requirements file mapping
        self.requirements_map = {
            'client': 'requirements-client.txt',
            'server': 'requirements-server.txt',
            'standalone': 'requirements-server.txt',  # Standalone uses server requirements
            'dev': 'requirements-dev.txt',
            'ml': 'requirements-ml.txt',
            'base': 'requirements.txt',
            'optional': 'requirements-optional.txt'
        }
    
    def check_python_version(self) -> bool:
        """Check if Python version is compatible."""
        if self.python_version < (3, 8):
            print(f"❌ Error: Python 3.8+ required, found {self.python_version.major}.{self.python_version.minor}")
            return False
        
        print(f"✅ Python {self.python_version.major}.{self.python_version.minor}.{self.python_version.micro} detected")
        return True
    
    def check_system_dependencies(self, mode: str) -> bool:
        """Check and install system dependencies based on platform and mode."""
        print(f"🔍 Checking system dependencies for {mode} mode on {self.platform}...")
        
        if self.platform == 'linux':
            return self._check_linux_dependencies(mode)
        elif self.platform == 'darwin':
            return self._check_macos_dependencies(mode)
        elif self.platform == 'windows':
            return self._check_windows_dependencies(mode)
        else:
            print(f"⚠️  Platform {self.platform} not fully supported")
            return True
    
    def _check_linux_dependencies(self, mode: str) -> bool:
        """Check Linux system dependencies."""
        required_packages = ['python3-pip', 'python3-venv']
        
        if mode in ['client', 'server', 'standalone']:
            required_packages.extend([
                'portaudio19-dev',
                'libasound2-dev',
                'ffmpeg',
                'libavcodec-extra'
            ])
        
        if mode in ['server', 'standalone']:
            required_packages.append('redis-server')
        
        if mode in ['dev', 'ml']:
            required_packages.extend([
                'build-essential',
                'cmake',
                'git'
            ])
        
        print(f"📦 Required packages: {', '.join(required_packages)}")
        print("💡 Install with: sudo apt-get install " + " ".join(required_packages))
        return True
    
    def _check_macos_dependencies(self, mode: str) -> bool:
        """Check macOS system dependencies."""
        required_packages = []
        
        if mode in ['client', 'server', 'standalone']:
            required_packages.extend(['portaudio', 'ffmpeg'])
        
        if mode in ['server', 'standalone']:
            required_packages.append('redis')
        
        if required_packages:
            print(f"📦 Required Homebrew packages: {', '.join(required_packages)}")
            print("💡 Install with: brew install " + " ".join(required_packages))
        
        return True
    
    def _check_windows_dependencies(self, mode: str) -> bool:
        """Check Windows system dependencies."""
        print("📦 Windows dependencies:")
        print("- Visual C++ Build Tools (for some packages)")
        print("- FFmpeg (add to PATH)")
        if mode in ['server', 'standalone']:
            print("- Redis (optional, can use Redis Cloud)")
        return True
    
    def create_virtual_environment(self, venv_name: str = "trixy-env") -> bool:
        """Create a Python virtual environment."""
        venv_path = self.cwd / venv_name
        
        if venv_path.exists():
            print(f"✅ Virtual environment {venv_name} already exists")
            return True
        
        try:
            print(f"🔨 Creating virtual environment: {venv_name}")
            subprocess.run([sys.executable, '-m', 'venv', str(venv_path)], check=True)
            print(f"✅ Virtual environment created: {venv_path}")
            return True
        except subprocess.CalledProcessError as e:
            print(f"❌ Failed to create virtual environment: {e}")
            return False
    
    def get_pip_command(self, venv_name: str = "trixy-env") -> List[str]:
        """Get the pip command for the virtual environment."""
        venv_path = self.cwd / venv_name
        
        if self.platform == 'windows':
            pip_path = venv_path / "Scripts" / "pip.exe"
        else:
            pip_path = venv_path / "bin" / "pip"
        
        return [str(pip_path)]
    
    def install_pytorch(self, cuda_version: Optional[str] = None, venv_name: str = "trixy-env") -> bool:
        """Install PyTorch with appropriate configuration."""
        pip_cmd = self.get_pip_command(venv_name)
        
        if cuda_version:
            print(f"🔥 Installing PyTorch with CUDA {cuda_version}")
            index_url = f"https://download.pytorch.org/whl/cu{cuda_version.replace('.', '')}"
            cmd = pip_cmd + ['install', 'torch', 'torchaudio', '--index-url', index_url]
        elif self.architecture in ['aarch64', 'arm64']:
            print("🔧 Installing PyTorch for ARM architecture")
            cmd = pip_cmd + ['install', 'torch', 'torchaudio', '--index-url', 
                           'https://download.pytorch.org/whl/cpu']
        else:
            print("🔧 Installing PyTorch CPU version")
            cmd = pip_cmd + ['install', 'torch', 'torchaudio', '--index-url', 
                           'https://download.pytorch.org/whl/cpu']
        
        try:
            subprocess.run(cmd, check=True)
            print("✅ PyTorch installed successfully")
            return True
        except subprocess.CalledProcessError as e:
            print(f"❌ Failed to install PyTorch: {e}")
            return False
    
    def install_requirements(self, mode: str, venv_name: str = "trixy-env") -> bool:
        """Install requirements for the specified mode."""
        requirements_file = self.requirements_map.get(mode)
        if not requirements_file:
            print(f"❌ Unknown mode: {mode}")
            return False
        
        requirements_path = self.cwd / requirements_file
        if not requirements_path.exists():
            print(f"❌ Requirements file not found: {requirements_path}")
            return False
        
        pip_cmd = self.get_pip_command(venv_name)
        
        try:
            print(f"📦 Installing requirements from {requirements_file}")
            # Upgrade pip first
            subprocess.run(pip_cmd + ['install', '--upgrade', 'pip'], check=True)
            
            # Install requirements
            subprocess.run(pip_cmd + ['install', '-r', str(requirements_path)], check=True)
            print(f"✅ Requirements installed successfully")
            return True
        except subprocess.CalledProcessError as e:
            print(f"❌ Failed to install requirements: {e}")
            return False
    
    def setup_development_tools(self, venv_name: str = "trixy-env") -> bool:
        """Set up development tools like pre-commit hooks."""
        try:
            # Get the python command from virtual environment
            venv_path = self.cwd / venv_name
            if self.platform == 'windows':
                python_cmd = str(venv_path / "Scripts" / "python.exe")
            else:
                python_cmd = str(venv_path / "bin" / "python")
            
            print("🔧 Setting up pre-commit hooks...")
            subprocess.run([python_cmd, '-m', 'pre_commit', 'install'], check=True)
            print("✅ Pre-commit hooks installed")
            return True
        except subprocess.CalledProcessError as e:
            print(f"⚠️  Failed to set up pre-commit hooks: {e}")
            return False
    
    def verify_installation(self, mode: str, venv_name: str = "trixy-env") -> bool:
        """Verify the installation is working correctly."""
        print("🔍 Verifying installation...")
        
        venv_path = self.cwd / venv_name
        if self.platform == 'windows':
            python_cmd = str(venv_path / "Scripts" / "python.exe")
        else:
            python_cmd = str(venv_path / "bin" / "python")
        
        # Test basic imports
        test_commands = [
            "import torch; print(f'PyTorch version: {torch.__version__}')",
            "import torchaudio; print(f'TorchAudio version: {torchaudio.__version__}')",
            "import numpy; print(f'NumPy version: {numpy.__version__}')",
        ]
        
        if mode in ['server', 'dev']:
            test_commands.append("import rich; print('Rich import successful')")
        
        for test_cmd in test_commands:
            try:
                result = subprocess.run([python_cmd, '-c', test_cmd], 
                                      capture_output=True, text=True, check=True)
                print(f"✅ {result.stdout.strip()}")
            except subprocess.CalledProcessError as e:
                print(f"❌ Test failed: {test_cmd}")
                print(f"   Error: {e.stderr}")
                return False
        
        # Test main application import
        try:
            subprocess.run([python_cmd, '-c', 
                          "print('Testing main application...'); import main; print('✅ Main application import successful')"], 
                         check=True)
        except subprocess.CalledProcessError as e:
            print("⚠️  Main application test failed (expected if core modules not implemented)")
        
        return True
    
    def create_config_files(self, mode: str) -> bool:
        """Create default configuration files."""
        config_dir = self.cwd / "config"
        config_dir.mkdir(exist_ok=True)
        
        configs = {
            'server': {
                'mode': 'server',
                'debug': False,
                'host': '0.0.0.0',
                'command_port': 2101,
                'audio_input_port': 2102,
                'audio_output_port': 2103,
                'music_output_port': 2104,
                'max_satellites': 10,
                'registration_timeout': 60
            },
            'client': {
                'mode': 'client',
                'debug': False,
                'server_host': 'localhost',
                'server_port': 2101,
                'room_id': 'default',
                'alias': 'client-device',
                'wakeword_model': 'models/wakeword/default.pth'
            },
            'standalone': {
                'mode': 'standalone',
                'debug': False,
                'enable_tui': True,
                'plugin_directory': 'plugins',
                'model_directory': 'models'
            }
        }
        
        if mode in configs:
            config_file = config_dir / f"{mode}_config.json"
            with open(config_file, 'w') as f:
                json.dump(configs[mode], f, indent=2)
            print(f"✅ Created configuration file: {config_file}")
        
        return True
    
    def print_next_steps(self, mode: str, venv_name: str = "trixy-env"):
        """Print next steps for the user."""
        print("\n" + "="*50)
        print("🎉 Installation Complete!")
        print("="*50)
        
        # Activation command
        venv_path = self.cwd / venv_name
        if self.platform == 'windows':
            activate_cmd = f"{venv_path}\\Scripts\\activate"
        else:
            activate_cmd = f"source {venv_path}/bin/activate"
        
        print(f"\n📝 Next Steps:")
        print(f"1. Activate the virtual environment:")
        print(f"   {activate_cmd}")
        print(f"\n2. Run Trixy in {mode} mode:")
        print(f"   python main.py {mode}")
        
        if mode == 'dev':
            print(f"\n3. Development commands:")
            print(f"   pytest              # Run tests")
            print(f"   black .             # Format code") 
            print(f"   mypy .              # Type checking")
            print(f"   pre-commit run --all-files  # Run all checks")
        
        print(f"\n📖 Documentation:")
        print(f"   - Installation guide: INSTALL.md")
        print(f"   - Project overview: CLAUDE.md")
        print(f"   - Configuration: config/{mode}_config.json")
        
        print(f"\n🔧 Configuration files created in: config/")
        print(f"💡 Edit configuration files before first run")


def main():
    """Main setup function."""
    parser = argparse.ArgumentParser(
        description="Trixy Voice Assistant Setup Script",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  python setup.py --mode client                    # Minimal client installation
  python setup.py --mode server --with-gpu         # Server with GPU support  
  python setup.py --mode dev --install-hooks       # Development environment
  python setup.py --mode ml --cuda-version 11.8    # ML training with CUDA
        """
    )
    
    parser.add_argument(
        '--mode',
        choices=['client', 'server', 'standalone', 'dev', 'ml'],
        required=True,
        help='Installation mode'
    )
    
    parser.add_argument(
        '--venv-name',
        default='trixy-env',
        help='Virtual environment name (default: trixy-env)'
    )
    
    parser.add_argument(
        '--cuda-version',
        help='CUDA version for GPU support (e.g., 11.8)'
    )
    
    parser.add_argument(
        '--with-gpu',
        action='store_true',
        help='Install GPU support (CUDA)'
    )
    
    parser.add_argument(
        '--install-hooks',
        action='store_true',
        help='Install pre-commit hooks (for dev mode)'
    )
    
    parser.add_argument(
        '--skip-system-check',
        action='store_true',
        help='Skip system dependency checks'
    )
    
    parser.add_argument(
        '--skip-pytorch',
        action='store_true',
        help='Skip PyTorch installation (assume already installed)'
    )
    
    args = parser.parse_args()
    
    # Initialize setup
    setup = TrixySetup()
    
    print("🚀 Trixy Voice Assistant Setup")
    print("="*50)
    print(f"Mode: {args.mode}")
    print(f"Platform: {setup.platform} ({setup.architecture})")
    print(f"Python: {setup.python_version.major}.{setup.python_version.minor}.{setup.python_version.micro}")
    print("="*50)
    
    # Check Python version
    if not setup.check_python_version():
        sys.exit(1)
    
    # Check system dependencies
    if not args.skip_system_check:
        if not setup.check_system_dependencies(args.mode):
            print("⚠️  System dependency check failed. Use --skip-system-check to continue anyway.")
            sys.exit(1)
    
    # Create virtual environment
    if not setup.create_virtual_environment(args.venv_name):
        sys.exit(1)
    
    # Install PyTorch
    if not args.skip_pytorch:
        cuda_version = args.cuda_version if args.with_gpu or args.cuda_version else None
        if not setup.install_pytorch(cuda_version, args.venv_name):
            sys.exit(1)
    
    # Install requirements
    if not setup.install_requirements(args.mode, args.venv_name):
        sys.exit(1)
    
    # Set up development tools
    if args.mode == 'dev' and args.install_hooks:
        setup.setup_development_tools(args.venv_name)
    
    # Create configuration files
    setup.create_config_files(args.mode)
    
    # Verify installation
    if not setup.verify_installation(args.mode, args.venv_name):
        print("⚠️  Installation verification had issues, but you can still try running the application")
    
    # Print next steps
    setup.print_next_steps(args.mode, args.venv_name)


if __name__ == "__main__":
    main()