from apscheduler.schedulers.background import BackgroundScheduler
import datetime
import os
from typing import Dict, Tuple
from uuid import UUID

import altair as alt
import argilla as rg
from argilla.feedback import FeedbackDataset
from argilla.client.feedback.dataset.remote.dataset import RemoteFeedbackDataset
import gradio as gr
import pandas as pd

# Translation of legends and titels
ANNOTATED = 'Đã dịch'
NUMBER_ANNOTATED = 'Tổng số mẫu đã dịch'
PENDING = 'Số mẫu còn lại'

NUMBER_ANNOTATORS = "Số thành viên tham gia"
NAME = 'Username'
NUMBER_ANNOTATIONS = 'Tỗng số mẫu'

CATEGORY = 'Danh mục'



def obtain_source_target_datasets() -> (
    Tuple[
        FeedbackDataset | RemoteFeedbackDataset, FeedbackDataset | RemoteFeedbackDataset
    ]
):
    """
    This function returns the source and target datasets to be used in the application.

    Returns:
        A tuple with the source and target datasets. The source dataset is filtered by the response status 'pending'.

    """

    # Obtain the public dataset and see how many pending records are there
    source_dataset = rg.FeedbackDataset.from_argilla(
        os.getenv("SOURCE_DATASET"), workspace=os.getenv("SOURCE_WORKSPACE")
    )
    filtered_source_dataset = source_dataset.filter_by(response_status=["pending"])

    # Obtain a list of users from the private workspace
    # target_dataset = rg.FeedbackDataset.from_argilla(
    #    os.getenv("RESULTS_DATASET"), workspace=os.getenv("RESULTS_WORKSPACE")
    # )

    target_dataset = source_dataset.filter_by(response_status=["submitted"])

    return filtered_source_dataset, target_dataset


def get_user_annotations_dictionary(
    dataset: FeedbackDataset | RemoteFeedbackDataset,
) -> Dict[str, int]:
    """
    This function returns a dictionary with the username as the key and the number of annotations as the value.

    Args:
        dataset: The dataset to be analyzed.
    Returns:
        A dictionary with the username as the key and the number of annotations as the value.
    """
    output = {}
    for record in dataset:
        for response in record.responses:
            if str(response.user_id) not in output.keys():
                output[str(response.user_id)] = 1
            else:
                output[str(response.user_id)] += 1

    # Changing the name of the keys, from the id to the username
    for key in list(output.keys()):
        output[rg.User.from_id(UUID(key)).username] = output.pop(key)

    return output


def donut_chart_total() -> alt.Chart:
    """
    This function returns a donut chart with the progress of the total annotations.
    Counts each record that has been annotated at least once.

    Returns:
        An altair chart with the donut chart.
    """

    # Load your data
    annotated_records = len(target_dataset)
    pending_records = int(os.getenv("TARGET_RECORDS")) - annotated_records

    # Prepare data for the donut chart
    source = pd.DataFrame(
        {
            "values": [annotated_records, pending_records],
            "category": [ANNOTATED, PENDING],
            "colors": ["#4CAF50", "#757575"],  # Green for Completed, Grey for Remaining
        }
    )

    base = alt.Chart(source).encode(
        theta=alt.Theta("values:Q", stack=True),
        radius=alt.Radius(
            "values", scale=alt.Scale(type="sqrt", zero=True, rangeMin=20)
        ),
        color=alt.Color("category:N", legend=alt.Legend(title=CATEGORY)),
    )

    c1 = base.mark_arc(innerRadius=20, stroke="#fff")

    c2 = base.mark_text(radiusOffset=20).encode(text="values:Q")

    chart = c1 + c2

    return chart


def donut_chart_target() -> alt.Chart:
    """
    This function returns a donut chart with the progress of the total annotations, in terms of the v1 objective.
    Counts each record that has been annotated at least once.

    Returns:
        An altair chart with the donut chart.
    """

    # Load your data
    annotated_records = len(target_dataset)
    pending_records = int(os.getenv("TARGET_ANNOTATIONS_V1")) - annotated_records

    # Prepare data for the donut chart
    source = pd.DataFrame(
        {
            "values": [annotated_records, pending_records],
            "category": [ANNOTATED, PENDING],
            "colors": ["#4CAF50", "#757575"],  # Green for Completed, Grey for Remaining
        }
    )

    base = alt.Chart(source).encode(
        theta=alt.Theta("values:Q", stack=True),
        radius=alt.Radius(
            "values", scale=alt.Scale(type="sqrt", zero=True, rangeMin=20)
        ),
        color=alt.Color("category:N", legend=alt.Legend(title="Category")),
    )

    c1 = base.mark_arc(innerRadius=20, stroke="#fff")

    c2 = base.mark_text(radiusOffset=20).encode(text="values:Q")

    chart = c1 + c2

    return chart


