import gradio as gr
from codecarbon import EmissionsTracker
import os
import json
from datetime import datetime
import requests
from huggingface_hub import HfApi
import tempfile
from dotenv import load_dotenv
import webbrowser

# Load environment variables
load_dotenv()

# Get environment variables
HF_TOKEN = os.getenv("HF_TOKEN")
if not HF_TOKEN:
    print("Warning: HF_TOKEN not found in environment variables. Submissions will not work.")

api = HfApi(token=HF_TOKEN)


DEFAULT_PARAMS = {
    "text":{
        "dataset_name": "QuotaClimat/frugalaichallenge-text-train",
        "test_size": 0.2,  # must be between 0 and 1
        "test_seed": 42,   # must be non-negative
    },
    "image":{
        "dataset_name": "pyronear/pyro-sdis",
        "test_size": 0.2,  # must be between 0 and 1
        "test_seed": 42,   # must be non-negative
    },
    "audio":{
        "dataset_name": "rfcx/frugalai",
        "test_size": 0.2,  # must be between 0 and 1
        "test_seed": 42,   # must be non-negative
    }
}


def evaluate_model(task: str, space_url: str):
    """
    Evaluate a model through its API endpoint
    """
    # username = space_url.split("/")[0]

    if "localhost" in space_url:
        api_url = f"{space_url}/{task}"
    else:

        try:
            info_space = api.space_info(repo_id=space_url)
        except:
            return None, None, None, gr.Warning(f"Space '{space_url}' not found, it needs to be in the format username/space-name")

        
        host = info_space.host
        api_url = f"{host}/{task}"
    
    try:
        # Make API call to the space
        params = DEFAULT_PARAMS[task]
        response = requests.post(api_url, json=params)
        if response.status_code != 200:
            return None, None, None, gr.Warning(f"API call failed with status {response.status_code}")
        
        results = response.json()
        
        # Check for required keys based on task
        base_required_keys = {
            "username", "space_url", "submission_timestamp", "model_description",
            "energy_consumed_wh", "emissions_gco2eq", "emissions_data",
            "api_route", "dataset_config"
        }

        # Add task-specific accuracy keys
        if task == "image":
            accuracy_keys = {"classification_accuracy", "mean_iou"}
        else:  # text and audio
            accuracy_keys = {"accuracy"}
        
        required_keys = base_required_keys | accuracy_keys
        
        missing_keys = required_keys - set(results.keys())
        if missing_keys:
            return None, None, None, gr.Warning(f"API response missing required keys: {', '.join(missing_keys)}")
        
        # Return appropriate accuracy metric based on task
        if task == "image":
            accuracy = results["classification_accuracy"]  # For display in UI
        else:
            accuracy = results["accuracy"]
        
        return (
            accuracy,
            results["emissions_gco2eq"],
            results["energy_consumed_wh"],
            results
        )
        
    except Exception as e:
        return None, None, None, gr.Warning(str(e))


def submit_results(task: str, results_json):
    if not results_json:
        return gr.Warning("No results to submit")
    
    if not HF_TOKEN:
        return gr.Warning("HF_TOKEN not found. Please set up your Hugging Face token.")
    
    try:
        results_str = json.dumps(results_json)
        
        # Create a temporary file with the results
        with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix='.json') as f:
            f.write(results_str)
            temp_path = f.name
        
        # Upload to the appropriate dataset based on task
        api = HfApi(token=HF_TOKEN)
        path_in_repo = f"submissions/{results_json['username']}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        
        dataset_mapping = {
            "text": "frugal-ai-challenge/public-leaderboard-text",
            "image": "frugal-ai-challenge/public-leaderboard-image",
            "audio": "frugal-ai-challenge/public-leaderboard-audio"
        }
        
        api.upload_file(
            path_or_fileobj=temp_path,
            path_in_repo=path_in_repo,
            repo_id=dataset_mapping[task],
            repo_type="dataset",
            token=HF_TOKEN
        )
        
        # Clean up
        os.unlink(temp_path)
        
        return gr.Info("Results submitted successfully to the leaderboard! 🎉")
    except Exception as e:
        return gr.Warning(f"Error submitting results: {str(e)}")

