import subprocess
from datetime import datetime
import sys
import argparse


class bcolors:
    OKGREEN = "\033[92m"
    WARNING = "\033[93m"
    FAIL = "\033[91m"
    ENDC = "\033[0m"


def wait_for_cmd(cmd_list, log_file, quiet=False):
    process = subprocess.Popen(
        cmd_list,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        text=True,
    )

    with open(log_file, "a") as f:
        for line in process.stdout:
            f.write(line)
            # Filter out verbose INFO logs if quiet mode is enabled
            if quiet and line.strip().startswith("INFO:modal_runner_eval:"):
                continue
            print(f"{line}", end="")
    status = process.wait()
    return status


def main(
    timestamp=datetime.now().strftime("%Y_%m_%d-%H_%M_%S"),
    config_file="task_eval/sample_config.json",
    log_file="automate_eval_logs.txt",
    eval_only=False,
    quiet=False,
):
    print(f"{bcolors.OKGREEN} TIMESTAMP: {timestamp} {bcolors.ENDC} ")

    # Overwrite the log file at the start of each run
    with open(log_file, "w") as f:
        f.write(f"Evaluation run started at {timestamp}\n")
        f.write("=" * 80 + "\n\n")

    cmd_eval = [
        sys.executable,
        "-m",
        "modal",
        "run",
        "suno_utils/worker/modal_runner_eval.py",
        "--config-file",
        config_file,
        "--time-label",
        timestamp,
    ]
    if eval_only:
        cmd_eval.append("--eval-only")
        cmd_eval.append("1")
    wait_for_cmd(cmd_eval, log_file, quiet)

    cmd_score = [
        sys.executable,
        "task_eval/score_tasks.py",
        "--timestamp",
        timestamp,
    ]
    status = wait_for_cmd(cmd_score, log_file, quiet)
    if status != 0:
        print(f"{bcolors.FAIL} Score tasks run failed with status code {status} {bcolors.ENDC}")
        return

    cmd_score = [
        sys.executable,
        "task_eval/plot_scores.py",
        "--input_dir",
        f"/app2/suno/data/ditto_evals/{timestamp}",
        "--baseline_config",
        "task_eval/baseline_config.json",
    ]
    status = wait_for_cmd(cmd_score, log_file, quiet)
    if status != 0:
        print(f"{bcolors.FAIL} Plot evals run failed with status code {status} {bcolors.ENDC}")
        return

    print("Evaluation Complete!")


# example usage: python task_eval/automate_evals.py -c task_eval/crow_config.json
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="generate and run ditto evals")
    parser.add_argument(
        "--timestamp",
        type=str,
        default=datetime.now().strftime("%Y_%m_%d-%H_%M_%S"),
        help="Timestamp for labeling the run (default: now)",
    )
    parser.add_argument(
        "-c",
        "--config_file",
        type=str,
        default="task_eval/sample_config.json",
        help="Path to the config file",
    )

    parser.add_argument(
        "--log_file", type=str, default="automate_eval_logs.txt", help="Path to write logs to"
    )
    parser.add_argument(
        "--eval_only",
        type=bool,
        default=False,
        help="Run evaluation on existing generations (all configs must have 'mapping': path/to/mapping/file)",
    )
    parser.add_argument(
        "--quiet",
        action="store_true",
        help="Reduce verbose logging output by filtering INFO logs from modal_runner_eval",
    )

    args = parser.parse_args()
    main(
        timestamp=args.timestamp,
        config_file=args.config_file,
        log_file=args.log_file,
        eval_only=args.eval_only,
        quiet=args.quiet,
    )