"""The UI file for the SynthGenAI package."""

import os
import asyncio
from huggingface_hub import HfFolder

import gradio as gr
from synthgenai import DatasetConfig, DatasetGeneratorConfig, LLMConfig, InstructionDatasetGenerator, PreferenceDatasetGenerator,RawDatasetGenerator,SentimentAnalysisDatasetGenerator, SummarizationDatasetGenerator, TextClassificationDatasetGenerator


def validate_inputs(*args):
    """
    Validate that all required inputs are filled.

    Args:
        *args: The input values to validate.

    Returns:
        bool: True if all required inputs are filled, False otherwise.
    """
    for arg in args:
        if not arg:
            return False
    return True

stop_event = asyncio.Event()

def stop_generation():
    """
    Stop the dataset generation process.
    """
    stop_event.set()

def get_hf_token():
    """
    Retrieve the Hugging Face token from the huggingface_hub.

    Returns:
        str: The Hugging Face token.
    """
    token = HfFolder.get_token()
    if not token:
        raise ValueError("Hugging Face token not found. Please login using the LoginButton.")
    return token

def generate_synthetic_dataset(
    llm_model,
    temperature,
    top_p,
    max_tokens,
    dataset_type,
    topic,
    domains,
    language,
    additional_description,
    num_entries,
    hf_repo_name,
    llm_env_vars,
):
    """
    Generate a dataset based on the provided parameters.

    Args:
        llm_model (str): The LLM model to use.
        temperature (float): The temperature for the LLM.
        top_p (float): The top_p value for the LLM.
        max_tokens (int): The maximum number of tokens for the LLM.
        dataset_type (str): The type of dataset to generate.
        topic (str): The topic of the dataset.
        domains (str): The domains for the dataset.
        language (str): The language of the dataset.
        additional_description (str): Additional description for the dataset.
        num_entries (int): The number of entries in the dataset.
        hf_repo_name (str): The Hugging Face repository name.
        llm_env_vars (str): Comma-separated environment variables for the LLM.

    Returns:
        str: A message indicating the result of the dataset generation.
    """

    hf_token = get_hf_token()
    os.environ["HF_TOKEN"] = hf_token

    for var in llm_env_vars.split(","):
        if "=" in var:
            key, value = var.split("=", 1)
            os.environ[key.strip()] = value.strip()

    # Validate inputs
    if not validate_inputs(
        llm_model,
        temperature,
        top_p,
        max_tokens,
        dataset_type,
        topic,
        domains,
        language,
        num_entries,
        hf_repo_name,
        llm_env_vars,
    ):
        return "All fields except API Base and API Key must be filled."


    llm_config = LLMConfig(
        model=llm_model,
        temperature=temperature,
        top_p=top_p,
        max_tokens=max_tokens,
    )

    dataset_config = DatasetConfig(
        topic=topic,
        domains=domains.split(","),
        language=language,
        additional_description=additional_description,
        num_entries=num_entries,
    )

    dataset_generator_config = DatasetGeneratorConfig(
        llm_config=llm_config,
        dataset_config=dataset_config,
    )

    if dataset_type == "Raw":
        generator = RawDatasetGenerator(dataset_generator_config)
    elif dataset_type == "Instruction":
        generator = InstructionDatasetGenerator(dataset_generator_config)
    elif dataset_type == "Preference":
        generator = PreferenceDatasetGenerator(dataset_generator_config)
    elif dataset_type == "Sentiment Analysis":
        generator = SentimentAnalysisDatasetGenerator(dataset_generator_config)
    elif dataset_type == "Summarization":
        generator = SummarizationDatasetGenerator(dataset_generator_config)
    elif dataset_type == "Text Classification":
        generator = TextClassificationDatasetGenerator(dataset_generator_config)
    else:
        return "Invalid dataset type"

    stop_event.clear()

    async def generate():
        if stop_event.is_set():
            return "Dataset generation stopped."
        dataset = await generator.agenerate_dataset()
        if stop_event.is_set():
            return "Dataset generation stopped."
        dataset.save_dataset(hf_repo_name=hf_repo_name)
        return "Dataset generated and saved successfully."

    return asyncio.run(generate())