def kpi_chart_remaining() -> alt.Chart:
    """
    This function returns a KPI chart with the remaining amount of records to be annotated.
    Returns:
        An altair chart with the KPI chart.
    """

    pending_records = int(os.getenv("TARGET_RECORDS")) - len(target_dataset)
    # Assuming you have a DataFrame with user data, create a sample DataFrame
    data = pd.DataFrame({"Category": [PENDING], "Value": [pending_records]})

    # Create Altair chart
    chart = (
        alt.Chart(data)
        .mark_text(fontSize=100, align="center", baseline="middle", color="steelblue")
        .encode(text="Value:N")
        .properties(title=PENDING, width=250, height=200)
    )

    return chart


def kpi_chart_submitted() -> alt.Chart:
    """
    This function returns a KPI chart with the total amount of records that have been annotated.
    Returns:
        An altair chart with the KPI chart.
    """

    total = len(target_dataset)

    # Assuming you have a DataFrame with user data, create a sample DataFrame
    data = pd.DataFrame({"Category": [NUMBER_ANNOTATED], "Value": [total]})

    # Create Altair chart
    chart = (
        alt.Chart(data)
        .mark_text(fontSize=100, align="center", baseline="middle", color="#e68b39")
        .encode(text="Value:N")
        .properties(title=NUMBER_ANNOTATED, width=250, height=200)
    )

    return chart


def kpi_chart() -> alt.Chart:
    """
    This function returns a KPI chart with the total amount of annotators.

    Returns:
        An altair chart with the KPI chart.
    """

    # Obtain the total amount of annotators
    total_annotators = len(user_ids_annotations)

    # Assuming you have a DataFrame with user data, create a sample DataFrame
    data = pd.DataFrame(
        {"Category": [NUMBER_ANNOTATORS], "Value": [total_annotators]}
    )

    # Create Altair chart
    chart = (
        alt.Chart(data)
        .mark_text(fontSize=100, align="center", baseline="middle", color="steelblue")
        .encode(text="Value:N")
        .properties(title=NUMBER_ANNOTATORS, width=250, height=200)
    )

    return chart


def render_hub_user_link(hub_id):
    link = f"https://huggingface.co/{hub_id}"
    return f'<a target="_blank" href="{link}" style="color: var(--link-text-color); text-decoration: underline;text-decoration-style: dotted;">{hub_id}</a>'


def obtain_top_users(user_ids_annotations: Dict[str, int], N: int = 50) -> pd.DataFrame:
    """
    This function returns the top N users with the most annotations.

    Args:
        user_ids_annotations: A dictionary with the user ids as the key and the number of annotations as the value.

    Returns:
        A pandas dataframe with the top N users with the most annotations.
    """

    dataframe = pd.DataFrame(
        user_ids_annotations.items(), columns=[NAME, NUMBER_ANNOTATIONS]
    )
    dataframe[NAME] = dataframe[NAME].apply(render_hub_user_link)
    dataframe = dataframe.sort_values(by=NUMBER_ANNOTATIONS, ascending=False)
    return dataframe.head(N)


def fetch_data() -> None:
    """
    This function fetches the data from the source and target datasets and updates the global variables.
    """

    print(f"Starting to fetch data: {datetime.datetime.now()}")

    global source_dataset, target_dataset, user_ids_annotations, annotated, remaining, percentage_completed, top_dataframe
    source_dataset, target_dataset = obtain_source_target_datasets()
    user_ids_annotations = get_user_annotations_dictionary(target_dataset)

    annotated = len(target_dataset)
    remaining = int(os.getenv("TARGET_RECORDS")) - annotated
    percentage_completed = round(
        (annotated / int(os.getenv("TARGET_RECORDS"))) * 100, 1
    )

    # Print the current date and time
    print(f"Data fetched: {datetime.datetime.now()}")


def get_top(N = 50) -> pd.DataFrame:
    return obtain_top_users(user_ids_annotations, N=N)


