import os, glob
import json
from datetime import datetime, timezone
from dataclasses import dataclass
from datasets import load_dataset, Dataset
import pandas as pd
import gradio as gr
from huggingface_hub import HfApi, snapshot_download, ModelInfo, list_models
from enum import Enum

OWNER = "AIEnergyScore"
COMPUTE_SPACE = f"{OWNER}/launch-computation-example"
TOKEN = os.environ.get("DEBUG")
API = HfApi(token=TOKEN)

def add_docker_eval_with_agreement(zip_file, agreement):
    if not agreement:
        gr.Warning("You must agree to the terms before submitting your energy score data.")
        return
    
    new_fid_list = zip_file.split("/")
    new_fid = new_fid_list[-1]
    if new_fid.endswith('.zip'):
        API.upload_file(
            path_or_fileobj=zip_file,
            repo_id="AIEnergyScore/tested_proprietary_models",
            path_in_repo='submitted_models/'+new_fid,
            repo_type="dataset",
            commit_message="Adding logs via submission Space.",
            token=TOKEN
        )
        gr.Info('Uploaded logs to dataset! We will validate their validity and add them to the next version of the leaderboard.')
    else:
        gr.Info('You can only upload .zip files here!')

with gr.Blocks() as demo:
    gr.Markdown("# AI Energy Score | Submission Portal")
    gr.Markdown("### The goal of the AI Energy Score project is to develop an energy-based rating system for AI model deployment that will guide members of the community in choosing models for different tasks based on energy efficiency.")
    
    with gr.Row():
        with gr.Column():
            with gr.Accordion("Submit log files from a Docker run:", open=False):
                gr.Markdown("If you've already benchmarked your model using the [Docker file](https://github.com/huggingface/EnergyStarAI/) provided, please upload the **entire run log directory** (in .zip format) below:")
                
                agreement_checkbox = gr.Checkbox(label="I agree to the following terms:")
                agreement_text = gr.Markdown("""
                By checking the box below and submitting your energy score data, you confirm and agree to the following:
                1. **Public Data Sharing**: You consent to the public sharing of the energy performance data derived from your submission. No additional information related to this model including proprietary configurations will be disclosed.
                2. **Data Integrity**: You validate that the log files submitted are accurate, unaltered, and generated directly from testing your model as per the specified benchmarking procedures.
                3. **Model Representation**: You verify that the model tested and submitted is representative of the production-level version of the model, including its level of quantization and any other relevant characteristics impacting energy efficiency and performance.
                """)
                
                file_output = gr.File(visible=False)
                u = gr.UploadButton("Upload a zip file with logs", file_count="single")
                
                u.upload(add_docker_eval_with_agreement, [u, agreement_checkbox], file_output)

demo.launch()