import requests
import gradio as gr
from urllib.parse import urlencode
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()


def create_image(stats, username):
    url = "https://argilla.imglab-cdn.net/dibt/dibt_v2.png"

    total_stats = stats["Total Statistics"]
    top_items = stats["Most Popular Items"]

    text = f"""<span size="12pt" weight="bold">Hugging Face  ❤️ {username} in 2024</span>

<span weight="bold">{total_stats['Model Downloads']:,}</span> model downloads
<span weight="bold">{total_stats['Model Likes']:,}</span> model likes
<span weight="bold">{total_stats['Dataset Downloads']:,}</span> dataset downloads
<span weight="bold">{total_stats['Dataset Likes']:,}</span> dataset likes

<span size="10pt">Most Popular Contributions:</span>
Model: <span weight="bold">{top_items['Top Model']['name']}</span>
  ({top_items['Top Model']['downloads']:,} downloads, {top_items['Top Model']['likes']} likes)
Dataset: <span weight="bold">{top_items['Top Dataset']['name']}</span>
  ({top_items['Top Dataset']['downloads']:,} downloads, {top_items['Top Dataset']['likes']} likes)
Space: <span weight="bold">{top_items['Top Space']['name']}</span>
  ({top_items['Top Space']['likes']} likes)"""

    params = {
        "width": "1200",
        "text": text,
        "text-width": "800",
        "text-height": "600",
        "text-padding": "60",
        "text-color": "39,71,111",
        "text-x": "460",
        "text-y": "40",
        "format": "png",
        "dpr": "2",
    }

    return f"{url}?{urlencode(params)}"


def get_user_stats(username):
    headers = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}

    # Get models stats
    models_response = requests.get(
        "https://huggingface.co/api/models",
        params={"author": username, "full": "True"},
        headers=headers,
    )
    models = models_response.json()

    # Get datasets stats
    datasets_response = requests.get(
        "https://huggingface.co/api/datasets",
        params={"author": username, "full": "True"},
        headers=headers,
    )
    datasets = datasets_response.json()

    # Get spaces stats
    spaces_response = requests.get(
        "https://huggingface.co/api/spaces",
        params={"author": username, "full": "True"},
        headers=headers,
    )
    spaces = spaces_response.json()

    # Calculate totals
    total_model_downloads = sum(model.get("downloads", 0) for model in models)
    total_model_likes = sum(model.get("likes", 0) for model in models)
    total_dataset_downloads = sum(dataset.get("downloads", 0) for dataset in datasets)
    total_dataset_likes = sum(dataset.get("likes", 0) for dataset in datasets)
    total_space_likes = sum(space.get("likes", 0) for space in spaces)

    # Find most liked items
    most_liked_model = max(models, key=lambda x: x.get("likes", 0), default=None)
    most_liked_dataset = max(datasets, key=lambda x: x.get("likes", 0), default=None)
    most_liked_space = max(spaces, key=lambda x: x.get("likes", 0), default=None)

    stats = {
        "Total Statistics": {
            "Model Downloads": total_model_downloads,
            "Model Likes": total_model_likes,
            "Dataset Downloads": total_dataset_downloads,
            "Dataset Likes": total_dataset_likes,
            "Space Likes": total_space_likes,
        },
        "Most Popular Items": {
            "Top Model": {
                "name": most_liked_model.get("modelId", "None")
                if most_liked_model
                else "None",
                "likes": most_liked_model.get("likes", 0) if most_liked_model else 0,
                "downloads": most_liked_model.get("downloads", 0)
                if most_liked_model
                else 0,
            },
            "Top Dataset": {
                "name": most_liked_dataset.get("id", "None")
                if most_liked_dataset
                else "None",
                "likes": most_liked_dataset.get("likes", 0)
                if most_liked_dataset
                else 0,
                "downloads": most_liked_dataset.get("downloads", 0)
                if most_liked_dataset
                else 0,
            },
            "Top Space": {
                "name": most_liked_space.get("id", "None")
                if most_liked_space
                else "None",
                "likes": most_liked_space.get("likes", 0) if most_liked_space else 0,
            },
        },
    }

    # Generate image URL
    image_url = create_image(stats, username)

    return image_url


with gr.Blocks(title="Hugging Face Community Stats") as demo:
    gr.Markdown("# Hugging Face Community Recap")
    gr.Markdown(
        "Enter a username to see their impact and top contributions across the Hugging Face Hub"
    )

    with gr.Row():
        username_input = gr.Textbox(
            label="Username", placeholder="Enter Hugging Face username...", scale=4
        )
        submit_btn = gr.Button("Get Stats", scale=1)

    with gr.Row():
        with gr.Column():
            stats_image = gr.Markdown()

    # Add example usernames
    gr.Examples(
        examples=[["merve"], ["mlabonne"], ["bartowski"]],
        inputs=username_input,
        label="Try these examples",
    )

    def format_markdown(image_url):
        return f"![Hugging Face Stats]({image_url})"

    # Handle submission
    submit_btn.click(
        fn=lambda x: format_markdown(get_user_stats(x)),
        inputs=username_input,
        outputs=stats_image,
        api_name="get_stats",
    )
    # Also trigger on enter key
    username_input.submit(
        fn=lambda x: format_markdown(get_user_stats(x)),
        inputs=username_input,
        outputs=stats_image,
    )

if __name__ == "__main__":
    demo.launch()