import gradio as gr
from openai import OpenAI
import os

# Retrieve the access token from the environment variable
ACCESS_TOKEN = os.getenv("HF_TOKEN")
print("Access token loaded.")

# Initialize the OpenAI client with the Hugging Face Inference API endpoint
client = OpenAI(
    base_url="https://api-inference.huggingface.co/v1/",
    api_key=ACCESS_TOKEN,
)
print("OpenAI client initialized.")

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    frequency_penalty,
    seed,
    model
):
    """
    This function handles the chatbot response. It takes in:
    - message: the user's new message
    - history: the list of previous messages, each as a tuple (user_msg, assistant_msg)
    - system_message: the system prompt
    - max_tokens: the maximum number of tokens to generate in the response
    - temperature: sampling temperature
    - top_p: top-p (nucleus) sampling
    - frequency_penalty: penalize repeated tokens in the output
    - seed: a fixed seed for reproducibility; -1 will mean 'random'
    - model: the model to use for text generation
    """

    print(f"Received message: {message}")
    print(f"History: {history}")
    print(f"System message: {system_message}")
    print(f"Max tokens: {max_tokens}, Temperature: {temperature}, Top-P: {top_p}")
    print(f"Frequency Penalty: {frequency_penalty}, Seed: {seed}")
    print(f"Model: {model}")

    # Convert seed to None if -1 (meaning random)
    if seed == -1:
        seed = None

    # Construct the messages array required by the API
    messages = [{"role": "system", "content": system_message}]

    # Add conversation history to the context
    for val in history:
        user_part = val[0]
        assistant_part = val[1]
        if user_part:
            messages.append({"role": "user", "content": user_part})
            print(f"Added user message to context: {user_part}")
        if assistant_part:
            messages.append({"role": "assistant", "content": assistant_part})
            print(f"Added assistant message to context: {assistant_part}")

    # Append the latest user message
    messages.append({"role": "user", "content": message})

    # Start with an empty string to build the response as tokens stream in
    response = ""
    print("Sending request to OpenAI API.")

    # Make the streaming request to the HF Inference API via openai-like client
    for message_chunk in client.chat.completions.create(
        model=model,   # Use the selected model
        max_tokens=max_tokens,
        stream=True,  # Stream the response
        temperature=temperature,
        top_p=top_p,
        frequency_penalty=frequency_penalty,  # <-- NEW
        seed=seed,                             # <-- NEW
        messages=messages,
    ):
        # Extract the token text from the response chunk
        token_text = message_chunk.choices[0].delta.content
        print(f"Received token: {token_text}")
        response += token_text
        yield response

    print("Completed response generation.")

# Create a Chatbot component with a specified height
chatbot = gr.Chatbot(height=600)
print("Chatbot interface created.")

# Create the Gradio ChatInterface
# We add two new sliders for Frequency Penalty and Seed
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="", label="System message"),
        gr.Slider(minimum=1,   maximum=4096, value=512, step=1,   label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0,  value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0,  value=0.95, step=0.05, label="Top-P"),
        gr.Slider(
            minimum=-2.0,
            maximum=2.0,
            value=0.0,
            step=0.1,
            label="Frequency Penalty"
        ),
        gr.Slider(
            minimum=-1,
            maximum=65535,  # Arbitrary upper limit for demonstration
            value=-1,
            step=1,
            label="Seed (-1 for random)"
        ),
        gr.Textbox(label="Custom Model", placeholder="Enter a custom model path"),
    ],
    fill_height=True,
    chatbot=chatbot,
    theme="Nymbo/Nymbo_Theme",
)
print("Gradio interface initialized.")

