ai-shell/ai_shell/config.py

"""
Configuration module for AI Shell
"""

import os
from pathlib import Path

try:
    from dotenv import load_dotenv
except ImportError:
    load_dotenv = None

# Load .env file if it exists
def load_env_file() -> None:
    """Load environment variables from .env file"""
    if load_dotenv is None:
        return

    # Try to find .env file in current directory or package directory
    env_paths = [
        Path.cwd() / ".env",
        Path(__file__).parent.parent / ".env",
        Path.home() / ".ai-shell" / ".env",
    ]

    for env_path in env_paths:
        if env_path.exists():
            load_dotenv(env_path)
            break

# Load .env file on import
load_env_file()

# Default API configuration (fallback values)
DEFAULT_API_KEY = "your_api_key_here"
DEFAULT_BASE_URL = "https://api.openai.com/v1/"
DEFAULT_MODEL = "gpt-3.5-turbo"

def get_api_key() -> str:
    """Get API key from environment or use default"""
    api_key = os.getenv("AI_SHELL_API_KEY", DEFAULT_API_KEY)
    if api_key == DEFAULT_API_KEY:
        raise ValueError(
            "API key not configured. Please set AI_SHELL_API_KEY in .env file or environment variable."
        )
    return api_key

def get_base_url() -> str:
    """Get base URL from environment or use default"""
    return os.getenv("AI_SHELL_BASE_URL", DEFAULT_BASE_URL)

def get_model() -> str:
    """Get model name from environment or use default"""
    return os.getenv("AI_SHELL_MODEL", DEFAULT_MODEL)

def get_timeout() -> int:
    """Get request timeout from environment"""
    return int(os.getenv("AI_SHELL_TIMEOUT", "30"))

def get_max_retries() -> int:
    """Get max retries from environment"""
    return int(os.getenv("AI_SHELL_MAX_RETRIES", "3"))

def get_modern_commands() -> dict:
    """Get modern command alternatives configuration"""
    # Default modern command alternatives
    default_alternatives = {
        "ls": "eza",           # Modern ls replacement
        "cat": "bat",          # Syntax highlighting cat
        "find": "fd",          # Fast and user-friendly find
        "grep": "rg",          # Ripgrep - faster grep
        "du": "ncdu",          # Interactive disk usage
        "df": "duf",           # Modern df with better output
        "ps": "procs",         # Modern ps replacement
        "top": "htop",         # Interactive process viewer
        "ping": "gping",       # Ping with graph
        "curl": "httpie",      # User-friendly HTTP client
        "wget": "aria2c",      # Multi-connection downloader
        "diff": "delta",       # Better diff with syntax highlighting
        "tree": "tree",        # Keep tree as is (already modern)
        "sed": "sd",           # Simpler sed alternative
        "awk": "choose",       # Human-friendly awk alternative
        "cut": "choose",       # Alternative to cut
        "sort": "sort",        # Keep sort as is
        "uniq": "uniq",        # Keep uniq as is
        "head": "head",        # Keep head as is
        "tail": "tail",        # Keep tail as is
        "less": "bat",         # Use bat for paging too
        "more": "bat",         # Use bat for paging
        "vim": "nvim",         # Modern vim
        "nano": "micro",       # Modern nano alternative
        "cd": "zoxide",        # Smart cd with frecency
        "cp": "cp",            # Keep cp as is (or could use rsync)
        "mv": "mv",            # Keep mv as is
        "rm": "trash",         # Safer deletion (trash-cli)
        "mkdir": "mkdir",      # Keep mkdir as is
        "rmdir": "rmdir",      # Keep rmdir as is
        "chmod": "chmod",      # Keep chmod as is
        "chown": "chown",      # Keep chown as is
        "tar": "ouch",         # Universal archive tool
        "zip": "ouch",         # Universal archive tool
        "unzip": "ouch",       # Universal archive tool
        "ssh": "ssh",          # Keep ssh as is
        "scp": "rsync",        # More efficient file transfer
        "rsync": "rsync",      # Keep rsync as is
        "git": "git",          # Keep git as is
        "docker": "docker",    # Keep docker as is
        "python": "python",    # Keep python as is
        "node": "node",        # Keep node as is
        "npm": "pnpm",         # Faster npm alternative
        "yarn": "pnpm",        # Faster yarn alternative
    }

    # Try to load custom alternatives from environment or config file
    custom_alternatives_str = os.getenv("AI_SHELL_MODERN_COMMANDS", "")
    custom_alternatives = {}

    if custom_alternatives_str:
        try:
            # Parse format: "old1:new1,old2:new2,old3:new3"
            pairs = custom_alternatives_str.split(",")
            for pair in pairs:
                if ":" in pair:
                    old_cmd, new_cmd = pair.strip().split(":", 1)
                    custom_alternatives[old_cmd.strip()] = new_cmd.strip()
        except Exception:
            pass  # Ignore parsing errors

    # Merge default and custom alternatives
    alternatives = default_alternatives.copy()
    alternatives.update(custom_alternatives)

    return alternatives

