import os, typer, sys, shlex, tempfile, subprocess
from typing import Optional
from rich.console import Console
from rich.syntax import Syntax
from jazari.slurm_generator import generate_sbatch_script
from jazari.config_generator import load_config, get_config_default
from jazari.utils import get_wandb_api_key

console = Console()
CONFIG = load_config()

def run_command(
    command:     list[str]      = typer.Argument(..., help="The command to run (e.g., python train.py). Use '--' before it."),
    nodes:       int            = typer.Option(get_config_default("nodes", CONFIG, 1), '--nodes', '-N', help='Number of nodes.'),
    gpus:        int            = typer.Option(get_config_default("gpus", CONFIG, 1), '--gpus', '-G', help='GPUs per node.'),
    cpus:        int            = typer.Option(get_config_default("cpus", CONFIG, 1), '--cpus', '-n', help='CPUs per task/GPU'), # Changed -c to -n per our earlier discussion
    time:        str            = typer.Option(get_config_default("time", CONFIG, '01:00:00'), '--time', '-t', help='Time limit (D-HH:MM)'),
    name:        str            = typer.Option('jazari_run', '--name', '-j', help='Job name'), # Changed -n to -j to allow CPU shortcut
    account:     Optional[str]  = typer.Option(get_config_default("account", CONFIG, None), "--account", "-A", help="Slurm account to charge."),
    track_wandb: bool           = typer.Option(get_config_default("track_wandb", CONFIG, False), "--track-wandb", help="Auto-configure W&B."),
    push_to_hub: Optional[str]  = typer.Option(None, "--push-to-hub", help="Hugging Face repo ID to upload model to."),   
    pull_data:   Optional[str]  = typer.Option(None, '--pull-data', help='Huggingface dataset ID to download.'),
    dry_run:     bool           = typer.Option(False, '--dry-run', help='Print sbatch script without submitting.'),
    ):
    '''
    Launch a distributed training job.
    '''
    if not command:
        console.print("[bold red]Error:[/bold red] You must provide a command to run.")
        sys.exit(1)

    log_dir = "logs"
    if not os.path.exists(log_dir):
        os.makedirs(log_dir, exist_ok=True)

    full_command_str = shlex.join(command)

    # --- Weights & Biases ---
    wandb_key = None
    if track_wandb:
        console.print("[dim]W&B tracking enabled.[/dim]")
        wandb_key = get_wandb_api_key()
        if not wandb_key:
            console.print("[bold yellow]Warning:[/bold yellow] Could not find W&B API key locally.")
            console.print("Please run [green]wandb login[/green] on this machine first.")
    # -----------------

    console.print(f'[bold]🐘 Generating Slurm script for: {name}[/bold]')

    sbatch_content = generate_sbatch_script(
        nodes           = nodes,
        gpus_per_node   = gpus,
        cpus_per_task   = cpus,
        time_limit      = time,
        job_name        = name,
        account_name    = account,
        wandb_api_key   = wandb_key,
        hf_repo_id      = push_to_hub,
        hf_dataset_id   = pull_data,
        user_command    = full_command_str
    )

    if dry_run:
        console.print('\n[yellow]--- DRY RUN: Generated #SBATCH Script ---[/yellow]')
        syntax = Syntax(sbatch_content, 'bash', theme = 'monokai', line_numbers = True)
        console.print(syntax)
        console.print('[yellow]-----------------------------------------------[/yellow]')
    else:
        console.print("[dim]Submitting to Slurm scheduler...[/dim]")

        with tempfile.NamedTemporaryFile(mode='w+', suffix=".sh", delete=False) as temp_file:
            temp_script_path = temp_file.name
            temp_file.write(sbatch_content)

        try:
            result = subprocess.run(
                ["sbatch", temp_script_path],
                check           = True,
                capture_output  = True,
                text            = True
            )
            job_id = result.stdout.strip().split()[-1]
            console.print(f"[bold green]✅ Job submitted successfully![/bold green] (ID: [bold]{job_id}[/bold])")
            console.print(f"[dim]View logs: cat logs/{name}-{job_id}.out[/dim]")

        except subprocess.CalledProcessError as e:
            console.print("[bold red]❌ Failed to submit job.[/bold red]")
            console.print(f"Sbatch error: {e.stderr}")
        except FileNotFoundError:
             console.print("\n[bold red]❌ Error:[/bold red] 'sbatch' command not found.")
             console.print("Are you running this on a Slurm login node?")
        finally:
            if os.path.exists(temp_script_path):
                os.unlink(temp_script_path)