# Define the Gradio interface
with gr.Blocks(theme='Nymbo/Nymbo_Theme_5') as textgen:
    # Tab for basic settings
    with gr.Tab("Basic Settings"):
        with gr.Row():
            with gr.Column(elem_id="prompt-container"):
                with gr.Row():
                    # Textbox for user to input the prompt
                    text_prompt = gr.Textbox(label="Prompt", placeholder="Enter a prompt here", lines=3, elem_id="prompt-text-input")
                with gr.Row():
                    # Textbox for custom model input
                    custom_model = gr.Textbox(label="Custom Model", info="Model Hugging Face path (optional)", placeholder="meta-llama/Llama-3.3-70B-Instruct")
                with gr.Row():
                    # Accordion for selecting the model
                    with gr.Accordion("Featured Models", open=True):
                        # Textbox for searching models
                        model_search = gr.Textbox(label="Filter Models", placeholder="Search for a featured model...", lines=1, elem_id="model-search-input")
                        models_list = (
                            "meta-llama/Llama-3.3-70B-Instruct",
                            "meta-llama/Llama-3.3-13B-Instruct",
                            "meta-llama/Llama-3.3-30B-Instruct",
                            "meta-llama/Llama-3.3-7B-Instruct",
                        )

                        # Radio buttons to select the desired model
                        model = gr.Radio(label="Select a model below", value="meta-llama/Llama-3.3-70B-Instruct", choices=models_list, interactive=True, elem_id="model-radio")

                        # Filtering models based on search input
                        def filter_models(search_term):
                            filtered_models = [m for m in models_list if search_term.lower() in m.lower()]
                            return gr.update(choices=filtered_models)

                        # Update model list when search box is used
                        model_search.change(filter_models, inputs=model_search, outputs=model)

    # Tab for advanced settings
    with gr.Tab("Advanced Settings"):
        with gr.Row():
            # Slider for setting the maximum number of new tokens
            max_tokens = gr.Slider(label="Max new tokens", value=512, minimum=1, maximum=4096, step=1)
        with gr.Row():
            # Slider for setting the temperature
            temperature = gr.Slider(label="Temperature", value=0.7, minimum=0.1, maximum=4.0, step=0.1)
        with gr.Row():
            # Slider for setting the top-p (nucleus) sampling
            top_p = gr.Slider(label="Top-P", value=0.95, minimum=0.1, maximum=1.0, step=0.05)
        with gr.Row():
            # Slider for setting the frequency penalty
            frequency_penalty = gr.Slider(label="Frequency Penalty", value=0.0, minimum=-2.0, maximum=2.0, step=0.1)
        with gr.Row():
            # Slider for setting the seed for reproducibility
            seed = gr.Slider(label="Seed", value=-1, minimum=-1, maximum=65535, step=1)

    # Tab to provide information to the user
    with gr.Tab("Information"):
        with gr.Row():
            # Display a sample prompt for guidance
            gr.Textbox(label="Sample prompt", value="{prompt} | ultra detail, ultra elaboration, ultra quality, perfect.")

        # Accordion displaying featured models
        with gr.Accordion("Featured Models (WiP)", open=False):
            gr.HTML(
                """
            <p><a href="https://huggingface.co/models?inference=warm&pipeline_tag=text-generation&sort=trending">See all available models</a></p>
            <table style="width:100%; text-align:center; margin:auto;">
                <tr>
                    <th>Model Name</th>
                    <th>Typography</th>
                    <th>Notes</th>
                </tr>
                <tr>
                    <td>meta-llama/Llama-3.3-70B-Instruct</td>
                    <td>✅</td>
                    <td></td>
                </tr>
                <tr>
                    <td>meta-llama/Llama-3.3-13B-Instruct</td>
                    <td>✅</td>
                    <td></td>
                </tr>
                <tr>
                    <td>meta-llama/Llama-3.3-30B-Instruct</td>
                    <td>✅</td>
                    <td></td>
                </tr>
                <tr>
                    <td>meta-llama/Llama-3.3-7B-Instruct</td>
                    <td>✅</td>
                    <td></td>
                </tr>
            </table>
            """
            )

        # Accordion providing an overview of advanced settings
        with gr.Accordion("Parameters Overview", open=False):
            gr.Markdown(
            """
            ## Max New Tokens
            ###### This slider allows you to specify the maximum number of tokens to generate in the response. The default value is 512, and the maximum output is 4096.

            ## Temperature
            ###### The temperature controls the randomness of the output. A higher temperature makes the output more random, while a lower temperature makes it more deterministic. The default value is 0.7.

            ## Top-P
            ###### Top-P (nucleus) sampling is a way to control the diversity of the output. A higher value allows for more diverse outputs, while a lower value makes the output more focused. The default value is 0.95.

            ## Frequency Penalty
            ###### The frequency penalty penalizes repeated tokens in the output. A higher value makes the output more diverse, while a lower value allows for more repetition. The default value is 0.0.

            ## Seed
            ###### The seed is a fixed value for reproducibility. If you find a seed that gives you a result you love, you can use it again to create a similar output. If you leave it at -1, the AI will generate a new seed every time.

            ### Remember, these settings are all about giving you control over the text generation process. Feel free to experiment and see what each one does. And if you're ever in doubt, the default settings are a great place to start. Happy creating!
            """
            )

    # Row containing the 'Run' button to trigger the text generation
    with gr.Row():
        text_button = gr.Button("Run", variant='primary', elem_id="gen-button")
    # Row for displaying the generated text output
    with gr.Row():
        text_output = gr.Textbox(label="Text Output", elem_id="text-output")

    # Set up button click event to call the respond function
    text_button.click(respond, inputs=[text_prompt, chatbot, gr.Textbox(value="", label="System message"), max_tokens, temperature, top_p, frequency_penalty, seed, model], outputs=text_output)

print("Launching Gradio interface...")  # Debug log
# Launch the Gradio interface without showing the API or sharing externally
textgen.launch(show_api=False, share=False)