def main() -> None:

    # Set the update interval
    update_interval = 300  # seconds
    update_interval_charts = 30  # seconds

    # Connect to the space with rg.init()
    rg.init(
        api_url=os.getenv("ARGILLA_API_URL"),
        api_key=os.getenv("ARGILLA_API_KEY"),
    )

    fetch_data()

    scheduler = BackgroundScheduler()
    scheduler.add_job(
        func=fetch_data, trigger="interval", seconds=update_interval, max_instances=1
    )
    scheduler.start()

    # To avoid the orange border for the Gradio elements that are in constant loading
    css = """
    .generating {
        border: none;
    }
    """

    with gr.Blocks(css=css) as demo:
        gr.Markdown(
            """
            # 🌍 Tiếng Việt - Dự án Đánh giá Prompt Đa ngôn ngữ

            Hugging Face và @argilla đang phát triển dự án [Dự án Đánh giá prompt Đa ngôn ngữ](https://github.com/huggingface/data-is-better-together/tree/main/prompt_translation). Đây là một chuẩn mực mở đa ngôn ngữ để đánh giá các mô hình ngôn ngữ, và tất nhiên, cũng dành cho [NGÔN NGỮ CỦA BẠN].

            ## Mục tiêu là dịch 500 Prompts
            Và như mọi khi: cần có dữ liệu cho việc đó! Cộng đồng đã chọn ra 500 prompt tốt nhất sẽ tạo nên chuẩn mực đánh giá. Bằng tiếng Anh, tất nhiên.
            **Đó là lý do chúng mình cần sự giúp đỡ của bạn**: nếu chúng ta cùng nhau dịch 500 prompts, chúng mình có thể thêm Tiếng Việt vào bảng xếp hạng.

            ## Cách tham gia
            Tham gia dễ dàng. Truy cập vào [AI-Vietnam/prompt-translation-for-vie][https://huggingface.co/spaces/AI-Vietnam/prompt-translation-for-vie], đăng nhập hoặc tạo một tài khoản Hugging Face, và bạn có thể bắt đầu.
            Cảm ơn các bạn rất nhiều! Bên cạnh đó, chúng mình đã dùng AI để chuẩn bị sẵn một đề xuất dịch giúp tăng tốc quá trình dịch thuật.

            # 🌍 Vietnamese - Multilingual Prompt Evaluation Project

            Hugging Face and @argilla are developing [Multilingual Prompt Evaluation Project](https://github.com/huggingface/data-is-better-together/tree/main/prompt_translation) project. It is an open multilingual benchmark for evaluating language models, and of course, also for [YOUR LANGUAGE].

            ## The goal is to translate 500 Prompts
            And as always: data is needed for that! The community selected the best 500 prompts that will form the benchmark. In English, of course.
            **That's why we need your help**: if we all translate the 500 prompts, we can add [YOUR LANGUAGE] to the leaderboard.

            ## How to participate
            Participating is easy. Go to the [annotation space][add a link to your annotation dataset], log in or create a Hugging Face account, and you can start working.
            Thanks in advance! Oh, and we'll give you a little push: Machine translation has already prepared a translation suggestion for you.
            """
        )

        gr.Markdown(
            f"""
            ## 🚀 Tiến độ hiện tại
            Cùng nhau xây dựng bộ dữ liệu này nhé!
            """
        )
        with gr.Row():

            kpi_submitted_plot = gr.Plot(label="Plot")
            demo.load(
                kpi_chart_submitted,
                inputs=[],
                outputs=[kpi_submitted_plot],
                every=update_interval_charts,
            )

            kpi_remaining_plot = gr.Plot(label="Plot")
            demo.load(
                kpi_chart_remaining,
                inputs=[],
                outputs=[kpi_remaining_plot],
                every=update_interval_charts,
            )

            donut_total_plot = gr.Plot(label="Plot")
            demo.load(
                donut_chart_total,
                inputs=[],
                outputs=[donut_total_plot],
                every=update_interval_charts,
            )

        gr.Markdown(
            """
            ## 👾 Bảng xếp hạng
            Tại đây bạn có thể thấy những người đóng góp hàng đầu và số lượng bản dịch họ đã thực hiện:
            """
        )

        with gr.Row():

            kpi_hall_plot = gr.Plot(label="Plot")
            demo.load(
                kpi_chart, inputs=[], outputs=[kpi_hall_plot], every=update_interval_charts
            )

            top_df_plot = gr.Dataframe(
                headers=[NAME, NUMBER_ANNOTATIONS],
                datatype=[
                    "markdown",
                    "number",
                ],
                row_count=50,
                col_count=(2, "fixed"),
                interactive=False,
                every=update_interval,
            )
            demo.load(get_top, None, [top_df_plot], every=update_interval_charts)

    # Launch the Gradio interface
    demo.launch()


if __name__ == "__main__":
    main()