import os
import gradio as gr
import asyncio
import pandas as pd
from datetime import datetime
from langchain.schema.output_parser import StrOutputParser
from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    SystemMessagePromptTemplate,
)

# Load OpenAI model with API key
openai_key = os.getenv('OPENAI_API_KEY')
model = ChatOpenAI(model_name="gpt-4-0125-preview", temperature=0.4)

# System prompt templates
raft_base_template  = os.getenv('RAFT_BASE_TEMPLATE')
raft_base_template_v2  = os.getenv('RAFT_BASE_TEMPLATE_V2')
raft_metaprompt_template = os.getenv('RAFT_METAPROMPT_TEMPLATE')
raft_metaprompt_template_v2 = os.getenv('RAFT_METAPROMPT_TEMPLATE_2')

# System prompt template
customer_template = os.getenv('CUSTOMER_TEMPLATE')

# Initialize DataFrame
results_df = pd.DataFrame(columns=['Experiment Number','Temperature', 'PromptID','Business Idea', 'Generated Domains'])

def extract_content_inside_brackets(json_string):
    # Находим позицию первой открывающей и последней закрывающей скобки
    start = json_string.find('[')
    end = json_string.find(']', start)

    # Если скобки найдены, извлекаем содержимое между ними
    if start != -1 and end != -1:
        content_inside_brackets = json_string[start+1:end]
        return content_inside_brackets
    else:
        # Возвращаем None, если скобки не найдены
        return None
list_of_prompts = ['customer_template', 'raft_base_template', 'raft_metaprompt_template', 'raft_base_template_v2', 'raft_metaprompt_template_v2']
# Function to generate domain names based on different templates
async def generate_domain_names(business_idea, temperature):
    model.temperature = temperature  # Update model temperature
    prompts = [customer_template, raft_base_template, raft_metaprompt_template, raft_base_template_v2, raft_metaprompt_template_v2]
    tasks = []
    for prompt in prompts:
        system_message_prompt = SystemMessagePromptTemplate.from_template(prompt)
        human_message_prompt = HumanMessagePromptTemplate.from_template("User's input: ~###~ {user_input}~###~,")
        combined_prompt = ChatPromptTemplate.from_messages([system_message_prompt, human_message_prompt])
        json_chain = combined_prompt | model | StrOutputParser() | extract_content_inside_brackets
        task = asyncio.to_thread(json_chain.invoke, {"user_input": business_idea})
        tasks.append(task)
    domains = await asyncio.gather(*tasks)
    # Update DataFrame and save to JSON
    experiment_number = datetime.now().strftime("%Y%m%d%H%M%S")
    for i, domain in enumerate(domains):
        new_row = {
            'Experiment Number': experiment_number,
            'Business Idea': business_idea,
            'Temperature': temperature,
            'PromptID': list_of_prompts[i],
            'Generated Domains': domain
        }
        global results_df
        results_df = pd.concat([results_df, pd.DataFrame([new_row])], ignore_index=True)
    save_to_json(results_df)
    return domains

# Function to save DataFrame to JSON
def save_to_json(df):
    json_path = 'domain_name_generation_results.json'
    df.to_json(json_path, orient='records', lines=True)

# Function to download the DataFrame as an Excel file
def download_excel():
    df = pd.read_json('domain_name_generation_results.json', lines=True)
    excel_path = 'domain_name_generation_results.xlsx'
    df.to_excel(excel_path, index=False)
    return excel_path

# Setup Gradio interface
def setup_interface():
    with gr.Blocks() as iface:
        with gr.Row():
            text_input = gr.Textbox(label="Enter your business idea", placeholder="Type here...", lines=2)
            temperature_slider = gr.Slider(minimum=0, maximum=1, step=0.01, label="Set Temperature")
            submit_button = gr.Button("Generate")
        with gr.Row():
            output1 = gr.Textbox(label="Output for Customer Template")
            output2 = gr.Textbox(label="Output for Base Raft Template")
            output3 = gr.Textbox(label="Output for MetaPrompt Raft Template")
            output4 = gr.Textbox(label="Output for Base Raft Template (Новые домены после точки)")
            output5 = gr.Textbox(label="Output for MetaPrompt Raft Template (Новые домены после точки)")
        with gr.Row():
            download_btn = gr.Button("Download Excel")
        submit_button.click(
            fn=generate_domain_names,
            inputs=[text_input, temperature_slider],
            outputs=[output1, output2, output3, output4, output5]
        )
        download_btn.click(
            fn=download_excel,
            inputs=None,
            outputs=gr.File(label="Download Excel")
        )
    return iface


iface = setup_interface()
iface.launch()