import argparse
import os
import re
import subprocess
import time


def get_job_data():
    # get job ids, partitions, runtime and total time
    result = subprocess.run(
        ["squeue", "--user=rr41qemu", "--format=%i,%P,%j,%T,%M,%l", "--noheader"],
        capture_output=True,
        text=True,
        # shell=True,
    )
    if result.returncode != 0:
        print("Error getting job ids.")
        return []
    job_data = dict()
    for line in result.stdout.splitlines():
        if line.strip():
            job_id, partition, job_name, state, runtime, total_time = line.split(",")
            job_data[job_id] = {
                "partition": partition,
                "job_name": job_name,
                "state": state,
                "runtime": runtime,
                "total_time": total_time,
            }
    return job_data


def tail_progress(log_file, n=1) -> str:
    # get last modified time
    output = ""
    c_time = os.path.getmtime(log_file)
    try:
        with open(log_file, "r") as f:
            lines = f.readlines()

        output += f"Last update time: {time.ctime(c_time)}\n"
        for line in lines[-n:]:
            if "%" in line:  # crude tqdm-like filter
                output = output + line.strip()
    except FileNotFoundError:
        output += f"Log file {log_file} not found.\n"
        pass
    except Exception as e:
        output += f"Error reading log file {log_file}: {e}\n"

    return output


def get_stats(log_file) -> str:
    try:
        with open(log_file, "r") as f:
            lines = f.readlines()

        # find num_gpus
        num_gpus = 0
        for line in lines:
            if "Training on " in line:
                match = re.search(r"Training on (\d+) GPUs", line)
                num_gpus = int(match.group(1))
                break

        # find num_gpu lines to get models
        gpu_models = {}
        searched = False
        for line in reversed(lines):
            # pattern is "Rank X: Current device: $device_name on local rank $rank"
            if "Rank " in line and "Current device: " in line:
                match = re.search(r"Rank (\d+): Current device: (.+) on local rank \d+", line)
                if match:
                    rank = int(match.group(1))
                    if rank == 0 and searched:
                        # if rank 0 is already searched, we can stop searching
                        break
                    elif rank == 0:
                        searched = True

                    if rank not in gpu_models:
                        gpu_models[rank] = dict()

                    device_name = match.group(2)
                    gpu_models[rank]["device_name"] = device_name

        # get batch size, pattern: "Rank 0: Estimated batch size: 64 for GPU xxx"
        for rank in gpu_models.keys():
            for line in reversed(lines):
                if f"Rank {rank}: Estimated batch size: " in line:
                    match = re.search(r"Rank (\d+): Estimated batch size: (\d+) for GPU", line)
                    if match:
                        gpu_models[rank]["batch_size"] = int(match.group(2))
                    break

        # find run_id
        run_id = None
        for line in lines:
            if "Run ID: " in line:
                match = re.search(r"Run ID: (.+)", line)
                if match:
                    run_id = match.group(1)
                break

        # find number of trainings
        num_trainings = 0
        for line in lines:
            if "Rank 0: Number of trainings: " in line:
                match = re.search(r"Rank 0: Number of trainings: (\d+)", line)
                if match:
                    num_trainings = int(match.group(1))
                break

        # get current model training name
        current_model_training_name = None
        for line in reversed(lines):
            # pattern is "Rank 0: running $model_name"
            if "Rank 0: Running: " in line:
                match = re.search(r"Rank 0: Running: (.+)", line)
                if match:
                    current_model_training_name = match.group(1)
                break

        # get model parameters, pattern: " Creating model with parameters: {'embed_dim': 128, 'num_enc_layers': 4, 'num_heads': 4, 'seq_len': 240}"
        model_parameters = None
        for line in reversed(lines):
            if "Creating model with parameters: " in line:
                match = re.search(r"Creating model with parameters: (.+)", line)
                if match:
                    model_parameters = match.group(1)
                break

        # get input parameters pattern is "Input_parameters: {'item_limit': 1000, 'num_dataloader_workers': 4}"
        input_parameters = None
        for line in reversed(lines):
            if "Input_parameters: " in line:
                match = re.search(r"Input_parameters: (.+)", line)
                if match:
                    input_parameters = match.group(1)
                break

        # find current training number
        current_training_number = 0
        for line in reversed(lines):
            if "Rank 0: Running: " in line:
                current_training_number += 1
        if num_trainings == 0:
            num_trainings = f"{current_training_number}+"

        # find last epoch
        current_epoch = None
        num_epochs = None
        for line in reversed(lines):
            # pattern is "Rank 0: Epoch: 1/10"
            if "Rank 0: Epoch" in line:
                match = re.search(r"Rank 0: Epoch (\d+)/(\d+)", line)
                if match:
                    current_epoch = int(match.group(1))
                    num_epochs = int(match.group(2))
                break

        # get last validation loss
        val_losses = list()
        if current_epoch > 1:
            for line in lines:
                # patttern is "Rank 0: Overall val loss: 0.4439"
                if "Rank 0: Overall val loss" in line:
                    match = re.search(r"Rank 0: Overall val loss: ([\d.]+)", line)
                    if match:
                        val_losses.append(float(match.group(1)))
            if val_losses:
                last_val_loss = val_losses[-1]
                best_val_loss = min(val_losses)
            else:
                last_val_loss = None
                best_val_loss = None
        else:
            last_val_loss = None
            best_val_loss = None

        is_validation = False
        is_training = False
        is_evaluation = False
        for line in reversed(lines):
            if "Evaluating model" in line:
                is_evaluation = True
                break
            if "Validation" in line:
                is_validation = True
                break
            if "Rank 0: Epoch " in line:
                is_training = True
                break
        if is_validation:
            status = "Validation"
        elif is_training:
            status = "Training"
        elif is_evaluation:
            status = "Evaluation"
        else:
            status = "Unknown"

        # create output
        output = ""
        output += f"Run ID: {run_id}, Current training: {current_training_number} of {num_trainings}\n"
        output += f"Number of GPUs: {num_gpus}\n"
        output += f"GPUs used: {gpu_models}\n"
        output += f"Current model training name: {current_model_training_name}\n"
        output += f"Model parameters: {model_parameters}\n"
        output += f"Input parameters: {input_parameters}\n"
        output += f"Current epoch: {current_epoch} of {num_epochs}\n"
        output += f"Last validation loss: {last_val_loss}, Best validation loss: {best_val_loss}\n"
        output += f"Status: {status}\n"

    except FileNotFoundError:
        output = f"Log file {log_file} not found."
        pass
    except Exception as e:
        output = f"Error reading log file {log_file}: {e}"

    return output


