import json
import os
from datetime import datetime, timezone
from typing import Optional
from src.display.formatting import styled_error, styled_message, styled_warning
from src.envs import API, EVAL_REQUESTS_PATH, TOKEN, QUEUE_REPO
from src.submission.check_validity import (
    already_submitted_models,
    check_model_card,
    get_model_size,
    is_model_on_hub,
)
import gradio as gr

REQUESTED_MODELS = None
USERS_TO_SUBMISSION_DATES = None

def add_new_eval_option1(
    benchmark: str,
    model: str,
    base_model: str,
    revision: str,
    precision: str,
    temperature: str,
    top_p: str,
    top_k: str,
    presence_penalty: str,
    frequency_penalty: str,
    repetition_penalty: str,
    vllm_version: str,
    user_state: str,
    organization_list: list
):
    global REQUESTED_MODELS
    global USERS_TO_SUBMISSION_DATES
    if not REQUESTED_MODELS:
        REQUESTED_MODELS, USERS_TO_SUBMISSION_DATES = already_submitted_models(EVAL_REQUESTS_PATH)

    user_name = ""
    model_path = model
    if "/" in model:
        user_name = model.split("/")[0]
        model_path = model.split("/")[1]

    precision = precision.split(" ")[0]
    current_time = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S %z")

    # Check submitter qualification
    if user_name != user_state and user_name not in organization_list:
        return styled_error("The submitter does not have submission rights for this model.")
        
    # Does the organization submit more than three times in a day?
    submission_times = [item['submitted_time'] for item in USERS_TO_SUBMISSION_DATES[user_name] if item['benchmark'] == benchmark]
    submission_cnt = 0
    for i in range(len(submission_times)):
        hours_diff = (datetime.strptime(current_time, "%Y-%m-%dT%H:%M:%S %z") - datetime.strptime(submission_times[i], "%Y-%m-%dT%H:%M:%S %z")).total_seconds() / 3600
        if hours_diff <= 24:
            submission_cnt += 1
    if submission_cnt > 3:
        return styled_error("The organization already submitted three times for this benchmark today.")

    # Does the model actually exist?
    if revision == "":
        revision = "main"

    # Is the model info correctly filled?
    try:
        model_info = API.model_info(repo_id=model, revision=revision)
    except Exception:
        return styled_error("Could not get your model information. Please fill it up properly.")

    model_size = get_model_size(model_info=model_info, precision=precision)

    # Were the model card and license filled?
    try:
        license = model_info.cardData["license"]
    except Exception:
        return styled_error("Please select a license for your model.")

    modelcard_OK, error_msg = check_model_card(model)
    if not modelcard_OK:
        return styled_error(error_msg)

    if temperature == "":
        temperature = "1.0"
    
    if top_p == "":
        top_p = "1.0"
    
    if top_k == "":
        top_k = "-1"
    
    if presence_penalty == "":
        presence_penalty = "0.0"
    
    if frequency_penalty == "":
        frequency_penalty = "0.0"

    if repetition_penalty == "":
        repetition_penalty = "1.0"
    
    # Seems good, creating the eval
    print("Adding new eval")

    eval_entry = {
        "benchmark": benchmark,
        "model": model,
        "base_model": base_model,
        "revision": revision,
        "precision": precision,
        "status": "PENDING",
        "submitted_time": current_time,
        "likes": model_info.likes,
        "params": model_size,
        "license": license,
        "private": False,
        "temperature": float(temperature),
        "top_p": float(top_p),
        "top_k": float(top_k),
        "vllm_version": vllm_version,
        "presence_penalty": float(presence_penalty),
        "frequency_penalty": float(frequency_penalty),
        "repetition_penalty": float(repetition_penalty),
        "load_model_code": "None",
        "inference_code": "None",
        "termination_code": "None",
    }

    # Check for duplicate submission
    submission_times = [item['submitted_time'] for item in USERS_TO_SUBMISSION_DATES[user_name] if item['benchmark'] == benchmark and item['model'] == model]
    submission_cnt = 0
    for i in range(len(submission_times)):
        hours_diff = (datetime.strptime(current_time, "%Y-%m-%dT%H:%M:%S %z") - datetime.strptime(submission_times[i], "%Y-%m-%dT%H:%M:%S %z")).total_seconds() / 3600
        if hours_diff <= 24:
            submission_cnt += 1
    if submission_cnt > 1:
        return styled_warning("This model has been already submitted within 24 hours.")

    print("Creating eval file")
    OUT_DIR = f"{EVAL_REQUESTS_PATH}/{user_name}"
    os.makedirs(OUT_DIR, exist_ok=True)
    out_path = f"{OUT_DIR}/{benchmark}_{model_path}_eval_request_False.json"

    with open(out_path, "w") as f:
        f.write(json.dumps(eval_entry))

    print("Uploading eval file")
    API.upload_file(
        path_or_fileobj=out_path,
        path_in_repo=out_path.split("eval-queue/")[1],
        repo_id=QUEUE_REPO,
        repo_type="dataset",
        commit_message=f"Add {model} to eval queue",
    )

    # Remove the local file
    os.remove(out_path)

    return styled_message(
        "Your request has been submitted to the evaluation queue!"

    )


