kvcache-ai-ktransformers/kt-kernel/python/cli/main.py

"""
Main entry point for kt-cli.

KTransformers CLI - A unified command-line interface for KTransformers.
"""

import sys
import warnings

# Suppress numpy subnormal warnings
warnings.filterwarnings("ignore", message="The value of the smallest subnormal")

import typer

from kt_kernel.cli import __version__
from kt_kernel.cli.commands import bench, chat, config, doctor, model, quant, run, sft, version
from kt_kernel.cli.i18n import t, set_lang, get_lang


def _get_app_help() -> str:
    """Get app help text based on current language."""
    lang = get_lang()
    if lang == "zh":
        return "KTransformers CLI - KTransformers 统一命令行界面"
    return "KTransformers CLI - A unified command-line interface for KTransformers."


def _get_help(key: str) -> str:
    """Get help text based on current language."""
    help_texts = {
        "version": {"en": "Show version information", "zh": "显示版本信息"},
        "run": {"en": "Start model inference server", "zh": "启动模型推理服务器"},
        "chat": {"en": "Interactive chat with running model", "zh": "与运行中的模型进行交互式聊天"},
        "quant": {"en": "Quantize model weights", "zh": "量化模型权重"},
        "edit": {"en": "Edit model information", "zh": "编辑模型信息"},
        "bench": {"en": "Run full benchmark", "zh": "运行完整基准测试"},
        "microbench": {"en": "Run micro-benchmark", "zh": "运行微基准测试"},
        "doctor": {"en": "Diagnose environment issues", "zh": "诊断环境问题"},
        "model": {"en": "Manage models and storage paths", "zh": "管理模型和存储路径"},
        "config": {"en": "Manage configuration", "zh": "管理配置"},
        "sft": {"en": "Fine-tuning with LlamaFactory", "zh": "使用 LlamaFactory 进行微调"},
    }
    lang = get_lang()
    return help_texts.get(key, {}).get(lang, help_texts.get(key, {}).get("en", key))


# Create main app with dynamic help
app = typer.Typer(
    name="kt",
    help="KTransformers CLI - A unified command-line interface for KTransformers.",
    no_args_is_help=False,  # Handle no-args case manually to support first-run setup
    add_completion=False,  # Use static completion scripts instead of dynamic completion
    rich_markup_mode="rich",
)


def _update_help_texts() -> None:
    """Update all help texts based on current language setting."""
    # Update main app help
    app.info.help = _get_app_help()

    # Update command help texts
    for cmd_info in app.registered_commands:
        # cmd_info is a CommandInfo object
        if hasattr(cmd_info, "name") and cmd_info.name:
            cmd_info.help = _get_help(cmd_info.name)

    # Update sub-app help texts
    for group_info in app.registered_groups:
        if hasattr(group_info, "name") and group_info.name:
            group_info.help = _get_help(group_info.name)


# Commands are registered later after tui_command is defined


def check_first_run() -> None:
    """Check if this is the first run and prompt for language setup."""
    import os

    # Skip if not running in interactive terminal
    if not sys.stdin.isatty():
        return

    from kt_kernel.cli.config.settings import DEFAULT_CONFIG_FILE

    # Only check if config file exists - don't create it yet
    if not DEFAULT_CONFIG_FILE.exists():
        # First run - show welcome and language selection
        from kt_kernel.cli.config.settings import get_settings

        settings = get_settings()
        _show_first_run_setup(settings)
    else:
        # Config exists - check if initialized
        from kt_kernel.cli.config.settings import get_settings

        settings = get_settings()
        if not settings.get("general._initialized"):
            _show_first_run_setup(settings)