def monitor_jobs(logdir, interval=5):
    try:
        while True:
            job_data = get_job_data()
            outputs = list()
            for job_id, job in job_data.items():
                # add header with job data
                outputs.append(f"\n--- Job {job_id} ---")
                outputs.append(f"Partition: {job['partition']}, Job Name: {job['job_name']}, State: {job['state']}, "
                               f"Runtime: {job['runtime']}, Total Time: {job['total_time']}\n")

                err_file_found = True
                out_file_found = True

                err_file = os.path.join(logdir, f"{job['job_name']}_{job_id}.err")
                if not os.path.exists(err_file):
                    # try alternative with job_id only
                    err_file = os.path.join(logdir, f"{job_id}.err")
                if not os.path.exists(err_file):
                    err_file_found = False

                out_file = os.path.join(logdir, f"{job['job_name']}_{job_id}.out")
                if not os.path.exists(out_file):
                    # try alternative with job_id only
                    out_file = os.path.join(logdir, f"{job_id}.out")
                if not os.path.exists(out_file):
                    out_file_found = False

                if err_file_found and out_file_found:
                    stat_output = get_stats(out_file)
                    progress_output = tail_progress(err_file)
                    outputs.append(stat_output)
                    outputs.append(progress_output)
                else:
                    if not err_file_found:
                        outputs.append(f"Error log file: {err_file} not found.\n")
                    if not out_file_found:
                        outputs.append(f"Output log file: {out_file} not found.\n")

                    outputs.append("Job might not have started yet.\n")

            os.system('clear')
            print(f"Monitoring training logs in: {logdir}")
            print("\n".join(outputs))

            time.sleep(interval)
    except KeyboardInterrupt:
        print("Stopped monitoring.")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Monitor SLURM job progress from log files.")
    parser.add_argument("--logdir", required=True, help="Directory containing SLURM .err log files")
    parser.add_argument("--interval", type=int, default=5, help="Update interval in seconds")
    args = parser.parse_args()

    monitor_jobs(args.logdir, args.interval)