import json
import os

import pandas as pd

from src.display.formatting import has_no_nan_values, make_clickable_model
# changes to be made here
from src.display.utils import AutoEvalColumn, EvalQueueColumn, OpenEndedColumns, MedSafetyColumns, MedicalSummarizationColumns, ACIColumns, SOAPColumns, HealthbenchColumns, HealthbenchHardColumns, OpenEndedArabicColumn, OpenEndedFrenchColumn, OpenEndedSpanishColumn, OpenEndedPortugueseColumn, OpenEndedRomanianColumn, OpenEndedGreekColumn, ClosedEndedMultilingualColumns
from src.leaderboard.read_evals import get_raw_eval_results
from src.envs import PRIVATE_REPO

def get_leaderboard_df(results_path: str, requests_path: str, cols: list, benchmark_cols: list, evaluation_metric:str, subset:str) -> pd.DataFrame:
    """Creates a dataframe from all the individual experiment results"""
    raw_data =  get_raw_eval_results(results_path, requests_path, evaluation_metric)
    # print(raw_data)
    # raise Exception("stop")
    # if subset.startswith("healthbench"):
    #     breakpoint()
    all_data_json = [v.to_dict(subset=subset) for v in raw_data if not v.full_model.startswith("/models_llm")]
    df = pd.DataFrame.from_records(all_data_json)
    # changes to be made here
    if subset == "datasets":
        df = df.sort_values(by=[AutoEvalColumn.average.name], ascending=False)
    elif subset == "med_safety":
        df = df.sort_values(by=["Harmfulness Score"], ascending=True)
    elif subset.startswith("open_ended"):
        df = df.sort_values(by=["ELO"], ascending=False)
    elif subset == "medical_summarization":
        df = df.sort_values(by=[AutoEvalColumn.overall.name], ascending=False)
    elif subset == "aci":
        df = df.sort_values(by=[AutoEvalColumn.overall.name], ascending=False)
    elif subset == "soap":
        df = df.sort_values(by=[AutoEvalColumn.overall.name], ascending=False)
    elif subset == "closed_ended_arabic":
        df = df.sort_values(by=[AutoEvalColumn.average.name], ascending=False)
    elif subset == "healthbench":
        df = df.sort_values(by=["Overall Score"], ascending=False)
    elif subset == "healthbench_hard":
        df = df.sort_values(by=["Overall Score"], ascending=False)
    elif subset == "closed_ended_multilingual":
        df = df.sort_values(by=[AutoEvalColumn.average.name], ascending=False)
    cols = list(set(df.columns).intersection(set(cols)))
    df = df[cols].round(decimals=2)
    # filter out if any of the benchmarks have not been produced
    df = df[has_no_nan_values(df, benchmark_cols)]
    return raw_data, df

def get_evaluation_queue_df(save_path: str, cols: list) -> list[pd.DataFrame]:
    """Creates the different dataframes for the evaluation queues requestes"""
    entries = [entry for entry in os.listdir(save_path) if not entry.startswith(".")]
    all_evals = []
    
    for entry in entries:
        full_path = os.path.join(save_path, entry)

        if os.path.isdir(full_path):
            continue

        if entry.endswith(".json"):
            with open(full_path) as fp:
                data = json.load(fp)

            data[EvalQueueColumn.model.name] = make_clickable_model(data["model_name"]) if not data["private"] else data["model_name"]
            data[EvalQueueColumn.revision.name] = data.get("revision", "main")
            data[EvalQueueColumn.closed_ended_status.name] = data["status"]["closed-ended"]
            data[EvalQueueColumn.open_ended_status.name] = data["status"]["open-ended"]
            data[EvalQueueColumn.med_safety_status.name] = data["status"]["med-safety"]
            data[EvalQueueColumn.medical_summarization_status.name] = data["status"]["medical-summarization"]
            data[EvalQueueColumn.note_generation_status.name] = data["status"]["note-generation"]
            if PRIVATE_REPO:
                data[EvalQueueColumn.closed_ended_arabic_status.name] = data["status"]["closed-ended-arabic"]
            all_evals.append(data)
    
    pending_list = []
    running_list = []
    finished_list = []
    for run in all_evals:
        status_list = [run["status"]["closed-ended"], run["status"]["open-ended"], run["status"]["med-safety"], run["status"]["medical-summarization"], run["status"]["note-generation"]]
        if PRIVATE_REPO:
            status_list.append(run["status"]["closed-ended-arabic"])
        
        if "RUNNING" in status_list:
            running_list.append(run)
        elif "PENDING" in status_list or "RERUN" in status_list:
            pending_list.append(run)
        else:
            finished_list.append(run)

    df_pending = pd.DataFrame.from_records(pending_list, columns=cols)
    df_running = pd.DataFrame.from_records(running_list, columns=cols)
    df_finished = pd.DataFrame.from_records(finished_list, columns=cols)
    
    return df_finished[cols], df_running[cols], df_pending[cols]