import typer, sys, os, tempfile, subprocess, netrc, shutil, getpass

from typing import Optional
from pathlib import Path

from rich.console import Console
from rich.syntax import Syntax
from rich.prompt import Prompt, Confirm
from rich.table import Table
from rich import box

from jazari.slurm_generator import generate_sbatch_script
from jazari.config_generator import load_config, get_config_default, make_config_dir, CONFIG_FILE

# Initializing Rich console for pretty printing
console = Console()

# Load config at start.
CONFIG = load_config()

# --- HELPER FUNCTIONS ---

def get_current_user() -> str:
    """Returns the current username, robustly on Linux environments."""
    try:
        return os.environ['USER']
    except KeyError:
        return getpass.getuser()

def get_wandb_api_key() -> Optional[str]:
    """
    Attempts to read the W&B API key from the local ~/.netrc file.
    """
    try:
        login_info = netrc.netrc().authenticators("api.wandb.ai")
        if login_info:
            username, _, api_key = login_info
            console.print(f"[dim]Found W&B API key for user: {username}[/dim]")
            return api_key
        else:
            return None
    except FileNotFoundError:
        return None
    except Exception as e:
        console.print(f"[dim bold red]Warning: Could not read W&B credentials: {e}[/dim]")
        return None

# --- COMMAND FUNCTIONS (No decorators here) ---

def status_command():
    """
    View the status of active Slurm jobs for the current user.
    """
    # 1. Check if squeue exists available
    if not shutil.which("squeue"):
        console.print("[bold red]Error:[/bold red] 'squeue' command not found.")
        console.print("Are you running this on a Slurm login node?")
        sys.exit(1)

    user = get_current_user()

    # 2. Define the squeue command
    # -u: Filter by user
    # -h: No header (we'll make our own)
    # -o: Custom format with '|' delimiter.
    #     %i=ID, %j=Name, %T=State, %M=Time Used, %D=Nodes, %R=Reason/Nodelist
    squeue_cmd = ["squeue", "-u", user, "-h", "-o", "%i|%j|%T|%M|%D|%R"]

    try:
        # 3. Run the command and capture output
        result = subprocess.run(squeue_cmd, check=True, capture_output=True, text=True)
        output_lines = result.stdout.strip().split('\n')

        # Remove empty lines if squeue returned nothing
        output_lines = [line for line in output_lines if line.strip()]

        if not output_lines:
             console.print(f"\n[dim]No active jobs found for user: [bold]{user}[/bold][/dim]\n")
             return

        # 4. Build the Rich Table
        table = Table(
            title=f"🐘 Cluster Activity for user: [bold green]{user}[/bold green]",
            box=box.ROUNDED,
            header_style="bold blue",
            expand=True
        )
        table.add_column("Job ID", style="cyan", no_wrap=True)
        table.add_column("Name")
        table.add_column("State")
        table.add_column("Time Used", justify="right")
        table.add_column("Nodes", justify="right")
        table.add_column("Nodelist / Reason", style="dim")

        # 5. Parse output and add rows
        for line in output_lines:
            try:
                job_id, name, state, time_used, nodes, reason = line.split('|')

                # Color-code the state
                if state == "RUNNING":
                    state_formatted = f"[bold green]{state}[/bold green]"
                elif state == "PENDING":
                    state_formatted = f"[yellow]{state}[/yellow]"
                else:
                    state_formatted = f"[red]{state}[/red]"

                table.add_row(job_id, name, state_formatted, time_used, nodes, reason)
            except ValueError:
                # Handle rare cases where splitting fails
                continue

        console.print(table)

    except subprocess.CalledProcessError as e:
         console.print(f"[bold red]Error running squeue:[/bold red] {e.stderr}")

def kill_command(
        job_id: Optional[str] = typer.Argument(None, help = 'Job ID to cancel.'),
        cancel_all: bool = typer.Option(False, '--all', help = 'Cancel ALL active jobs.')
):
    # Cancel a specific job or all jobs
    if not shutil.which("scancel"):
        console.print("[bold red]Error:[/bold red] 'scancel' command not found.")
        console.print("Are you running this on a Slurm login node?")
        sys.exit(1)
    
    user = get_current_user()

    if cancel_all:
        if not Confirm.ask(f"[bold red]⚠️  DANGER:[/bold red] Cancel ALL jobs for '{user}'?"):
             sys.exit(0)
        subprocess.run(["scancel", "-u", user], check=True)
        console.print(f"[bold green]💥 All jobs cancelled.[/bold green]")
        return

    if job_id:
        subprocess.run(["scancel", job_id], check=True)
        console.print(f"[bold green]💥 Signal sent to cancel job {job_id}.[/bold green]")
        return

    console.print("[yellow]Usage: jazari kill <JOB_ID> or jazari kill --all[/yellow]")

def logs_command(
    job_id: Optional[str] = typer.Argument(None, help="Job ID to view."),
    tail: bool = typer.Option(False, "--tail", "-f", help="Follow output live."),
    error: bool = typer.Option(False, "--error", "-e", help="View .err log instead of .out")
):
    """View or follow job logs."""
    log_dir = Path("logs")
    if not log_dir.exists():
        console.print("[red]Error: 'logs/' directory not found.[/red]")
        sys.exit(1)

    suffix = ".err" if error else ".out"
    target_file = None

    if job_id:
        matching = list(log_dir.glob(f"*{job_id}{suffix}"))
        if not matching:
             console.print(f"[red]No log found for Job {job_id}[/red]")
             sys.exit(1)
        target_file = matching[0]
    else:
        # Find latest
        files = list(log_dir.glob(f"*{suffix}"))
        if not files:
             console.print(f"[yellow]No {suffix} files found.[/yellow]")
             return
        files.sort(key=os.path.getmtime)
        target_file = files[-1]

    console.print(f"[dim]Viewing: {target_file}[/dim]")
    cmd = ["tail", "-f", str(target_file)] if tail else ["cat", str(target_file)]
    
    try:
        subprocess.run(cmd, check=True)
    except KeyboardInterrupt:
        console.print("\n[dim]Stopped.[/dim]")

def init_command():
    """
    Initialize Jazari config for this cluster.
    Interactively asks the user for defaults and saves them to ~/.jazari/config.yaml.
    """
    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
    # We use the current config value as the default prompt if it exists.
    default_account = 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 = 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 = CONFIG.get("track_wandb", False)
    track_wandb = Confirm.ask(
        "Enable W&B tracking by default?",
        default=default_wandb
    )

    # Prepare the config dictionary
    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.")

def run_command(
    # The default values for these arguments are now pulled dynamically from the
    # CONFIG dictionary using the get_config_default helper.
    command:     list[str]      = typer.Argument(..., help="The command to run (e.g., python train.py --batch 64). Use '--' before it if it has flags."),
    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', '-c', help='CPUs per task/GPU'),
    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', '-n', help='Job name'),
    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 (e.g. 'my-org/my-model)."),   
    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 = ' '.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]")

        # --- Submission ---
        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
            )
            # Extract the job ID from the output "Submitted batch job 12345"
            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)

# --- APP DEFINITION AND COMMAND REGISTRATION ---

# Initializing the Typer app
app = typer.Typer(
    name="jazari",
    help="🐘 The orchestration layer for modern Slurm clusters.",
    add_completion=False,
    no_args_is_help=True
)

# Explicitly register the commands
app.command(name="status")(status_command)
app.command(name="kill")(kill_command)
app.command(name="logs")(logs_command)
app.command(name="init")(init_command)
app.command(name="run")(run_command)

if __name__ == "__main__":
    app()