import gradio as gr
import requests
import json
import huggingface_hub
from huggingface_hub import HfApi
from gradio_client import Client
import os

HF_TOKEN = os.environ["HF_TOKEN"]
HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"}

tulu = "https://tonic1-tulu.hf.space/--replicas/fxqjz/"


welcome_message = """
Hi! I'm using [Tulu from AlenAi](https://huggingface.co/spaces/Tonic1/Tulu) I'll help you **build a GPT**. You can say something like, "make a bot that gives advice on how to grow your startup."

What would you like to make?
"""

welcome_preview_message = """
Welcome to **{}**! Say something like: 
"{}"
"""

# sample_response = """
# Certainly! Here we go:

# Title: Recipe Recommender
# System Prompt: Utilize your language model abilities to suggest delicious recipes based on user preferences such as ingredients, cuisine type, cooking time, etc. Ensure accuracy and variety while maintaining a conversational style with the user. 
# Example User Input: Vegetarian dinner ideas under 30 minutes
# """

system_prompt = """
I am an AI whose job it is to help users create their own chatbots. In particular, I respond using titles and subtiles in a friendly tone, write a system prompt for an LLM, a catchy title for the chatbot, and a very short example user input. I make sure each part is included.
I only respond in the following format :
# Title:
# System prompt:  
# Example input: 
<|user|>
"make a bot that gives advice on how to grow your startup", 

<|assistant|>
I first do a friendly response, then I add the title, system prompt, and example user input. I Immediately STOP after the example input. It should be EXACTLY in this format:

Sure, I'd be happy to help you build a bot! I'm generating a title, system prompt, and an example input. How do they sound? Feel free to give me feedback!
# Title: Startup Coach
# System prompt: My job as an LLM is to provide good startup advice. Do not provide extraneous comments on other topics. Be succinct but useful. 
# Example input: What are the risks of setting up a non-profit board in my startup?

<|user|>
Make a chatbot that roasts tech ceos

<|assistant|>
Sure, I'd be happy to help you build a bot! I'm generating a title, system prompt, and an example input. How do they sound? Feel free to give me feedback!
# Title: Tech Roaster
# System prompt: As an LLM, my primary function is to deliver hilarious and biting critiques of technology CEOs. I Keep it witty and entertaining, but also make sure my jokes aren't too mean-spirited or factually incorrect. 
# Example input: Roast Elon Musk for me.

<|user|>
Make an app that producesses assessments

<|assistant|>
Sure, I'd be happy to help you build an app! I'm generating a title, system prompt, and an example input. How do they sound? Feel free to give me feedback!
# Title: Assessment Genius
# System prompt: My primary function is to provide assessments for users. These assessments are relevant, useful, and accurate. Keep in mind that I am user-friendly and professional. 
# Example input: I would like a Personality Assessment


<|user|>
make a gpt that helps to create mutants and masterminds characters

<|assistant|>
Sure, I'd be happy to help you build a bot! I'm generating a title, system prompt, and an example input. How do they sound? Feel free to give me feedback!
# Title: Mutants and Masterminds Character Creator
# System prompt: As an LLM, my job is to help users create characters for the Mutants and Masterminds tabletop RPG. My prompts should be clear and concise, and should help users make characters that are both fun and balanced. 
# Example input: I would like to create a character with the Power Level 10
"""

def predict_beta(message, chatbot=[], system_prompt=system_prompt, max_new_tokens=500, temperature=0.4, top_p=0.90, repetition_penalty=0.90, advanced=False):
    client = Client(tulu)
    try:
        result = client.predict(
            message,  
            system_prompt,  
            max_new_tokens,  
            temperature, 
            top_p,  
            repetition_penalty,  
            advanced,  
            fn_index=0
        )
        print("Raw API Response:", result)  # Debugging print
        if result is not None:
            print("Processed bot_message:", result)  # Debugging print
            return result
        else:
            print("No response or empty response from the model.")  # Debugging print
            return None
            
    except Exception as e:
        error_msg = f"An error occurred: {str(e)}"
        print(error_msg)  # Debugging print
        return None
        
