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/cdnbn/"


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 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.

<|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: Your job as an LLM is to provide good startup advice. Do not provide extraneous comments on other topics. Be succinct but useful. 
# Example input: Risks of setting up a non-profit board

<|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, your primary function is to deliver hilarious and biting critiques of technology CEOs. Keep it witty and entertaining, but also make sure your jokes aren't too mean-spirited or factually incorrect. 
# Example input: Elon Musk

"""

def predict_beta(message, chatbot=[], system_prompt=system_prompt, max_new_tokens=650, temperature=0.4, top_p=0.90, repetition_penalty=0.90, advanced=True):
    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)
        if result is not None and len(result) > 0:
            bot_message = result[0]
            print(bot_message)
            return bot_message
        else:
            raise gr.Error("No response received from the model.")
            
    except Exception as e:
        error_msg = f"An error occurred: {str(e)}"
        raise gr.Error(error_msg)

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."

    # Find the start indices of each section
    title_start = text.find("# Title:")
    prompt_start = text.find("# System prompt:")
    example_start = text.find("# Example input:")

    # Extract Title
    if title_start != -1:
        title_start += len("# Title:")
        title_end = prompt_start if prompt_start != -1 else len(text)
        title = text[title_start:title_end].strip()
    else:
        title = default_title

    # Extract System Prompt
    if prompt_start != -1:
        prompt_start += len("# System prompt:")
        prompt_end = example_start if example_start != -1 else len(text)
        system_prompt = text[prompt_start:prompt_end].strip()
    else:
        system_prompt = default_system_prompt

    # Extract Example Input
    if example_start != -1:
        example_start += len("# Example input:")
        example_input = text[example_start:].strip().split("\n", 1)[0]
    else:
        example_input = default_example_input

    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):
    print("Response before extraction:", response)  # Print the response before extraction
    response = predict_beta(message, history, system_prompt)
    response, title, system_prompt, example_input = extract_title_prompt_example(response)

    print("Extracted Title:", title)
    print("Extracted System Prompt:", system_prompt)
    print("Extracted Example Input:", example_input)
    # Ensure all expected outputs are returned
    return (
        "",  # Placeholder for textbox
        history + [(message, response)],  # Updated chatbot history
        title,  # Extracted or default title
        system_prompt,  # Extracted or default system prompt
        example_input,  # Extracted or default example input
        [(None, welcome_preview_message.format(title, example_input))],  # Updated chatbot preview
        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()