from rich.console import Console
from rich.prompt import Prompt, Confirm
from jazari.config_generator import load_config, make_config_dir, CONFIG_FILE

console = Console()

def init_command():
    """
    Initialize Jazari config for this cluster.
    """
    current_config = load_config()

    console.print('[bold blue]Welcome to Jazari setup.[/bold blue]')
    console.print(f"This will create a configuration file at [dim]{CONFIG_FILE}[/dim]\n")

    # 1. Ask for Slurm Account
    default_account = current_config.get("account", "")
    account = Prompt.ask("Enter your default Slurm account (e.g., def-user)",default=default_account)

    # 2. Ask for default time limit
    default_time = current_config.get("time", "01:00:00")
    time_limit = Prompt.ask("Enter default time limit (D-HH:MM)",default=default_time)

    # 3. Ask for W&B tracking default
    default_wandb = current_config.get("track_wandb", False)
    track_wandb = Confirm.ask("Enable W&B tracking by default?",default=default_wandb)

    new_config = {
        "account":      account,
        "time":         time_limit,
        "track_wandb":  track_wandb
    }

    make_config_dir(new_config)

    console.print(f"\n[bold green]✅ Configuration saved![/bold green]")
    console.print("You can now run jobs without specifying these flags.")