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 x in ["Not Found", "Unknown"]:
        return "Not Found"
    if which_one == "models":
        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 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>'

# 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 to models and authors
    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):
    authors_df = models_df.copy()
    # Extract the author name from the href in the clickable link
    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"),
        Total_Likes=("Likes", "sum")
    ).reset_index()

    grouped.rename(columns={"Clean Author Name": "Author Name"}, inplace=True)
    grouped["Author Name"] = grouped["Author Name"].apply(lambda x: clickable(x, "models"))
    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 apply_model_filters(search_query, min_downloads, min_likes):
    df = all_models_df.copy()

    # Extract visible text for filtering purposes:
    visible_model_id = df["Model ID"].str.extract(r'>(.*?)<')[0]
    visible_author_name = df["Author Name"].str.extract(r'>(.*?)<')[0]

    # Search filter
    if search_query.strip():
        mask = (visible_model_id.str.contains(search_query, case=False, na=False)) | \
               (visible_author_name.str.contains(search_query, case=False, na=False))
        df = df[mask]

    # Minimum downloads filter
    if min_downloads is not None and min_downloads > 0:
        df = df[df["Downloads (30d)"] >= min_downloads]

    # Minimum likes filter
    if min_likes is not None and min_likes > 0:
        df = df[df["Likes"] >= min_likes]

    return df

def filter_models(search_query, min_downloads, min_likes):
    filtered = apply_model_filters(search_query, min_downloads, min_likes)
    return filtered.iloc[:CHUNK_SIZE], CHUNK_SIZE, filtered

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

def apply_author_filters(search_query, min_author_downloads, min_author_likes):
    df = authors_df.copy()

    # Extract visible text for author filtering:
    visible_author_name = df["Author Name"].str.extract(r'>(.*?)<')[0]

    # Search filter for authors
    if search_query.strip():
        mask = visible_author_name.str.contains(search_query, case=False, na=False)
        df = df[mask]

    # Minimum total downloads filter
    if min_author_downloads is not None and min_author_downloads > 0:
        df = df[df["Total_Downloads"] >= min_author_downloads]

    # Minimum total likes filter
    if min_author_likes is not None and min_author_likes > 0:
        df = df[df["Total_Likes"] >= min_author_likes]

    return df


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"):
            with gr.Row():
                search_query = gr.Textbox(label="Search (by Model ID or Author Name)")
                min_downloads = gr.Number(label="Min Downloads (30d)", value=0)
                min_likes = gr.Number(label="Min Likes", value=0)

            filter_button = gr.Button("Apply Filters")
            model_table = gr.DataFrame(
                value=all_models_df.iloc[:CHUNK_SIZE],
                interactive=False,
                label="GGUF Models (Click column headers to sort)",
                wrap=True,
                datatype=["markdown", "markdown", "number", "number", "str", "str"]
            )
            load_more_button = gr.Button("Load More Models")
            
            # States
            start_idx = gr.State(value=CHUNK_SIZE)
            filtered_df_state = gr.State(value=all_models_df)  # holds the currently filtered df

            filter_button.click(
                fn=filter_models,
                inputs=[search_query, min_downloads, min_likes],
                outputs=[model_table, start_idx, filtered_df_state]
            )
            load_more_button.click(fn=update_model_table, inputs=[start_idx, filtered_df_state], outputs=[model_table, start_idx])

        with gr.TabItem("Authors"):
            with gr.Row():
                author_search_query = gr.Textbox(label="Search by Author Name")
                min_author_downloads = gr.Number(label="Min Total Downloads", value=0)
                min_author_likes = gr.Number(label="Min Total Likes", value=0)

            author_filter_button = gr.Button("Apply Filters")
            author_table = gr.DataFrame(
                value=authors_df,
                interactive=False,
                label="Authors (Click column headers to sort)",
                wrap=True,
                datatype=["markdown", "number", "number", "number"]
            )

            def filter_authors(author_search_query, min_author_downloads, min_author_likes):
                filtered_authors = apply_author_filters(author_search_query, min_author_downloads, min_author_likes)
                return filtered_authors

            author_filter_button.click(
                fn=filter_authors,
                inputs=[author_search_query, min_author_downloads, min_author_likes],
                outputs=author_table
            )

    demo.launch()