import os
import sys
import subprocess
from dotenv import find_dotenv, load_dotenv


def evaluate_model_all_epochs_v2(
    model_name,
    adapter_path_base=None,
    start_epoch=0,
    load_in_4bit=True,
    num_of_entries=-1,
    result_file=None,
):
    new_env = os.environ.copy()
    new_env["MODEL_NAME"] = model_name
    model = model_name.split("/")[-1]

    new_env["LOAD_IN_4BIT"] = "true" if load_in_4bit else "false"
    if result_file is not None:
        new_env["LOGICAL_REASONING_RESULTS_PATH"] = result_file

    if adapter_path_base is None:
        num_train_epochs = 0
        print(f"No adapter path provided. Running with base model:{model_name}")
    else:
        # find subdirectories in adapter_path_base
        # and sort them by epoch number
        subdirs = [
            d
            for d in os.listdir(adapter_path_base)
            if os.path.isdir(os.path.join(adapter_path_base, d))
        ]

        subdirs = sorted(subdirs, key=lambda x: int(x.split("-")[-1]))
        num_train_epochs = len(subdirs)
        print(f"found {num_train_epochs} checkpoints: {subdirs}")
        end_epoch = os.getenv("END_EPOCH")
        if end_epoch:
            num_train_epochs = int(end_epoch)

    for i in range(start_epoch, num_train_epochs + 1):
        print(f"Epoch {i}")
        if i == 0:
            os.unsetenv("ADAPTER_NAME_OR_PATH")
        else:
            adapter_path = adapter_path_base + "/" + subdirs[i - 1]
            new_env["ADAPTER_NAME_OR_PATH"] = adapter_path

        print(f"adapter path: {new_env.get('ADAPTER_NAME_OR_PATH')}")

        log_file = "./logs/{}_epoch_{}.txt".format(model, i)
        with open(log_file, "w") as f_obj:
            subprocess.run(
                f"python llm_toolkit/eval_logical_reasoning.py {num_of_entries}",
                shell=True,
                env=new_env,
                stdout=f_obj,
                text=True,
            )


if __name__ == "__main__":
    found_dotenv = find_dotenv(".env")

    if len(found_dotenv) == 0:
        found_dotenv = find_dotenv(".env.example")
    print(f"loading env vars from: {found_dotenv}")
    load_dotenv(found_dotenv, override=False)

    workding_dir = os.path.dirname(found_dotenv)
    os.chdir(workding_dir)
    sys.path.append(workding_dir)
    print("workding dir:", workding_dir)
    print(f"adding {workding_dir} to sys.path")
    sys.path.append(workding_dir)

    model_name = os.getenv("MODEL_NAME")
    adapter_path_base = os.getenv("ADAPTER_PATH_BASE")
    start_epoch = int(os.getenv("START_EPOCH", 0))
    load_in_4bit = os.getenv("LOAD_IN_4BIT", "true").lower() == "true"
    result_file = os.getenv("LOGICAL_REASONING_RESULTS_PATH", None)

    num_of_entries = int(sys.argv[1]) if len(sys.argv) > 1 else -1

    evaluate_model_all_epochs_v2(
        model_name,
        adapter_path_base=adapter_path_base,
        start_epoch=start_epoch,
        load_in_4bit=load_in_4bit,
        num_of_entries=num_of_entries,
        result_file=result_file,
    )