"""
Shell completion module for MiniMax CLI

Generates completion scripts for bash, zsh, and fish shells with dynamic
completion based on available commands, arguments, and file paths.
"""

import os
import sys
from typing import List, Dict, Optional, Set
from pathlib import Path


class CompletionGenerator:
    """Generates shell completion scripts for the MiniMax CLI."""
    
    def __init__(self):
        self.commands = {
            'analyze': {
                'subcommands': ['review', 'explain', 'debug', 'optimize'],
                'options': ['--focus', '--function', '--lines', '--error', '--context', '--target'],
                'description': 'Analyze code files and projects'
            },
            'generate': {
                'subcommands': ['function', 'class', 'test', 'project'],
                'options': ['--language', '--output', '--tests', '--framework', '--type'],
                'description': 'Generate new code'
            },
            'edit': {
                'subcommands': ['modify', 'refactor', 'format'],
                'options': ['--backup', '--preview', '--type', '--target', '--style'],
                'description': 'Edit and modify existing files'
            },
            'chat': {
                'subcommands': ['start', 'context'],
                'options': ['--context', '--project', '--file'],
                'description': 'Interactive coding assistance'
            },
            'project': {
                'subcommands': ['init', 'scan', 'summary'],
                'options': ['--directory', '--depth', '--include-files'],
                'description': 'Project-level operations'
            },
            'complete': {
                'subcommands': [],
                'options': ['--stream', '--max-tokens', '--temperature'],
                'description': 'Simple text completion'
            }
        }
        
        self.global_options = ['--help', '--verbose', '--model', '--perf', '--profile']
        
        # File extensions for different contexts
        self.code_extensions = ['.py', '.js', '.ts', '.jsx', '.tsx', '.java', '.cpp', '.c', '.h']
        self.config_extensions = ['.json', '.yaml', '.yml', '.toml', '.ini']
        self.all_extensions = self.code_extensions + self.config_extensions + ['.txt', '.md']
        
    def generate_bash_completion(self) -> str:
        """Generate bash completion script."""
        script = '''#!/bin/bash
# Bash completion for MiniMax CLI
# Install: source this file or copy to /etc/bash_completion.d/

_minimax_completion() {
    local cur prev words cword
    _init_completion || return

    local commands="analyze generate edit chat project complete"
    local global_opts="--help --verbose --model --perf --profile"

    # Handle command completion
    if [[ $cword -eq 1 ]]; then
        COMPREPLY=($(compgen -W "$commands $global_opts" -- "$cur"))
        return 0
    fi

    local command="${words[1]}"
    
    case "$command" in
        analyze)
            if [[ $cword -eq 2 ]]; then
                COMPREPLY=($(compgen -W "review explain debug optimize" -- "$cur"))
            elif [[ $cword -ge 3 ]]; then
                case "${words[2]}" in
                    review|explain|debug|optimize)
                        if [[ "$cur" == -* ]]; then
                            local opts="--focus --function --lines --error --context --target --help"
                            COMPREPLY=($(compgen -W "$opts" -- "$cur"))
                        else
                            # Complete with files
                            _filedir
                        fi
                        ;;
                esac
            fi
            ;;
        generate)
            if [[ $cword -eq 2 ]]; then
                COMPREPLY=($(compgen -W "function class test project" -- "$cur"))
            elif [[ $cword -ge 3 ]]; then
                case "${words[2]}" in
                    function|class)
                        if [[ "$cur" == -* ]]; then
                            local opts="--language --output --tests --help"
                            COMPREPLY=($(compgen -W "$opts" -- "$cur"))
                        fi
                        ;;
                    test)
                        if [[ "$cur" == -* ]]; then
                            local opts="--framework --output --help"
                            COMPREPLY=($(compgen -W "$opts" -- "$cur"))
                        else
                            _filedir
                        fi
                        ;;
                    project)
                        if [[ "$cur" == -* ]]; then
                            local opts="--type --language --help"
                            COMPREPLY=($(compgen -W "$opts" -- "$cur"))
                        fi
                        ;;
                esac
            fi
            ;;
        edit)
            if [[ $cword -eq 2 ]]; then
                COMPREPLY=($(compgen -W "modify refactor format" -- "$cur"))
            elif [[ $cword -ge 3 ]]; then
                case "${words[2]}" in
                    modify|refactor|format)
                        if [[ "$cur" == -* ]]; then
                            local opts="--backup --preview --type --target --style --help"
                            COMPREPLY=($(compgen -W "$opts" -- "$cur"))
                        else
                            _filedir
                        fi
                        ;;
                esac
            fi
            ;;
        chat)
            if [[ $cword -eq 2 ]]; then
                COMPREPLY=($(compgen -W "start context" -- "$cur"))
            elif [[ $cword -ge 3 ]]; then
                case "${words[2]}" in
                    start)
                        if [[ "$cur" == -* ]]; then
                            local opts="--context --project --help"
                            COMPREPLY=($(compgen -W "$opts" -- "$cur"))
                        else
                            _filedir -d  # Directories for --project
                        fi
                        ;;
                    context)
                        if [[ "$cur" == -* ]]; then
                            local opts="--file --help"
                            COMPREPLY=($(compgen -W "$opts" -- "$cur"))
                        else
                            COMPREPLY=($(compgen -W "add remove list clear" -- "$cur"))
                        fi
                        ;;
                esac
            fi
            ;;
        project)
            if [[ $cword -eq 2 ]]; then
                COMPREPLY=($(compgen -W "init scan summary" -- "$cur"))
            elif [[ $cword -ge 3 ]]; then
                if [[ "$cur" == -* ]]; then
                    local opts="--directory --depth --include-files --help"
                    COMPREPLY=($(compgen -W "$opts" -- "$cur"))
                else
                    _filedir -d  # Directories
                fi
            fi
            ;;
        complete)
            if [[ "$cur" == -* ]]; then
                local opts="--stream --max-tokens --temperature --help"
                COMPREPLY=($(compgen -W "$opts" -- "$cur"))
            fi
            ;;
    esac
}

complete -F _minimax_completion minimax
'''
        return script
        
    def generate_zsh_completion(self) -> str:
        """Generate zsh completion script."""
        script = '''#compdef minimax
# Zsh completion for MiniMax CLI
# Install: copy to a directory in your $fpath

_minimax() {
    local context state line
    typeset -A opt_args

    _arguments -C \\
        '(--help -h)'{--help,-h}'[Show help]' \\
        '--verbose[Enable verbose output]' \\
        '--model[Override default model]:model:' \\
        '--perf[Enable performance monitoring]' \\
        '--profile[Use configuration profile]:profile:' \\
        '1: :_minimax_commands' \\
        '*::arg:->args' && return 0

    case $state in
        args)
            case $words[1] in
                analyze)
                    _arguments \\
                        '1: :(review explain debug optimize)' \\
                        '--focus[Focus area]:focus:(security performance style bugs)' \\
                        '--function[Specific function to analyze]:function:' \\
                        '--lines[Line range]:lines:' \\
                        '--error[Error message]:error:' \\
                        '--context[Additional context]:context:' \\
                        '--target[Optimization target]:target:(speed memory readability)' \\
                        '*:file:_files'
                    ;;
                generate)
                    _arguments \\
                        '1: :(function class test project)' \\
                        '--language[Programming language]:language:(python javascript typescript java cpp go rust)' \\
                        '--output[Output file]:output:_files' \\
                        '--tests[Generate tests]' \\
                        '--framework[Test framework]:framework:' \\
                        '--type[Project type]:type:(web api cli library)' \\
                        '*::arg:_files'
                    ;;
                edit)
                    _arguments \\
                        '1: :(modify refactor format)' \\
                        '--backup[Create backup]' \\
                        '--preview[Preview changes only]' \\
                        '--type[Refactoring type]:type:(extract inline rename move)' \\
                        '--target[Target element]:target:' \\
                        '--style[Code style]:style:' \\
                        '*:file:_files'
                    ;;
                chat)
                    _arguments \\
                        '1: :(start context)' \\
                        '--context[Initial context file]:context:_files' \\
                        '--project[Project directory]:project:_directories' \\
                        '--file[File to add/remove]:file:_files'
                    ;;
                project)
                    _arguments \\
                        '1: :(init scan summary)' \\
                        '--directory[Project directory]:directory:_directories' \\
                        '--depth[Scan depth]:depth:' \\
                        '--include-files[Include file contents]'
                    ;;
                complete)
                    _arguments \\
                        '--stream[Stream response]' \\
                        '--max-tokens[Maximum tokens]:tokens:' \\
                        '--temperature[Response temperature]:temperature:'
                    ;;
            esac
            ;;
    esac
}

_minimax_commands() {
    local commands=(
        'analyze:Analyze code files and projects'
        'generate:Generate new code'
        'edit:Edit and modify existing files'
        'chat:Interactive coding assistance'
        'project:Project-level operations'
        'complete:Simple text completion'
    )
    _describe 'commands' commands
}

_minimax "$@"
'''
        return script
        
    def generate_fish_completion(self) -> str:
        """Generate fish completion script."""
        script = '''# Fish completion for MiniMax CLI
# Install: copy to ~/.config/fish/completions/minimax.fish

# Remove old completions
complete -c minimax -e

# Global options
complete -c minimax -s h -l help -d "Show help"
complete -c minimax -l verbose -d "Enable verbose output"
complete -c minimax -l model -d "Override default model" -r
complete -c minimax -l perf -d "Enable performance monitoring"
complete -c minimax -l profile -d "Use configuration profile" -r

# Main commands
complete -c minimax -f -n "__fish_use_subcommand" -a "analyze" -d "Analyze code files and projects"
complete -c minimax -f -n "__fish_use_subcommand" -a "generate" -d "Generate new code"
complete -c minimax -f -n "__fish_use_subcommand" -a "edit" -d "Edit and modify existing files"
complete -c minimax -f -n "__fish_use_subcommand" -a "chat" -d "Interactive coding assistance"
complete -c minimax -f -n "__fish_use_subcommand" -a "project" -d "Project-level operations"
complete -c minimax -f -n "__fish_use_subcommand" -a "complete" -d "Simple text completion"

# Analyze subcommands
complete -c minimax -f -n "__fish_seen_subcommand_from analyze" -a "review" -d "Review code for issues"
complete -c minimax -f -n "__fish_seen_subcommand_from analyze" -a "explain" -d "Explain how code works"
complete -c minimax -f -n "__fish_seen_subcommand_from analyze" -a "debug" -d "Help debug issues"
complete -c minimax -f -n "__fish_seen_subcommand_from analyze" -a "optimize" -d "Suggest optimizations"

# Analyze options
complete -c minimax -n "__fish_seen_subcommand_from analyze" -l focus -d "Focus area" -r -a "security performance style bugs"
complete -c minimax -n "__fish_seen_subcommand_from analyze" -l function -d "Specific function" -r
complete -c minimax -n "__fish_seen_subcommand_from analyze" -l lines -d "Line range" -r
complete -c minimax -n "__fish_seen_subcommand_from analyze" -l error -d "Error message" -r
complete -c minimax -n "__fish_seen_subcommand_from analyze" -l context -d "Additional context" -r
complete -c minimax -n "__fish_seen_subcommand_from analyze" -l target -d "Optimization target" -r -a "speed memory readability"

# Generate subcommands
complete -c minimax -f -n "__fish_seen_subcommand_from generate" -a "function" -d "Generate a function"
complete -c minimax -f -n "__fish_seen_subcommand_from generate" -a "class" -d "Generate a class"
complete -c minimax -f -n "__fish_seen_subcommand_from generate" -a "test" -d "Generate tests"
complete -c minimax -f -n "__fish_seen_subcommand_from generate" -a "project" -d "Scaffold project"

# Generate options
complete -c minimax -n "__fish_seen_subcommand_from generate" -s l -l language -d "Programming language" -r -a "python javascript typescript java cpp go rust"
complete -c minimax -n "__fish_seen_subcommand_from generate" -s o -l output -d "Output file" -r -F
complete -c minimax -n "__fish_seen_subcommand_from generate" -l tests -d "Generate tests"
complete -c minimax -n "__fish_seen_subcommand_from generate" -l framework -d "Test framework" -r
complete -c minimax -n "__fish_seen_subcommand_from generate" -l type -d "Project type" -r -a "web api cli library"

# Edit subcommands
complete -c minimax -f -n "__fish_seen_subcommand_from edit" -a "modify" -d "Modify existing code"
complete -c minimax -f -n "__fish_seen_subcommand_from edit" -a "refactor" -d "Refactor code structure"
complete -c minimax -f -n "__fish_seen_subcommand_from edit" -a "format" -d "Format and style code"

# Edit options
complete -c minimax -n "__fish_seen_subcommand_from edit" -l backup -d "Create backup"
complete -c minimax -n "__fish_seen_subcommand_from edit" -l preview -d "Preview changes only"
complete -c minimax -n "__fish_seen_subcommand_from edit" -l type -d "Refactoring type" -r -a "extract inline rename move"
complete -c minimax -n "__fish_seen_subcommand_from edit" -l target -d "Target element" -r
complete -c minimax -n "__fish_seen_subcommand_from edit" -l style -d "Code style" -r

# Chat subcommands
complete -c minimax -f -n "__fish_seen_subcommand_from chat" -a "start" -d "Start interactive session"
complete -c minimax -f -n "__fish_seen_subcommand_from chat" -a "context" -d "Manage conversation context"

# Chat options
complete -c minimax -n "__fish_seen_subcommand_from chat" -l context -d "Initial context file" -r -F
complete -c minimax -n "__fish_seen_subcommand_from chat" -l project -d "Project directory" -r -a "(__fish_complete_directories)"
complete -c minimax -n "__fish_seen_subcommand_from chat" -l file -d "File to add/remove" -r -F

# Project subcommands
complete -c minimax -f -n "__fish_seen_subcommand_from project" -a "init" -d "Initialize project"
complete -c minimax -f -n "__fish_seen_subcommand_from project" -a "scan" -d "Scan project structure"
complete -c minimax -f -n "__fish_seen_subcommand_from project" -a "summary" -d "Generate project summary"

# Project options
complete -c minimax -n "__fish_seen_subcommand_from project" -s d -l directory -d "Project directory" -r -a "(__fish_complete_directories)"
complete -c minimax -n "__fish_seen_subcommand_from project" -l depth -d "Scan depth" -r
complete -c minimax -n "__fish_seen_subcommand_from project" -l include-files -d "Include file contents"

# Complete subcommand options
complete -c minimax -n "__fish_seen_subcommand_from complete" -l stream -d "Stream response"
complete -c minimax -n "__fish_seen_subcommand_from complete" -l max-tokens -d "Maximum tokens" -r
complete -c minimax -n "__fish_seen_subcommand_from complete" -l temperature -d "Response temperature" -r
'''
        return script
        
    def write_completion_script(self, shell: str, output_path: Optional[str] = None) -> bool:
        """
        Write completion script for specified shell.
        
        Args:
            shell: Shell type (bash, zsh, fish)
            output_path: Optional output path, defaults to appropriate location
            
        Returns:
            True if successful
        """
        try:
            if shell == 'bash':
                script = self.generate_bash_completion()
                default_path = Path.home() / '.minimax_completion.bash'
            elif shell == 'zsh':
                script = self.generate_zsh_completion()
                default_path = Path.home() / '.minimax_completion.zsh'
            elif shell == 'fish':
                script = self.generate_fish_completion()
                fish_config = Path.home() / '.config' / 'fish' / 'completions'
                fish_config.mkdir(parents=True, exist_ok=True)
                default_path = fish_config / 'minimax.fish'
            else:
                raise ValueError(f"Unsupported shell: {shell}")
                
            output = Path(output_path) if output_path else default_path
            
            with open(output, 'w', encoding='utf-8') as f:
                f.write(script)
                
            # Make executable for bash/zsh
            if shell in ['bash', 'zsh']:
                output.chmod(0o755)
                
            return True
            
        except Exception as e:
            print(f"Error writing completion script: {e}")
            return False
            
    def install_completion(self, shell: str) -> bool:
        """
        Install completion for the specified shell.
        
        Args:
            shell: Shell type to install for
            
        Returns:
            True if successful
        """
        try:
            if shell == 'bash':
                script_path = Path.home() / '.minimax_completion.bash'
                if self.write_completion_script('bash', str(script_path)):
                    print(f"Bash completion installed to {script_path}")
                    print("To enable, add this to your ~/.bashrc:")
                    print(f"source {script_path}")
                    return True
                    
            elif shell == 'zsh':
                script_path = Path.home() / '.minimax_completion.zsh'
                if self.write_completion_script('zsh', str(script_path)):
                    print(f"Zsh completion installed to {script_path}")
                    print("To enable, add this to your ~/.zshrc:")
                    print(f"source {script_path}")
                    return True
                    
            elif shell == 'fish':
                if self.write_completion_script('fish'):
                    print("Fish completion installed to ~/.config/fish/completions/minimax.fish")
                    print("Completion will be available in new fish sessions")
                    return True
                    
            else:
                print(f"Unsupported shell: {shell}")
                return False
                
        except Exception as e:
            print(f"Error installing completion: {e}")
            return False
            
        return False
        
    def generate_dynamic_completion(self, command_line: List[str]) -> List[str]:
        """
        Generate dynamic completion suggestions based on current command line.
        
        Args:
            command_line: Current command line tokens
            
        Returns:
            List of completion suggestions
        """
        if not command_line:
            return list(self.commands.keys()) + self.global_options
            
        if len(command_line) == 1:
            cmd = command_line[0]
            if cmd.startswith('-'):
                return [opt for opt in self.global_options if opt.startswith(cmd)]
            else:
                return [cmd for cmd in self.commands.keys() if cmd.startswith(cmd)]
                
        main_cmd = command_line[0]
        if main_cmd not in self.commands:
            return []
            
        cmd_info = self.commands[main_cmd]
        
        if len(command_line) == 2:
            subcmd = command_line[1]
            if subcmd.startswith('-'):
                return [opt for opt in cmd_info['options'] if opt.startswith(subcmd)]
            else:
                return [sub for sub in cmd_info['subcommands'] if sub.startswith(subcmd)]
                
        # Context-aware file completion
        current_token = command_line[-1]
        if current_token.startswith('-'):
            return [opt for opt in cmd_info['options'] if opt.startswith(current_token)]
        else:
            # Return file suggestions based on context
            return self._get_file_suggestions(command_line)
            
    def _get_file_suggestions(self, command_line: List[str]) -> List[str]:
        """Get file suggestions based on command context."""
        try:
            cwd = Path.cwd()
            files = []
            
            # Determine appropriate file extensions based on command
            main_cmd = command_line[0]
            if main_cmd in ['analyze', 'edit']:
                extensions = self.code_extensions
            elif main_cmd == 'generate' and 'test' in command_line:
                extensions = self.code_extensions
            else:
                extensions = self.all_extensions
                
            for item in cwd.iterdir():
                if item.is_file() and item.suffix in extensions:
                    files.append(str(item.name))
                elif item.is_dir() and not item.name.startswith('.'):
                    files.append(str(item.name) + '/')
                    
            return sorted(files)
            
        except Exception:
            return []


def main():
    """Main function for completion management."""
    import argparse
    
    parser = argparse.ArgumentParser(description='MiniMax CLI completion management')
    parser.add_argument('action', choices=['install', 'generate', 'complete'])
    parser.add_argument('--shell', choices=['bash', 'zsh', 'fish'], 
                       default=os.environ.get('SHELL', '').split('/')[-1])
    parser.add_argument('--output', help='Output file path')
    parser.add_argument('--command-line', nargs='*', help='Current command line for completion')
    
    args = parser.parse_args()
    
    generator = CompletionGenerator()
    
    if args.action == 'install':
        success = generator.install_completion(args.shell)
        sys.exit(0 if success else 1)
        
    elif args.action == 'generate':
        success = generator.write_completion_script(args.shell, args.output)
        sys.exit(0 if success else 1)
        
    elif args.action == 'complete':
        if args.command_line:
            suggestions = generator.generate_dynamic_completion(args.command_line)
            for suggestion in suggestions:
                print(suggestion)
        sys.exit(0)


if __name__ == '__main__':
    main()