def ui_main():
    """
    Launch the Gradio UI for the SynthGenAI dataset generator.
    """
    with gr.Blocks(
        title="SynthGenAI Dataset Generator",
        css="""
            .gradio-container .gr-block {
                margin-bottom: 10px;
            }
        """,
        theme="ParityError/Interstellar",
    ) as demo:
        gr.Markdown(
            """
            <div style="text-align: center;">
                <img src="./assets/logo_header.png" alt="Header Image" style="display: block; margin-left: auto; margin-right: auto; width: 50%;"/>
                <h1>SynthGenAI Dataset Generator</h1>
            </div>

            ## Overview 🧐
            SynthGenAI is designed to be modular and can be easily extended to include different API providers for LLMs and new features.

            ## Why SynthGenAI? 🤔
            Interest in synthetic data generation has surged recently, driven by the growing recognition of data as a critical asset in AI development. Synthetic data generation addresses challenges by allowing us to create diverse and useful datasets using current pre-trained Large Language Models (LLMs).

            [GitHub Repository](https://github.com/Shekswess/synthgenai/tree/main) | [Documentation](https://shekswess.github.io/synthgenai/)

            For more information on which LLMs are allowed and how they can be used, please refer to the [documentation](https://shekswess.github.io/synthgenai/llm_providers/).
            """
        )

        gr.LoginButton(
            value="Login with Hugging Face",
        )

        with gr.Row():
            llm_model = gr.Textbox(
                label="LLM Model", placeholder="model_provider/model_name", value="huggingface/mistralai/Mistral-7B-Instruct-v0.3"
            )
            temperature = gr.Slider(
                label="Temperature", minimum=0.0, maximum=1.0, step=0.1, value=0.5
            )
            top_p = gr.Slider(
                label="Top P", minimum=0.0, maximum=1.0, step=0.1, value=0.9
            )
            max_tokens = gr.Number(label="Max Tokens", value=2048)

        with gr.Row():
            dataset_type = gr.Dropdown(
                label="Dataset Type",
                choices=[
                    "Raw",
                    "Instruction",
                    "Preference",
                    "Sentiment Analysis",
                    "Summarization",
                    "Text Classification",
                ],
            )
            topic = gr.Textbox(label="Topic", placeholder="Dataset topic", value="Artificial Intelligence")
            domains = gr.Textbox(label="Domains", placeholder="Comma-separated domains", value="Machine Learning, Deep Learning")
            language = gr.Textbox(
                label="Language", placeholder="Language", value="English"
            )
            additional_description = gr.Textbox(
                label="Additional Description",
                placeholder="Additional description",
                value="This dataset must be more focused on healthcare implementations of AI, Machine Learning, and Deep Learning.",
            )
            num_entries = gr.Number(label="Number of Entries", value=1000)

        with gr.Row():
            hf_repo_name = gr.Textbox(
                label="Hugging Face Repo Name",
                placeholder="organization_or_user_name/dataset_name",
                value="Shekswess/synthgenai-dataset",
            )
            llm_env_vars = gr.Textbox(
                label="LLM Environment Variables",
                placeholder="Comma-separated environment variables (e.g., KEY1=VALUE1, KEY2=VALUE2)",
                value="HUGGINGFACE_API_KEY=hf_1234566789912345677889, OPENAI_API_KEY=sk-1234566789912345677889",
            )

        generate_button = gr.Button("Generate Dataset")
        stop_button = gr.Button("Stop Generation", visible=False)
        output = gr.Textbox(label="Operation Result", value="")

        def on_generate_click(*args):
            generate_button.visible = False
            stop_button.visible = True
            return generate_synthetic_dataset(*args)

        def on_stop_click():
            stop_generation()
            generate_button.visible = True
            stop_button.visible = False
            return "Dataset generation stopped."

        generate_button.click(
            on_generate_click,
            inputs=[
                llm_model,
                temperature,
                top_p,
                max_tokens,
                dataset_type,
                topic,
                domains,
                language,
                additional_description,
                num_entries,
                hf_repo_name,
                llm_env_vars,
            ],
            outputs=output,
        )

        stop_button.click(
            on_stop_click,
            outputs=output,
        )

    demo.launch(inbrowser=True, favicon_path=None)


if __name__ == "__main__":
    ui_main()