def _show_first_run_setup(settings) -> None:
    """Show first-run setup wizard."""
    from rich.console import Console
    from rich.panel import Panel
    from rich.prompt import Prompt, Confirm
    from rich.spinner import Spinner
    from rich.live import Live

    from kt_kernel.cli.utils.environment import scan_storage_locations, format_size_gb

    console = Console()

    # Welcome message
    console.print()
    console.print(
        Panel.fit(
            "[bold cyan]Welcome to KTransformers CLI! / 欢迎使用 KTransformers CLI![/bold cyan]\n\n"
            "Let's set up your preferences.\n"
            "让我们设置您的偏好。",
            title="kt-cli",
            border_style="cyan",
        )
    )
    console.print()

    # Language selection
    console.print("[bold]Select your preferred language / 选择您的首选语言:[/bold]")
    console.print()
    console.print("  [cyan][1][/cyan] English")
    console.print("  [cyan][2][/cyan] 中文 (Chinese)")
    console.print()

    choice = Prompt.ask("Enter choice / 输入选择", choices=["1", "2"], default="1")
    lang = "en" if choice == "1" else "zh"

    # Save language setting
    settings.set("general.language", lang)
    set_lang(lang)

    # Confirmation message
    console.print()
    if lang == "zh":
        console.print("[green]✓[/green] 语言已设置为中文")
    else:
        console.print("[green]✓[/green] Language set to English")

    # Model discovery section
    console.print()
    if lang == "zh":
        console.print("[bold]发现模型权重[/bold]")
        console.print()
        console.print("[dim]扫描系统中已有的模型权重文件，以便快速添加到模型列表。[/dim]")
        console.print()
        console.print("  [cyan][1][/cyan] 全局扫描 (自动扫描所有非系统路径)")
        console.print("  [cyan][2][/cyan] 手动指定路径 (可添加多个)")
        console.print("  [cyan][3][/cyan] 跳过 (稍后手动添加)")
        console.print()
        scan_choice = Prompt.ask("选择扫描方式", choices=["1", "2", "3"], default="1")
    else:
        console.print("[bold]Discover Model Weights[/bold]")
        console.print()
        console.print("[dim]Scan existing model weights on your system to quickly add them to the model list.[/dim]")
        console.print()
        console.print("  [cyan][1][/cyan] Global scan (auto-scan all non-system paths)")
        console.print("  [cyan][2][/cyan] Manual paths (add multiple paths)")
        console.print("  [cyan][3][/cyan] Skip (add manually later)")
        console.print()
        scan_choice = Prompt.ask("Select scan method", choices=["1", "2", "3"], default="1")

    if scan_choice == "1":
        # Global scan
        from kt_kernel.cli.utils.model_discovery import discover_and_register_global, format_discovery_summary

        console.print()
        try:
            total_found, new_found, registered = discover_and_register_global(
                min_size_gb=2.0, max_depth=6, show_progress=True, lang=lang
            )

            format_discovery_summary(
                total_found=total_found,
                new_found=new_found,
                registered=registered,
                lang=lang,
                show_models=True,
                max_show=10,
            )

        except Exception as e:
            console.print(f"[yellow]Warning: Scan failed - {e}[/yellow]")

    elif scan_choice == "2":
        # Manual path specification
        from kt_kernel.cli.utils.model_discovery import discover_and_register_path
        import os

        discovered_paths = set()  # Track paths discovered in this session
        total_registered = []

        while True:
            console.print()
            if lang == "zh":
                path = Prompt.ask("输入要扫描的路径 (例如: /mnt/data/models)")
            else:
                path = Prompt.ask("Enter path to scan (e.g., /mnt/data/models)")

            # Expand and validate path
            path = os.path.expanduser(path)

            if not os.path.exists(path):
                if lang == "zh":
                    console.print(f"[yellow]警告: 路径不存在: {path}[/yellow]")
                else:
                    console.print(f"[yellow]Warning: Path does not exist: {path}[/yellow]")
                continue

            if not os.path.isdir(path):
                if lang == "zh":
                    console.print(f"[yellow]警告: 不是一个目录: {path}[/yellow]")
                else:
                    console.print(f"[yellow]Warning: Not a directory: {path}[/yellow]")
                continue

            # Scan this path
            console.print()
            try:
                total_found, new_found, registered = discover_and_register_path(
                    path=path, min_size_gb=2.0, existing_paths=discovered_paths, show_progress=True, lang=lang
                )

                # Update discovered paths
                for model in registered:
                    discovered_paths.add(model.path)
                total_registered.extend(registered)

                console.print()
                if lang == "zh":
                    console.print(f"[green]✓[/green] 在此路径找到 {total_found} 个模型，其中 {new_found} 个为新模型")
                else:
                    console.print(f"[green]✓[/green] Found {total_found} models in this path, {new_found} are new")

                if new_found > 0:
                    for model in registered[:5]:
                        console.print(f"  • {model.name} ({model.format})")

                    if len(registered) > 5:
                        if lang == "zh":
                            console.print(f"  [dim]... 还有 {len(registered) - 5} 个新模型[/dim]")
                        else:
                            console.print(f"  [dim]... and {len(registered) - 5} more new models[/dim]")

            except Exception as e:
                console.print(f"[red]Error scanning path: {e}[/red]")

            # Ask if continue
            console.print()
            if lang == "zh":
                continue_scan = Confirm.ask("是否继续添加其他路径?", default=False)
            else:
                continue_scan = Confirm.ask("Continue adding more paths?", default=False)

            if not continue_scan:
                break

        if total_registered:
            console.print()
            if lang == "zh":
                console.print(f"[green]✓[/green] 总共发现 {len(total_registered)} 个新模型")
            else:
                console.print(f"[green]✓[/green] Total {len(total_registered)} new models discovered")

    # Model storage path selection
    console.print()
    console.print(f"[bold]{t('setup_model_path_title')}[/bold]")
    console.print()
    console.print(f"[dim]{t('setup_model_path_desc')}[/dim]")
    console.print()

    # Scan for storage locations
    console.print(f"[dim]{t('setup_scanning_disks')}[/dim]")
    locations = scan_storage_locations(min_size_gb=50.0)
    console.print()

    if locations:
        # Show storage location options
        for i, loc in enumerate(locations[:5], 1):  # Show top 5 options
            available = format_size_gb(loc.available_gb)
            total = format_size_gb(loc.total_gb)

            # Build the option string
            if i == 1:
                option_str = t("setup_disk_option_recommended", path=loc.path, available=available, total=total)
            else:
                option_str = t("setup_disk_option", path=loc.path, available=available, total=total)

            console.print(f"  [cyan][{i}][/cyan] {option_str}")

        # Custom path option
        custom_idx = min(len(locations), 5) + 1
        console.print(f"  [cyan][{custom_idx}][/cyan] {t('setup_custom_path')}")
        console.print()

        valid_choices = [str(i) for i in range(1, custom_idx + 1)]
        path_choice = Prompt.ask(t("prompt_select"), choices=valid_choices, default="1")

        if path_choice == str(custom_idx):
            # Custom path
            selected_path = _prompt_custom_path(console, settings)
        else:
            selected_path = locations[int(path_choice) - 1].path
    else:
        # No large storage found, ask for custom path
        console.print(f"[yellow]{t('setup_no_large_disk')}[/yellow]")
        console.print()
        selected_path = _prompt_custom_path(console, settings)

    # Ensure the path exists
    import os
    from pathlib import Path

    if not os.path.exists(selected_path):
        if Confirm.ask(t("setup_path_not_exist"), default=True):
            try:
                Path(selected_path).mkdir(parents=True, exist_ok=True)
            except (OSError, PermissionError) as e:
                console.print(f"[red]{t('error')}: {e}[/red]")
                # Fall back to default
                selected_path = str(Path.home() / ".ktransformers" / "models")
                Path(selected_path).mkdir(parents=True, exist_ok=True)

    # Check available space and warn if low
    from kt_kernel.cli.utils.environment import detect_disk_space_gb

    available_gb, _ = detect_disk_space_gb(
        selected_path if os.path.exists(selected_path) else str(Path(selected_path).parent)
    )
    if available_gb < 100:
        console.print(f"[yellow]{t('setup_path_low_space')}[/yellow]")

    # Save the path
    settings.set("paths.models", selected_path)
    settings.set("general._initialized", True)

    console.print()
    console.print(f"[green]✓[/green] {t('setup_model_path_set', path=selected_path)}")
    console.print()

    # Tips
    if lang == "zh":
        console.print("[dim]提示: 运行 'kt config show' 查看所有配置[/dim]")
    else:
        console.print("[dim]Tip: Run 'kt config show' to view all settings[/dim]")

    console.print()