def add_new_eval_option2(
    benchmark: str,
    model: str,
    base_model: str,
    revision: str,
    precision: str,
    temperature: str,
    top_p: str,
    top_k: str,
    presence_penalty: str,
    frequency_penalty: str,
    repetition_penalty: str,
    load_model_code: str,
    inference_code: str,
    termination_code: str,
    user_state: str,
    organization_list: list
):
    global REQUESTED_MODELS
    global USERS_TO_SUBMISSION_DATES
    if not REQUESTED_MODELS:
        REQUESTED_MODELS, USERS_TO_SUBMISSION_DATES = already_submitted_models(EVAL_REQUESTS_PATH)

    user_name = ""
    model_path = model
    if "/" in model:
        user_name = model.split("/")[0]
        model_path = model.split("/")[1]

    precision = precision.split(" ")[0]
    current_time = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%S %z")
    
    # Check submitter qualification
    if user_name != user_state and user_name not in organization_list:
        return styled_error("The submitter does not have submission rights for this model.")
    
    # Does the organization submit more than three times in a day?
    submission_times = [item['submitted_time'] for item in USERS_TO_SUBMISSION_DATES[user_name] if item['benchmark'] == benchmark]
    submission_cnt = 0
    for i in range(len(submission_times)):
        hours_diff = (datetime.strptime(current_time, "%Y-%m-%dT%H:%M:%S %z") - datetime.strptime(submission_times[i], "%Y-%m-%dT%H:%M:%S %z")).total_seconds() / 3600
        if hours_diff <= 24:
            submission_cnt += 1
    if submission_cnt > 3:
        return styled_error("The organization already submitted three times for this benchmark today.")

    # Does the model actually exist?
    if revision == "":
        revision = "main"

    # Is the model info correctly filled?
    try:
        model_info = API.model_info(repo_id=model, revision=revision)
    except Exception:
        return styled_error("Could not get your model information. Please fill it up properly.")

    model_size = get_model_size(model_info=model_info, precision=precision)

    # Were the model card and license filled?
    try:
        license = model_info.cardData["license"]
    except Exception:
        return styled_error("Please select a license for your model.")

    modelcard_OK, error_msg = check_model_card(model)
    if not modelcard_OK:
        return styled_error(error_msg)

    if temperature == "":
        temperature = "1.0"
    
    if top_p == "":
        top_p = "1.0"
    
    if top_k == "":
        top_k = "-1"
    
    if presence_penalty == "":
        presence_penalty = "0.0"
    
    if frequency_penalty == "":
        frequency_penalty = "0.0"

    if repetition_penalty == "":
        repetition_penalty = "1.0"
    
    # Seems good, creating the eval
    print("Adding new eval")

    eval_entry = {
        "benchmark": benchmark,
        "model": model,
        "base_model": base_model,
        "revision": revision,
        "precision": precision,
        "status": "PENDING",
        "submitted_time": current_time,
        "likes": model_info.likes,
        "params": model_size,
        "license": license,
        "private": False,
        "temperature": float(temperature),
        "top_p": float(top_p),
        "top_k": float(top_k),
        "vllm_version": "None",
        "presence_penalty": float(presence_penalty),
        "frequency_penalty": float(frequency_penalty),
        "repetition_penalty": float(repetition_penalty),
        "load_model_code": load_model_code,
        "inference_code": inference_code,
        "termination_code": termination_code
    }
    
    # Check for duplicate submission
    submission_times = [item['submitted_time'] for item in USERS_TO_SUBMISSION_DATES[user_name] if item['benchmark'] == benchmark and item['model'] == model]
    submission_cnt = 0
    for i in range(len(submission_times)):
        hours_diff = (datetime.strptime(current_time, "%Y-%m-%dT%H:%M:%S %z") - datetime.strptime(submission_times[i], "%Y-%m-%dT%H:%M:%S %z")).total_seconds() / 3600
        if hours_diff <= 24:
            submission_cnt += 1
    if submission_cnt > 1:
        return styled_warning("This model has been already submitted within 24 hours.")

    print("Creating eval file")
    OUT_DIR = f"{EVAL_REQUESTS_PATH}/{user_name}"
    os.makedirs(OUT_DIR, exist_ok=True)
    out_path = f"{OUT_DIR}/{benchmark}_{model_path}_eval_request_False.json"

    with open(out_path, "w") as f:
        f.write(json.dumps(eval_entry))

    print("Uploading eval file")
    API.upload_file(
        path_or_fileobj=out_path,
        path_in_repo=out_path.split("eval-queue/")[1],
        repo_id=QUEUE_REPO,
        repo_type="dataset",
        commit_message=f"Add {model} to eval queue",
    )

    # Remove the local file
    os.remove(out_path)

    return styled_message(
        "Your request has been submitted to the evaluation queue!"
    )