ai-shell/ai_shell/main.py

"""
Main entry point for AI Shell
"""

import os
import sys
import argparse
from typing import List, Optional

from .agent import create_agent
from . import __version__

def execute_command(command: str) -> None:
    """Execute a shell command"""
    os.system(command)

def get_user_confirmation(command: str) -> bool:
    """Get user confirmation before executing command"""
    print(f"(AI Answer): {command}")
    response = input("Execute? [Y/n]: ").strip().lower()
    return response in ["", "y", "yes"]

def process_prompt(prompt: str) -> None:
    """Process user prompt and generate shell command"""
    if not prompt.strip():
        print("Error: No prompt provided")
        sys.exit(1)
    
    # Create AI agent
    agent = create_agent()
    
    try:
        # Generate response
        resp = agent.run_sync(prompt)
        answer = resp.output
        
        if answer.success and answer.cmd is not None:
            if get_user_confirmation(answer.cmd):
                execute_command(answer.cmd)
        else:
            print(f"(AI Answer): {answer.failure}")
            print("Command generation failed")
            sys.exit(1)
            
    except Exception as e:
        print(f"Error: {e}")
        sys.exit(1)

def create_parser() -> argparse.ArgumentParser:
    """Create command line argument parser"""
    parser = argparse.ArgumentParser(
        prog="ai",
        description="AI-powered shell command generator using DeepSeek V3",
        epilog="Examples:\n"
               "  ai list files\n"
               "  ai \"show disk usage\"\n"
               "  ai 显示当前目录\n",
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    parser.add_argument(
        "prompt",
        nargs="*",
        help="Natural language description of the command you want"
    )
    
    parser.add_argument(
        "--version",
        action="version",
        version=f"ai-shell {__version__}"
    )
    
    parser.add_argument(
        "--config",
        action="store_true",
        help="Show current configuration"
    )
    
    return parser

def show_config() -> None:
    """Show current configuration"""
    from .config import (get_api_key, get_base_url, get_model, get_timeout,
                        get_max_retries, validate_config, get_available_modern_commands,
                        get_missing_modern_commands, get_installation_hint)

    print("AI Shell Configuration:")
    print(f"  Model: {get_model()}")
    print(f"  Base URL: {get_base_url()}")

    try:
        api_key = get_api_key()
        print(f"  API Key: {api_key[:8]}...{api_key[-4:]}")
    except ValueError as e:
        print(f"  API Key: ❌ {e}")

    print(f"  Timeout: {get_timeout()}s")
    print(f"  Max Retries: {get_max_retries()}")

    print(f"\nConfiguration Status: {'✅ Valid' if validate_config() else '❌ Invalid'}")

    print("\nConfiguration Sources (in priority order):")
    print("  1. Environment variables")
    print("  2. .env file in current directory")
    print("  3. .env file in package directory")
    print("  4. ~/.ai-shell/.env file")
    print("  5. Default values")

    print("\nEnvironment Variables:")
    print("  AI_SHELL_API_KEY - API key")
    print("  AI_SHELL_BASE_URL - API base URL")
    print("  AI_SHELL_MODEL - Model name")
    print("  AI_SHELL_TIMEOUT - Request timeout (seconds)")
    print("  AI_SHELL_MAX_RETRIES - Maximum retry attempts")
    print("  AI_SHELL_MODERN_COMMANDS - Custom modern command alternatives")

    # Show modern commands configuration
    modern_commands = get_available_modern_commands()
    missing_commands = get_missing_modern_commands()

    print(f"\n现代化命令替代:")

    if modern_commands:
        # Group commands for better display
        active_alternatives = {k: v for k, v in modern_commands.items() if k != v}
        unchanged_commands = {k: v for k, v in modern_commands.items() if k == v}

        if active_alternatives:
            print(f"  ✅ 已启用的替代 ({len(active_alternatives)} 个):")
            for old_cmd, new_cmd in sorted(active_alternatives.items()):
                print(f"    {old_cmd} → {new_cmd}")

        if unchanged_commands:
            print(f"  ➡️ 保持原样: {', '.join(sorted(unchanged_commands.keys()))}")

    if missing_commands:
        print(f"\n  ⚠️ 可配置但未安装的工具 ({len(missing_commands)} 个):")
        for old_cmd, new_cmd in sorted(missing_commands.items()):
            hint = get_installation_hint(new_cmd)
            if hint:
                print(f"    {old_cmd} → {new_cmd}")
                print(f"      安装: {hint.split('#')[0].strip()}")
            else:
                print(f"    {old_cmd} → {new_cmd}")

    if not modern_commands and not missing_commands:
        print("  未检测到可用的现代化工具配置")

    print("\n  配置方法:")
    print("    1. 编辑配置文件: ai_shell/modern_commands.toml")
    print("    2. 环境变量: export AI_SHELL_MODERN_COMMANDS=\"ls:eza,cat:bat\"")
    print("    3. .env 文件: AI_SHELL_MODERN_COMMANDS=\"ls:eza,cat:bat\"")

def main() -> None:
    """Main entry point"""
    parser = create_parser()
    args = parser.parse_args()

    if args.config:
        show_config()
        return

    if not args.prompt:
        parser.print_help()
        sys.exit(1)

    # Validate configuration before processing
    from .config import validate_config
    if not validate_config():
        print("❌ Configuration error: API key not configured.")
        print("Please set AI_SHELL_API_KEY in .env file or environment variable.")
        print("Run 'ai --config' to see current configuration.")
        sys.exit(1)

    # Join all prompt arguments into a single string
    prompt = " ".join(args.prompt)
    process_prompt(prompt)

if __name__ == "__main__":
    main()