def _prompt_custom_path(console, settings) -> str:
    """Prompt user to enter a custom path."""
    from rich.prompt import Prompt
    from pathlib import Path
    import os

    default_path = str(Path.home() / ".ktransformers" / "models")

    while True:
        custom_path = Prompt.ask(t("setup_enter_custom_path"), default=default_path)

        # Expand user home
        custom_path = os.path.expanduser(custom_path)

        # Check if path exists or parent is writable
        if os.path.exists(custom_path):
            if os.access(custom_path, os.W_OK):
                return custom_path
            else:
                console.print(f"[red]{t('setup_path_no_write')}[/red]")
        else:
            # Check if we can create it (parent writable)
            parent = str(Path(custom_path).parent)
            while not os.path.exists(parent) and parent != "/":
                parent = str(Path(parent).parent)

            if os.access(parent, os.W_OK):
                return custom_path
            else:
                console.print(f"[red]{t('setup_path_no_write')}[/red]")


def _install_shell_completion() -> None:
    """Install shell completion scripts to user directories.

    Uses standard locations that are auto-loaded by shell completion systems:
    - Bash: ~/.local/share/bash-completion/completions/kt (auto-loaded by bash-completion 2.0+)
    - Zsh: ~/.zfunc/_kt (requires fpath setup, but commonly used)
    - Fish: ~/.config/fish/completions/kt.fish (auto-loaded)
    """
    import os
    import shutil
    from pathlib import Path

    from kt_kernel.cli.config.settings import get_settings

    settings = get_settings()

    # Check if already installed
    if settings.get("general._completion_installed", False):
        return

    # Detect current shell
    shell = os.environ.get("SHELL", "")
    shell_name = "zsh" if "zsh" in shell else "fish" if "fish" in shell else "bash"

    try:
        cli_dir = Path(__file__).parent
        completions_dir = cli_dir / "completions"
        home = Path.home()

        def install_completion(src_name: str, dest_dir: Path, dest_name: str) -> None:
            """Install completion file from source to destination."""
            src_file = completions_dir / src_name
            if src_file.exists():
                dest_dir.mkdir(parents=True, exist_ok=True)
                shutil.copy2(src_file, dest_dir / dest_name)

        if shell_name == "bash":
            install_completion(
                "kt-completion.bash", home / ".local" / "share" / "bash-completion" / "completions", "kt"
            )
        elif shell_name == "zsh":
            install_completion("_kt", home / ".zfunc", "_kt")
        elif shell_name == "fish":
            install_completion("kt.fish", home / ".config" / "fish" / "completions", "kt.fish")

        # Mark as installed
        settings.set("general._completion_installed", True)

        # For bash/zsh, completion will work in new terminals automatically
        # (bash-completion 2.0+ auto-loads from ~/.local/share/bash-completion/completions/)

    except (OSError, IOError):
        # Silently ignore errors - completion is not critical
        pass