def get_available_modern_commands() -> dict:
    """Get only the modern commands that are actually installed on the system"""
    import shutil

    all_alternatives = get_modern_commands()
    available_alternatives = {}

    for old_cmd, new_cmd in all_alternatives.items():
        # Check if the modern command is actually available
        if shutil.which(new_cmd):
            available_alternatives[old_cmd] = new_cmd
        # If modern command not available, keep the original
        elif shutil.which(old_cmd):
            available_alternatives[old_cmd] = old_cmd

    return available_alternatives

def generate_modern_commands_prompt() -> str:
    """Generate prompt text about available modern command alternatives"""
    available_commands = get_available_modern_commands()

    if not available_commands:
        return ""

    prompt_parts = [
        "\n\n**IMPORTANT: Modern Command Preferences**",
        "When generating shell commands, prefer these modern alternatives when available:",
        ""
    ]

    # Group by categories for better readability
    categories = {
        "文件操作": ["ls", "cat", "find", "tree", "cp", "mv", "rm"],
        "文本处理": ["grep", "sed", "awk", "cut", "sort", "uniq", "head", "tail", "less", "more", "diff"],
        "系统监控": ["ps", "top", "du", "df"],
        "网络工具": ["ping", "curl", "wget"],
        "编辑器": ["vim", "nano"],
        "导航": ["cd"],
        "压缩工具": ["tar", "zip", "unzip"],
        "包管理": ["npm", "yarn"],
        "其他": []
    }

    for category, commands in categories.items():
        category_commands = []
        for cmd in commands:
            if cmd in available_commands and available_commands[cmd] != cmd:
                category_commands.append(f"  - Use `{available_commands[cmd]}` instead of `{cmd}`")

        if category_commands:
            prompt_parts.append(f"**{category}:**")
            prompt_parts.extend(category_commands)
            prompt_parts.append("")

    # Add remaining commands not in categories
    other_commands = []
    categorized_commands = set()
    for commands in categories.values():
        categorized_commands.update(commands)

    for old_cmd, new_cmd in available_commands.items():
        if old_cmd not in categorized_commands and new_cmd != old_cmd:
            other_commands.append(f"  - Use `{new_cmd}` instead of `{old_cmd}`")

    if other_commands:
        prompt_parts.append("**其他工具:**")
        prompt_parts.extend(other_commands)
        prompt_parts.append("")

    prompt_parts.extend([
        "These modern tools typically provide:",
        "- Better performance and speed",
        "- More intuitive syntax and output",
        "- Enhanced features and functionality",
        "- Better error messages and user experience",
        "",
        "Always prioritize these modern alternatives when they fulfill the user's requirements."
    ])

    return "\n".join(prompt_parts)

def setup_environment() -> None:
    """Setup environment variables for OpenAI client"""
    os.environ["OPENAI_API_KEY"] = get_api_key()
    os.environ["OPENAI_BASE_URL"] = get_base_url()

def validate_config() -> bool:
    """Validate configuration"""
    try:
        get_api_key()
        get_base_url()
        get_model()
        return True
    except ValueError:
        return False