import gradio as gr
import pandas as pd
from huggingface_hub import HfApi

# Initialize Hugging Face API
api = HfApi()

# Constants
GGUF_TAG = "gguf"
CHUNK_SIZE = 1000

# Clickable links function
def clickable(x, which_one):
    if which_one == "models":
        if x not in ["Not Found", "Unknown"]:
            return f'<a target="_blank" href="https://huggingface.co/{x}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">{x}</a>'
        else:
            return "Not Found"
    else:
        if x != "Not Found":
            return f'<a target="_blank" href="https://huggingface.co/{which_one}/{x}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">{x}</a>'
        return "Not Found"

# Fetch models and return a DataFrame with clickable links
def fetch_models():
    models = api.list_models(filter=GGUF_TAG, full=True)
    data = []
    for model in models:
        model_id = model.id if model.id else "Not Found"
        author = model.author if model.author else "Unknown"
        data.append({
            "Model ID": model_id,
            "Author Name": author,
            "Downloads (30d)": model.downloads or 0,
            "Likes": model.likes or 0,
            "Created At": model.created_at.isoformat() if model.created_at else "N/A",
            "Last Modified": model.last_modified.isoformat() if model.last_modified else "N/A",
        })
    df = pd.DataFrame(data)
    # Apply clickable links
    df["Model ID"] = df["Model ID"].apply(lambda x: clickable(x, "models"))
    df["Author Name"] = df["Author Name"].apply(lambda x: clickable(x, "models"))
    return df

# Prepare authors DataFrame
def prepare_authors_df(models_df):
    # Extract the actual author name from the link (if needed)
    authors_df = models_df.copy()
    authors_df["Clean Author Name"] = authors_df["Author Name"].str.extract(r'href="https://huggingface\.co/(.*?)"')
    grouped = authors_df.groupby("Clean Author Name").agg(
        Models_Count=("Model ID", "count"),
        Total_Downloads=("Downloads (30d)", "sum")
    ).reset_index()
    grouped.rename(columns={"Clean Author Name": "Author Name"}, inplace=True)
    return grouped.sort_values(by="Models_Count", ascending=False)

all_models_df = fetch_models().sort_values(by="Downloads (30d)", ascending=False)
authors_df = prepare_authors_df(all_models_df)

# Calculate totals
total_models_count = len(all_models_df)
total_downloads = all_models_df["Downloads (30d)"].sum()
total_likes = all_models_df["Likes"].sum()

def update_model_table(start_idx):
    new_end = start_idx + CHUNK_SIZE
    combined_df = all_models_df.iloc[:new_end].copy()
    return combined_df, new_end

with gr.Blocks() as demo:
    gr.Markdown(f"""
    # 🚀GGUF Tracker🚀
    Welcome to 🚀**GGUF Tracker**🚀, a live-updating leaderboard for all things GGUF on 🚀Hugging Face.  
    Stats refresh every hour, giving you the latest numbers.

    By the way, I’m 🚀Richard Erkhov, and you can check out more of what I’m working on at my [🌟**github**](https://github.com/RichardErkhov), 
    [🌟**huggingface**](https://huggingface.co/RichardErkhov) or [🌟**erkhov.com**](https://erkhov.com). Go take a look—I think you’ll like what you find.
    """)

    gr.Markdown(f"""
    # GGUF Models and Authors Leaderboard
    **Total Models:** {total_models_count} | **Total Downloads (30d):** {total_downloads} | **Total Likes:** {total_likes}
    """)

    with gr.Tabs():
        with gr.TabItem("Models"):
            model_table = gr.DataFrame(
                value=all_models_df.iloc[:CHUNK_SIZE],
                interactive=True,
                label="GGUF Models",
                wrap=True,
                datatype=["markdown", "markdown", "number", "number", "str", "str"]
            )
            load_more_button = gr.Button("Load More Models")
            start_idx = gr.State(value=CHUNK_SIZE)

            load_more_button.click(fn=update_model_table, inputs=start_idx, outputs=[model_table, start_idx])

        with gr.TabItem("Authors"):
            gr.DataFrame(
                value=authors_df,
                interactive=False,
                label="Authors",
                wrap=True,
                datatype=["str", "number", "number"]
            )

    demo.launch()