# Create the demo interface
with gr.Blocks(
    css="""
    .button-link > a {
        display: inline-block;
        padding: 0.5rem 1.5rem;
        background-color: #FF7C01;
        color: white !important;
        text-decoration: none;
        border-radius: 0.5rem;
        border: none;
        cursor: pointer;
        text-align: center;
        font-weight: 600;
        width: 100%;
        box-shadow: 0 1px 3px rgba(0,0,0,0.12), 0 1px 2px rgba(0,0,0,0.24);
        transition: all 0.3s cubic-bezier(.25,.8,.25,1);
    }
    .button-link > a:hover {
        background-color: #E66E00;
        box-shadow: 0 3px 6px rgba(0,0,0,0.16), 0 3px 6px rgba(0,0,0,0.23);
        text-decoration: none;
    }
    """
).queue(default_concurrency_limit=20) as demo:  # Allow up to 20 concurrent requests by default

    gr.Image("./logo.png", show_label=False, container=False)
    
    gr.Markdown("""
    # Frugal AI Challenge - Submission Portal
    Submit your model results for any of the three tasks: Text, Image, or Audio classification.
    """)
    
    with gr.Tabs():


        with gr.Tab("Instructions"):

            gr.Markdown("""
To submit your results in one of the three tasks, please follow the steps below:
                        
## Prepare your model submission
1. Duplicate the template of the submission API by duplicating this space https://huggingface.co/spaces/frugal-ai-challenge/submission-template on your own Hugging Face account.
2. In ``tasks/text.py``, ``tasks/image.py``, or ``tasks/audio.py``, modify the ``evaluate_model`` function to replace the baseline by your model loading and inference within the inference pass where the energy consumption and emissions are tracked.
3. Eventually complete the requirements and/or any necessaries dependencies in your space.
4. Write down your model card in the ``README.md`` file.
5. Deploy your space (FastAPI) and verify that it works.
6. (Optional) You can change the Space hardware to use any GPU directly on Hugging Face.
                        
## Submit your model to the leaderboard in the ``Model Submission`` tab
When your API is deployed : 

0. Fill out the [submission form](https://framaforms.org/2025-frugal-ai-challenge-submission-form-1736883260-0) with all the details regarding your team and project.
1. Select the task you want to submit your model to
2. Enter the Space URL of your API
3. (Optional) Precise the API route (default is ``/text``, ``/image``, or ``/audio``)
4. Step 1 - Evaluate model: Click on the button to evaluate your model. This will run you model on your API, computes the accuracy on the test set (20% of the train set), and track the energy consumption and emissions.
5. Step 2 - Submit to leaderboard (optional): Click on the button to submit your results to the leaderboard. This will upload the results to the leaderboard dataset and update the leaderboard.
6. Step 3 - Submit to final evaluation (as a form): [Click on the button to submit your results to the challenge](https://framaforms.org/2025-frugal-ai-challenge-submission-form-1736883260-0). This will open a form to submit your results to the challenge.
7. You can see the public leaderboards at the following links - they are mostly informational because we will rank the models on the private dataset after the challenge ended, but you can see the current state of the leaderboard.
    - Text - https://huggingface.co/datasets/frugal-ai-challenge/public-leaderboard-text
    - Image - https://huggingface.co/datasets/frugal-ai-challenge/public-leaderboard-image
    - Audio - https://huggingface.co/datasets/frugal-ai-challenge/public-leaderboard-audio

## About
> You can find more information about the Frugal AI Challenge 2025 on the [Frugal AI Challenge website](https://frugalaichallenge.org/).
> Or directly on the organization page on Hugging Face: [Frugal AI Challenge](https://huggingface.co/frugal-ai-challenge)
                        
This portal is a submission portal for the Frugal AI Challenge 2025. It is a simple interface to evaluate and submit your model to the leaderboard.
The challenge is organized by Hugging Face, Data For Good, and the French Ministry of Environment. 

The goal of the Frugal AI Challenge is to encourage both academic and industry actors to keep efficiency in mind when deploying AI models. By tracking both energy consumption and performance for different AI tasks, we can incentivize frugality in AI deployment while also addressing real-world challenges.
""")

        # Text Classification Task
        with gr.Tab("📜 Text Classification"):
            with gr.Row():
                text_space_url = gr.Textbox(
                    label="Space URL",
                    placeholder="username/your-space",
                    lines=1
                )
                text_route = gr.Textbox(
                    label="API route (Advanced)",
                    value="/text",
                    lines=1
                )
            
            with gr.Row():
                with gr.Column(scale=1):
                    text_evaluate_btn = gr.Button("1. Evaluate model", variant="secondary")
                with gr.Column(scale=1):
                    text_submit_btn = gr.Button("2. Submit to public leaderboard (optional)", variant="secondary")
                with gr.Column(scale=1):
                    text_form = gr.Button(value="3. Submit to final evaluation form", link="https://framaforms.org/2025-frugal-ai-challenge-submission-form-1736883260-0")
        
            with gr.Row():
                text_accuracy = gr.Number(label="Accuracy", precision=4)
                text_energy = gr.Number(label="Energy Consumed (Wh)", precision=12)
                text_emissions = gr.Number(label="Emissions (gCO2eq)", precision=12)
            with gr.Row():
                text_results_json = gr.JSON(label="Detailed Results", visible=True)
        
        # Image Classification Task
        with gr.Tab("🎥 Image Classification"):
            with gr.Row():
                image_space_url = gr.Textbox(
                    label="Space URL",
                    placeholder="username/your-space",
                    lines=1
                )
                image_route = gr.Textbox(
                    label="API route (Advanced)",
                    value="/image",
                    lines=1
                )
            
            with gr.Row():
                with gr.Column(scale=1):
                    image_evaluate_btn = gr.Button("1. Evaluate model", variant="secondary")
                with gr.Column(scale=1):
                    image_submit_btn = gr.Button("2. Submit to public leaderboard (optional)", variant="secondary")
                with gr.Column(scale=1):
                     image_form = gr.Button(value="3. Submit to final evaluation form", link="https://framaforms.org/2025-frugal-ai-challenge-submission-form-1736883260-0")
        
            with gr.Row():
                image_accuracy = gr.Number(label="Accuracy", precision=4)
                image_energy = gr.Number(label="Energy Consumed (Wh)", precision=12)
                image_emissions = gr.Number(label="Emissions (gCO2eq)", precision=12)
            with gr.Row():
                image_results_json = gr.JSON(label="Detailed Results", visible=True)
        
        # Audio Classification Task
        with gr.Tab("🔊 Audio Classification"):
            with gr.Row():
                audio_space_url = gr.Textbox(
                    label="Space URL",
                    placeholder="username/your-space",
                    lines=1
                )
                audio_route = gr.Textbox(
                    label="API route (Advanced)",
                    value="/audio",
                    lines=1
                )
            
            with gr.Row():
                with gr.Column(scale=1):
                    audio_evaluate_btn = gr.Button("1. Evaluate model", variant="secondary")
                with gr.Column(scale=1):
                    audio_submit_btn = gr.Button("2. Submit to public leaderboard (optional)", variant="secondary")
                with gr.Column(scale=1):
                    audio_form = gr.Button(value="3. Submit to final evaluation form", link="https://framaforms.org/2025-frugal-ai-challenge-submission-form-1736883260-0")
        
                    
            with gr.Row():
                audio_accuracy = gr.Number(label="Accuracy", precision=4)
                audio_energy = gr.Number(label="Energy Consumed (Wh)", precision=12)
                audio_emissions = gr.Number(label="Emissions (gCO2eq)", precision=12)
            with gr.Row():
                audio_results_json = gr.JSON(label="Detailed Results", visible=True)


    FORM_URL = "https://framaforms.org/2025-frugal-ai-challenge-submission-form-1736883260-0"

    def open_form():
        webbrowser.open_new_tab(FORM_URL)
        return gr.Info("Opening submission form in new tab...")

    # Set up event handlers with specific queue configurations
    text_evaluate_btn.click(
        lambda url, route: evaluate_model(route.strip("/"), url),
        inputs=[text_space_url, text_route],
        outputs=[text_accuracy, text_emissions, text_energy, text_results_json],
        concurrency_limit=5,  # Allow 5 concurrent model evaluations
        concurrency_id="eval_queue"  # Share evaluation queue across tasks
    )
    
    text_submit_btn.click(
        lambda results: submit_results("text", results),
        inputs=[text_results_json],
        outputs=None,
        concurrency_limit=10,  # Allow 10 concurrent submissions
        concurrency_id="submit_queue"  # Share submission queue across tasks
    )
    
    image_evaluate_btn.click(
        lambda url, route: evaluate_model(route.strip("/"), url),
        inputs=[image_space_url, image_route],
        outputs=[image_accuracy, image_emissions, image_energy, image_results_json],
        concurrency_limit=5,  # Share same limit with text evaluation
        concurrency_id="eval_queue"
    )
    
    image_submit_btn.click(
        lambda results: submit_results("image", results),
        inputs=[image_results_json],
        outputs=None,
        concurrency_limit=10,
        concurrency_id="submit_queue"
    )
    
    audio_evaluate_btn.click(
        lambda url, route: evaluate_model(route.strip("/"), url),
        inputs=[audio_space_url, audio_route],
        outputs=[audio_accuracy, audio_emissions, audio_energy, audio_results_json],
        concurrency_limit=5,
        concurrency_id="eval_queue"
    )
    
    audio_submit_btn.click(
        lambda results: submit_results("audio", results),
        inputs=[audio_results_json],
        outputs=None,
        concurrency_limit=10,
        concurrency_id="submit_queue"
    )


if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )