from typing import Any, List
import gradio as gr
from toolz import concat
import httpx
import plotly.express as px
import polars as pl
from pathlib import Path
from datasets import load_dataset
from cachetools import TTLCache, cached
from datetime import datetime, timedelta
from datasets import Dataset
import os
from functools import lru_cache
import pandas as pd
from toolz import frequencies
from dotenv import load_dotenv
from typing import List, Any
from toolz import concat
import httpx
from tqdm.auto import tqdm


load_dotenv()
token = os.environ["HUGGINGFACE_TOKEN"]
user_agent = os.environ["USER_AGENT"]
user = os.environ["USER_TO_TRACK"]
assert token
assert user_agent
assert user

headers = {"user-agent": user_agent, "authorization": f"Bearer {token}"}


def get_hub_community_activity(user: str) -> List[Any]:
    with tqdm() as pbar:
        all_data = []
        i = 1
        while True:
            r = httpx.get(
                f"https://huggingface.co/api/recent-activity?limit=100&type=discussion&skip={i}&user={user}",
                headers=headers,
            )
            activity = r.json()["recentActivity"]
            if not activity:
                break
            all_data.append(activity)
            if len(all_data) % 1000 == 0:
                # print(f"Length of all_data: {len(all_data)}")
                pbar.write(f"Length of all_data: {len(all_data)}")
            i += 100
            pbar.update(100)

    return list(concat(all_data))


# def get_hub_community_activity(user: str) -> List[Any]:
#     all_data = []
#     for i in range(1, 2000, 100):
#         r = httpx.get(
#             f"https://huggingface.co/api/recent-activity?limit=100&type=discussion&skip={i}&user={user}"
#         )
#         activity = r.json()["recentActivity"]
#         all_data.append(activity)
#     return list(concat(all_data))


def parse_date_time(date_time: str) -> datetime:
    return datetime.strptime(date_time, "%Y-%m-%dT%H:%M:%S.%fZ")


def parse_pr_data(data):
    data = data["discussionData"]
    createdAt = parse_date_time(data["createdAt"])
    pr_number = data["num"]
    status = data["status"]
    repo_id = data["repo"]["name"]
    repo_type = data["repo"]["type"]
    isPullRequest = data["isPullRequest"]
    return {
        "createdAt": createdAt,
        "pr_number": pr_number,
        "status": status,
        "repo_id": repo_id,
        "type": repo_type,
        "isPullRequest": isPullRequest,
    }


@cached(cache=TTLCache(maxsize=1000, ttl=timedelta(minutes=30), timer=datetime.now))
def update_data():
    try:
        previous_df = pl.DataFrame(
            load_dataset(f"librarian-bot/{user}-stats", split="train").data.table
        )
    except FileNotFoundError:
        previous_df = pl.DataFrame()
    data = get_hub_community_activity(user)
    data = [parse_pr_data(d) for d in data]
    update_df = pl.DataFrame(data)
    df = pl.concat([previous_df, update_df]).unique()
    if len(df) != len(previous_df):
        Dataset(df.to_arrow()).push_to_hub(f"{user}-stats", token=token)
    return df


# def get_pr_status():
#     df = update_data()
#     df = df.filter(pl.col("isPullRequest") is True)
#     return df.select(pl.col("status").value_counts())
#     # return frequencies(x["status"] for x in pr_data)


@lru_cache(maxsize=512)
def get_pr_status(user: str):
    all_data = get_hub_community_activity(user)
    pr_data = (
        x["discussionData"] for x in all_data if x["discussionData"]["isPullRequest"]
    )
    return frequencies(x["status"] for x in pr_data)


def create_pie():
    frequencies = get_pr_status(user)
    df = pd.DataFrame({"status": frequencies.keys(), "number": frequencies.values()})
    return px.pie(df, values="number", names="status", template="seaborn")


def group_status_by_pr_number():
    all_data = get_hub_community_activity(user)
    all_data = [parse_pr_data(d) for d in all_data]
    return (
        pl.DataFrame(all_data).groupby("status").agg(pl.mean("pr_number")).to_pandas()
    )


def plot_over_time():
    all_data = get_hub_community_activity(user)
    all_data = [parse_pr_data(d) for d in all_data]
    df = pl.DataFrame(all_data).with_columns(pl.col("createdAt").cast(pl.Date))
    df = df.pivot(
        values=["status"],
        index=["createdAt"],
        columns=["status"],
        aggregate_function="count",
    )
    df = df.fill_null(0)
    df = df.with_columns(pl.sum(["open", "closed", "merged"])).sort("createdAt")
    df = df.to_pandas().set_index("createdAt").cumsum()
    return px.line(df, x=df.index, y=[c for c in df.columns if c != "sum"])


create_pie()

with gr.Blocks() as demo:
    # frequencies = get_pr_status("librarian-bot")
    gr.Markdown(f"# {user} PR Stats")
    gr.Markdown(f"Total prs and issues opened by {user}: {len(update_data()):,}")
    # gr.Markdown(f"Total PRs opened: {sum(frequencies.values())}")
    with gr.Column():
        gr.Markdown("## Pull requests status")
        gr.Markdown(
            "The below pie chart shows the percentage of pull requests made by"
            " librarian bot that are open, closed or merged"
        )
        gr.Plot(create_pie())
    with gr.Column():
        gr.Markdown("Pull requests opened, closed and merged over time (cumulative)")
        gr.Plot(plot_over_time())
    with gr.Column():
        gr.Markdown("## Pull requests status by PR number")
        gr.DataFrame(group_status_by_pr_number())
demo.launch(debug=True)