ai-shell/ai_shell/config.py

"""
Configuration module for AI Shell
"""

import os
import shutil
from pathlib import Path
from typing import Dict, List, Optional

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

try:
    import tomllib  # Python 3.11+
except ImportError:
    try:
        import tomli as tomllib  # fallback for older Python
    except ImportError:
        tomllib = 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 load_modern_commands_config() -> Dict[str, any]:
    """Load modern commands configuration from TOML file"""
    config_file = Path(__file__).parent / "modern_commands.toml"

    if not config_file.exists():
        return {"commands": {}, "descriptions": {}, "categories": {}, "installation_hints": {}}

    if tomllib is None:
        # Fallback to basic parsing if tomllib not available
        return {"commands": {}, "descriptions": {}, "categories": {}, "installation_hints": {}}

    try:
        with open(config_file, "rb") as f:
            return tomllib.load(f)
    except Exception:
        return {"commands": {}, "descriptions": {}, "categories": {}, "installation_hints": {}}

def get_modern_commands() -> Dict[str, str]:
    """Get modern command alternatives configuration"""
    # Load from TOML config file
    config = load_modern_commands_config()
    default_alternatives = config.get("commands", {})

    # Try to load custom alternatives from environment
    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[str, str]:
    """Get only the modern commands that are actually installed on the system"""
    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 (if it exists)
        elif shutil.which(old_cmd):
            available_alternatives[old_cmd] = old_cmd

    return available_alternatives

def get_missing_modern_commands() -> Dict[str, str]:
    """Get modern commands that are configured but not installed"""
    all_alternatives = get_modern_commands()
    missing_commands = {}

    for old_cmd, new_cmd in all_alternatives.items():
        # If modern command is not available but old command exists
        if not shutil.which(new_cmd) and shutil.which(old_cmd) and old_cmd != new_cmd:
            missing_commands[old_cmd] = new_cmd

    return missing_commands

def get_installation_hint(command: str) -> Optional[str]:
    """Get installation hint for a modern command"""
    config = load_modern_commands_config()
    installation_hints = config.get("installation_hints", {})
    return installation_hints.get(command)

def get_command_description(command: str) -> Optional[str]:
    """Get description for a modern command"""
    config = load_modern_commands_config()
    descriptions = config.get("descriptions", {})
    return descriptions.get(command)

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

    # Only include commands that have modern alternatives available
    active_alternatives = {k: v for k, v in available_commands.items() if k != v}

    if not active_alternatives:
        return ""

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

    # Use categories from config file
    categories = config.get("categories", {})
    descriptions = config.get("descriptions", {})

    for category_name, commands in categories.items():
        category_commands = []
        for cmd in commands:
            if cmd in active_alternatives:
                new_cmd = active_alternatives[cmd]
                desc = descriptions.get(new_cmd, "")
                if desc:
                    category_commands.append(f"  - Use `{new_cmd}` instead of `{cmd}` ({desc})")
                else:
                    category_commands.append(f"  - Use `{new_cmd}` instead of `{cmd}`")

        if category_commands:
            # Convert category name to display format
            display_name = category_name.replace("_", " ").title()
            prompt_parts.append(f"**{display_name}:**")
            prompt_parts.extend(category_commands)
            prompt_parts.append("")

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

    other_commands = []
    for old_cmd, new_cmd in active_alternatives.items():
        if old_cmd not in categorized_commands:
            desc = descriptions.get(new_cmd, "")
            if desc:
                other_commands.append(f"  - Use `{new_cmd}` instead of `{old_cmd}` ({desc})")
            else:
                other_commands.append(f"  - Use `{new_cmd}` instead of `{old_cmd}`")

    if other_commands:
        prompt_parts.append("**Other Tools:**")
        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