def extract_title_prompt_example(text):
    default_title = "Custom GPT Agent"
    default_system_prompt = "This is a custom GPT agent."
    default_example_input = "Type your query here."

    # Split the text into lines and reverse it to start from the end
    lines = text.split('\n')
    lines.reverse()

    title = default_title
    system_prompt = default_system_prompt
    example_input = default_example_input

    # Flags to check if we have found the sections
    found_title, found_prompt, found_example = False, False, False

    for line in lines:
        if not found_example and line.startswith("# Example input:"):
            example_input = line.replace("# Example input:", "").strip()
            found_example = True
        elif not found_prompt and line.startswith("# System prompt:"):
            system_prompt = line.replace("# System prompt:", "").strip()
            found_prompt = True
        elif not found_title and line.startswith("# Title:"):
            title = line.replace("# Title:", "").strip()
            found_title = True

        # Break the loop if all sections are found
        if found_title and found_prompt and found_example:
            break

    return text, title, system_prompt, example_input

def make_open_gpt(message, history, current_title, current_system_prompt, current_example_input, system_prompt=system_prompt):
    try:
        response = predict_beta(message, history, system_prompt)
        if not response:
            raise ValueError("Empty response from predict_beta")
        print("Response from predict_beta:", response)  # Debugging print
    except Exception as e:
        response = f"Error in predict_beta: {str(e)}"
        print("Error in predict_beta:", response)  # Debugging print
        # Set error values
        title = "Error"
        system_prompt = "Error in predict_beta"
        example_input = "Error"
    else:
        try:
            _, title, system_prompt, example_input = extract_title_prompt_example(response)
        except Exception as e:
            title = "Error"
            system_prompt = "Error in extraction"
            example_input = "Error"
            print(f"Error in extract_title_prompt_example: {str(e)}")

    # Ensure all expected outputs are returned
    return (
        "",  # Placeholder for textbox
        history + [(message, response)],  # Updated chatbot history
        title or current_title,  # Extracted or default title
        system_prompt or current_system_prompt,  # Extracted or default system prompt
        example_input or current_example_input,  # Extracted or default example input
        [(None, welcome_preview_message.format(title or current_title, example_input or current_example_input))],  # Updated chatbot preview
        example_input or current_example_input,  # Example input for textbox_preview
        gr.Column(visible=True),  # Column visibility control
        gr.Group(visible=True)  # Group visibility control
    )

    
def set_title_example(title, example):
    return [(None, welcome_preview_message.format(title, example))], example, gr.Column(visible=True), gr.Group(visible=True)

chatbot_preview = gr.Chatbot(layout="panel")
textbox_preview = gr.Textbox(scale=7, container=False)

def test_preview_chatbot(message, history, system_prompt):
    response = predict_beta(message, history, system_prompt)
    return response


def strip_invalid_filename_characters(filename: str, max_bytes: int = 200) -> str:
    """Strips invalid characters from a filename and ensures that the file_length is less than `max_bytes` bytes."""
    filename = filename.replace(" ", "-")
    filename = "".join([char for char in filename if char.isalnum() or char in "_-"])
    filename_len = len(filename.encode())
    if filename_len > max_bytes:
        while filename_len > max_bytes:
            if len(filename) == 0:
                break
            filename = filename[:-1]
            filename_len = len(filename.encode())
    return filename


constants = """
SYSTEM_PROMPT = "{}"
TITLE = "{}"
EXAMPLE_INPUT = "{}"
"""


def publish(textbox_system_prompt, textbox_title, textbox_example, textbox_token):
    source_file = 'app_template.py'
    destination_file = 'app.py'
    constants_formatted = constants.format(textbox_system_prompt, textbox_title, textbox_example)
    with open(source_file, 'r') as file:
        original_content = file.read()
    with open(destination_file, 'w') as file:
        file.write(constants_formatted + original_content)
    title = strip_invalid_filename_characters(textbox_title, max_bytes=30)
    api = HfApi(token=textbox_token)
    new_space = api.create_repo(
        repo_id=f"open-gpt-{title}",
        repo_type="space",
        exist_ok=True,
        private=False,
        space_sdk="gradio",
        token=textbox_token,
    )
    api.upload_file(
        repo_id=new_space.repo_id,
        path_or_fileobj='app.py',
        path_in_repo='app.py',
        token=textbox_token,
        repo_type="space",
    )
    api.upload_file(
        repo_id=new_space.repo_id,
        path_or_fileobj='README_template.md',
        path_in_repo='README.md',
        token=textbox_token,
        repo_type="space",
    )
    huggingface_hub.add_space_secret(
        new_space.repo_id, "HF_TOKEN", textbox_token, token=textbox_token
    )

    return gr.Markdown(f"Published to https://huggingface.co/spaces/{new_space.repo_id} ✅", visible=True), gr.Button("Publish", interactive=True)
    
    