def _apply_saved_language() -> None:
    """Apply the saved language setting.

    Priority:
    1. KT_LANG environment variable (if already set, don't override)
    2. Config file setting
    3. System locale (auto)
    """
    import os

    # Don't override if KT_LANG is already set by user
    if os.environ.get("KT_LANG"):
        return

    from kt_kernel.cli.config.settings import get_settings

    settings = get_settings()
    lang = settings.get("general.language", "auto")

    if lang != "auto":
        set_lang(lang)


app.command(name="version", help="Show version information")(version.version)
app.command(name="chat", help="Interactive chat with running model")(chat.chat)
app.command(name="quant", help="Quantize model weights")(quant.quant)
app.command(name="edit", help="Edit model information")(model.edit_model)
app.command(name="bench", help="Run full benchmark")(bench.bench)
app.command(name="microbench", help="Run micro-benchmark")(bench.microbench)
app.command(name="doctor", help="Diagnose environment issues")(doctor.doctor)

# Register sub-apps
app.add_typer(model.app, name="model", help="Manage models and storage paths")
app.add_typer(config.app, name="config", help="Manage configuration")
app.add_typer(sft.app, name="sft", help="Fine-tuning with LlamaFactory")


def main():
    """Main entry point."""
    # Apply saved language setting first (before anything else for correct help display)
    _apply_saved_language()

    # Update help texts based on language
    _update_help_texts()

    # Check for first run (but not for certain commands)
    # Skip first-run check for: --help, config commands, version
    args = sys.argv[1:] if len(sys.argv) > 1 else []
    skip_commands = ["--help", "-h", "config", "version", "--version", "--no-tui"]

    should_check_first_run = True
    for arg in args:
        if arg in skip_commands:
            should_check_first_run = False
            break

    # Handle no arguments case
    if not args:
        # Check if this is first run
        from kt_kernel.cli.config.settings import DEFAULT_CONFIG_FILE, get_settings

        is_first_run = False
        if not DEFAULT_CONFIG_FILE.exists():
            is_first_run = True
        else:
            settings = get_settings()
            if not settings.get("general._initialized"):
                is_first_run = True

        if is_first_run:
            # First run - start initialization
            _install_shell_completion()
            check_first_run()
            return
        else:
            # Not first run - show help
            app(["--help"])
            return

    # Auto-install shell completion on first run
    if should_check_first_run:
        _install_shell_completion()

    # Check first run before running commands
    if should_check_first_run:
        check_first_run()

    # Handle "run" command specially to pass through unknown options
    if args and args[0] == "run":
        # Get args after "run"
        run_args = args[1:]
        # Use click command directly with ignore_unknown_options
        from kt_kernel.cli.commands import run as run_module

        sys.exit(run_module.run.main(args=run_args, standalone_mode=False))

    app()


if __name__ == "__main__":
    main()