import requests
import gradio as gr
from urllib.parse import urlencode
import os
from datetime import datetime

# Load environment variables

DEFAULT_IMAGE = "https://hub-recap.imglab-cdn.net/default.jpg?width=1200&text=%3Cspan+size%3D%2212pt%22+weight%3D%22bold%22%3EHugging+Face++%E2%9D%A4%EF%B8%8F+bartowski+in+2024%3C%2Fspan%3E%0A%0A%3Cspan+weight%3D%22bold%22%3E2%2C020%2C552%3C%2Fspan%3E+model+downloads%0A%3Cspan+weight%3D%22bold%22%3E5%2C407%3C%2Fspan%3E+model+likes%0A%3Cspan+weight%3D%22bold%22%3E0%3C%2Fspan%3E+dataset+downloads%0A%3Cspan+weight%3D%22bold%22%3E0%3C%2Fspan%3E+dataset+likes%0A%0A%3Cspan+size%3D%2210pt%22%3EMost+Popular+Contributions%3A%3C%2Fspan%3E%0AModel%3A+%3Cspan+weight%3D%22bold%22%3Ebartowski%2Fgemma-2-9b-it-GGUF%3C%2Fspan%3E%0A++%2843%2C949+downloads%2C+196+likes%29%0ADataset%3A+%3Cspan+weight%3D%22bold%22%3ENone%3C%2Fspan%3E%0A++%280+downloads%2C+0+likes%29%0ASpace%3A+%3Cspan+weight%3D%22bold%22%3Ebartowski%2Fgguf-metadata-updater%3C%2Fspan%3E%0A++%287+likes%29&text-width=800&text-height=600&text-padding=60&text-color=39%2C71%2C111&text-x=460&text-y=40&format=png&dpr=2"


def create_image(stats, username):
    # Determine which image to use based on highest value
    total_stats = stats["Total Statistics"]
    model_activity = total_stats["Model Downloads"] + total_stats["Model Likes"]
    dataset_activity = total_stats["Dataset Downloads"] + total_stats["Dataset Likes"]
    space_activity = total_stats["Space Likes"]

    # Choose base image URL based on highest activity
    # if everything is 0, we show the empty image
    if model_activity == 0 and dataset_activity == 0 and space_activity == 0:
        url = "https://hub-recap.imglab-cdn.net/images/empty.png"
        avatar = "new! We couldn't find your stats on the Hub, maybe in 2025?"
    elif model_activity >= max(dataset_activity, space_activity):
        url = "https://hub-recap.imglab-cdn.net/images/model.png"
        avatar = "Model Pro"
    elif dataset_activity >= max(model_activity, space_activity):
        url = "https://hub-recap.imglab-cdn.net/images/dataset.png"
        avatar = "Dataset Guru"
    else:
        url = "https://hub-recap.imglab-cdn.net/images/space.png"
        avatar = "Space Artiste"

    # Build text content with proper formatting
    text_parts = []

    text_parts.append(
        f'<span size="10pt" weight="bold">Hugging Face  ❤️ {username} in 2024</span>'
    )
    text_parts.append("")  # Empty line for spacing

    # Stats section
    stats_lines = []
    if total_stats["Model Downloads"] > 0:
        stats_lines.append(
            f'<span weight="bold">{total_stats["Model Downloads"]:,}</span> model downloads'
        )
    if total_stats["Model Likes"] > 0:
        stats_lines.append(
            f'<span weight="bold">{total_stats["Model Likes"]:,}</span> model likes'
        )
    if total_stats["Dataset Downloads"] > 0:
        stats_lines.append(
            f'<span weight="bold">{total_stats["Dataset Downloads"]:,}</span> dataset downloads'
        )
    if total_stats["Dataset Likes"] > 0:
        stats_lines.append(
            f'<span weight="bold">{total_stats["Dataset Likes"]:,}</span> dataset likes'
        )
    if total_stats["Space Likes"] > 0:
        stats_lines.append(
            f'<span weight="bold">{total_stats["Space Likes"]:,}</span> space likes'
        )

    if stats_lines:
        text_parts.extend(stats_lines)
        text_parts.append("")  # Empty line for spacing

    # Popular items section
    top_items = stats["Most Popular Items"]
    if any(
        item["likes"] > 0 or item.get("downloads", 0) > 0 for item in top_items.values()
    ):
        text_parts.append('<span size="10pt">Most Popular Contributions:</span>')

        if top_items["Top Model"]["downloads"] > 0:
            text_parts.append(
                f'Model: <span weight="bold">{top_items["Top Model"]["name"]}</span>'
            )
            text_parts.append(
                f'  ({top_items["Top Model"]["downloads"]:,} downloads, {top_items["Top Model"]["likes"]} likes)'
            )

        if top_items["Top Dataset"]["downloads"] > 0:
            text_parts.append(
                f'Dataset: <span weight="bold">{top_items["Top Dataset"]["name"]}</span>'
            )
            text_parts.append(
                f'  ({top_items["Top Dataset"]["downloads"]:,} downloads, {top_items["Top Dataset"]["likes"]} likes)'
            )

        if top_items["Top Space"]["likes"] > 0:
            text_parts.append(
                f'Space: <span weight="bold">{top_items["Top Space"]["name"]}</span>'
            )
            text_parts.append(f'  ({top_items["Top Space"]["likes"]} likes)')

    # Add avatar type at the end
    text_parts.append("")  # Empty line for spacing
    text_parts.append(f"You are {avatar}! 🎉")

    # Join all parts with newlines
    text = "\n".join(text_parts)

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

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


def is_from_2024(created_at_str):
    if not created_at_str:
        return False
    created_at = datetime.strptime(created_at_str, "%Y-%m-%dT%H:%M:%S.%fZ")
    return created_at.year == 2024


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,
    )
    # Filter for 2024 models only
    models = [
        model
        for model in models_response.json()
        if is_from_2024(model.get("createdAt"))
    ]

    # Get datasets stats
    datasets_response = requests.get(
        "https://huggingface.co/api/datasets",
        params={"author": username, "full": "True"},
        headers=headers,
    )
    # Filter for 2024 datasets only
    datasets = [
        dataset
        for dataset in datasets_response.json()
        if is_from_2024(dataset.get("createdAt"))
    ]

    # Get spaces stats
    spaces_response = requests.get(
        "https://huggingface.co/api/spaces",
        params={"author": username, "full": "True"},
        headers=headers,
    )
    # Filter for 2024 spaces only
    spaces = [
        space
        for space in spaces_response.json()
        if is_from_2024(space.get("createdAt"))
    ]

    # Calculate totals for 2024 items only
    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 from 2024
    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="Hub username",
            placeholder="Enter Hugging Face username...",
            scale=6,
            value="bartowski",
        )
        submit_btn = gr.Button("Get Stats", scale=6)

    with gr.Row():
        with gr.Column():
            stats_image = gr.Markdown(f"![Hugging Face Stats]({DEFAULT_IMAGE})")

    # 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()