css = """
#preview-tab-button{
    font-weight: bold;
}
"""

with gr.Blocks(css=css) as demo:
    gr.Markdown(""" # 👋🏻Welcome to 🕵🏻‍♂️Agent🌷Tulu
    **A🕵🏻‍♂️Agent🌷Tulu** lets you create your own **open-source GPTs** using [allenai/tulu-2-dpo-13b](https://huggingface.co/allenai/tulu-2-dpo-13b). Start chatting to automatically below to automatically bake your GPT (or you can manually configure the recipe in the second tab). You can build and test them for free & publish them on Spaces (as Open GPTs are powered by the [Tulu DPO model](https://huggingface.co/allenai/tulu-2-dpo-70b) ).
    You think this is cool + want to make your own ? check out [GPTBaker](https://huggingface.co/abidlabs/GPT-Baker) from [AbidLabs](https://huggingface.co/abidlabs) of 🤗[Gradio](https://www.gradio.app/)
    ### Join us: 
    TeamTonic is always making cool demos! Join our active builder's community on Discord: [Discord](https://discord.gg/GWpVpekp) On Huggingface: [TeamTonic](https://huggingface.co/TeamTonic) & [MultiTransformer](https://huggingface.co/MultiTransformer) On Github: [Polytonic](https://github.com/tonic-ai) & contribute to [PolyGPT](https://github.com/tonic-ai/polygpt-alpha) """
               )
    with gr.Row():
        with gr.Column(scale=3):
            with gr.Tab("Create"):
                chatbot_maker = gr.Chatbot([(None, welcome_message)], layout="panel", elem_id="chatbot-maker")
                with gr.Group():
                    with gr.Row():
                        textbox_maker = gr.Textbox(placeholder="Make a bot that roasts tech CEOs", scale=7, container=False, autofocus=True)
                        submit_btn = gr.Button("Bake 👩‍🍳", variant="secondary")
            with gr.Tab("Configure Recipe"):
                textbox_title = gr.Textbox("GPT Preview", label="Title")
                textbox_system_prompt = gr.Textbox(label="System prompt", lines=6)
                textbox_example = gr.Textbox(label="Placeholder example", lines=2)
            with gr.Tab("Files"):
                gr.Markdown("RAG coming soon!")
        with gr.Column(visible=False, scale=5) as preview_column:
            with gr.Tab("🪄 Preview of your Open GPT", elem_id="preview-tab") as preview_tab:
                gr.ChatInterface(test_preview_chatbot, chatbot=chatbot_preview, textbox=textbox_preview, autofocus=False, submit_btn="Test", additional_inputs=[textbox_system_prompt])
    with gr.Group(visible=False) as publish_row:
        with gr.Row():
            textbox_token = gr.Textbox(show_label=False, placeholder="Ready to publish to Spaces? Enter your HF token here", scale=7)
            publish_btn = gr.Button("Publish", variant="primary")

    published_status = gr.Markdown(visible=False)
    
    gr.on([submit_btn.click, textbox_maker.submit], make_open_gpt, [textbox_maker, chatbot_maker, textbox_title, textbox_system_prompt, textbox_example], [textbox_maker, chatbot_maker, textbox_title, textbox_system_prompt, textbox_example, chatbot_preview, textbox_preview, preview_column, publish_row])
    gr.on([textbox_title.blur, textbox_example.blur], set_title_example, [textbox_title, textbox_example], [chatbot_preview, textbox_preview, preview_column, publish_row])

    publish_btn.click(lambda : gr.Button("Publishing...", interactive=False), None, publish_btn).then(publish, [textbox_system_prompt, textbox_title, textbox_example, textbox_token], [published_status, publish_btn])

demo.launch()