import gradio as gr
import requests
import os
import json
import pandas as pd
import time
from langchain.schema import SystemMessage
from langchain_community.chat_models.gigachat import GigaChat
from openpyxl import load_workbook
import base64

# Установка ключей API для OpenAI, GigaChat и GitHub
openai_api_key = os.getenv('GPT_KEY')
gc_key = os.getenv('GC_KEY')
token = os.getenv('GITHUB_TOKEN')

# Класс конфигурации модели
class ModelConfig:
    def __init__(self, temperature):
        self.max_tokens = 68  # фиксированное значение
        self.temperature = temperature

# Класс для генерации сообщений
class MessageGenerator:
    def __init__(self, config: ModelConfig):
        self.config = config
        self.chat_pro = self.authenticate_gigachat('GigaChat-Pro')
        self.chat_lite = self.authenticate_gigachat('GigaChat')
        self.chat_plus = self.authenticate_gigachat('GigaChat-Plus')

    def authenticate_gigachat(self, model):
        return GigaChat(credentials=gc_key, model=model, max_tokens=self.config.max_tokens, temperature=self.config.temperature, verify_ssl_certs=False)

    def clean_message(self, message):
        if not message.endswith(('.', '!', '?')):
            last_period = message.rfind('.')
            if last_period != -1:
                message = message[:last_period + 1]
        return message

    def generate_message_gpt4o(self, prompt):
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {openai_api_key}"
        }
        data = {
            "model": "chatgpt-4o-latest",
            "messages": [{"role": "system", "content": prompt}],
            "max_tokens": self.config.max_tokens,
            "temperature": self.config.temperature
        }
        response = requests.post("https://api.openai.com/v1/chat/completions", json=data, headers=headers)
        response_data = response.json()
        return self.clean_message(response_data["choices"][0]["message"]["content"].strip())

    def generate_message_gigachat(self, prompt, model):
        messages = [SystemMessage(content=prompt)]
        if model == 'GigaChat-Pro':
            res = self.chat_pro(messages)
        elif model == 'GigaChat':
            res = self.chat_lite(messages)
        elif model == 'GigaChat-Plus':
            res = self.chat_plus(messages)
        return self.clean_message(res.content.strip())

# Загрузка данных из Excel-файла
data = pd.read_excel('Признаки.xlsx', sheet_name=None)
features = {sheet_name: df.set_index(df.columns[0]).to_dict()[df.columns[1]] for sheet_name, df in data.items()}

# Функция для генерации стандартного промпта
def generate_standard_prompt(description, advantages, *selected_values):
    prompt = (
        "Сгенерируй смс-сообщение для клиента.\n"
        "Начни сообщение с призыва к действию с продуктом.\n"
        f"Описание предложения: {description}\n"
        f"Преимущества: {advantages}\n"
        "Вклад на короткий срок.\n"
        "В тексте смс запрещено использование:\n"
        "- Запрещенные слова: № один, номер один, № 1, вкусный, дешёвый, продукт, спам, доступный, банкротство, долги, займ, срочно, сейчас, лучший, главный, номер 1, гарантия, успех, лидер;\n"
        "- Обращение к клиенту;\n"
        "- Приветствие клиента;\n"
        "- Обещания и гарантии;\n"
        "- Использовать составные конструкции из двух глаголов;\n"
        "- Причастия и причастные обороты;\n"
        "- Деепричастия и деепричастные обороты;\n"
        "- Превосходная степень прилагательных;\n"
        "- Страдательный залог;\n"
        "- Порядковые числительные от 10 прописью;\n"
        "- Цепочки с придаточными предложениями;\n"
        "- Разделительные повторяющиеся союзы;\n"
        "- Вводные конструкции;\n"
        "- Усилители;\n"
        "- Паразиты времени;\n"
        "- Несколько существительных подряд, в том числе отглагольных;\n"
        "- Производные предлоги;\n"
        "- Сложные предложения, в которых нет связи между частями;\n"
        "- Сложноподчинённые предложения;\n"
        "- Даты прописью;\n"
        "- Близкие по смыслу однородные члены предложения;\n"
        "- Шокирующие, экстравагантные, кликбейтные фразы;\n"
        "- Абстрактные заявления без поддержки фактами и отсутствие доказательства пользы для клиента;\n"
        "- Гарантирующие фразы;\n"
        "- Узкоспециализированные термины;\n"
        "- Фразы, способные создать двойственное ощущение, обидеть;\n"
        "- Речевые клише, рекламные штампы, канцеляризмы;\n"
        "Убедись, что в готовом тексте до 250 знаков с пробелами."
    )
    return prompt

# Функция для генерации сообщений с учетом повторных попыток
def generate_message_with_retry(generator, prompt, model):
    for _ in range(10):
        message = generator.generate_message_gigachat(prompt, model)
        if len(message) <= 250:
            return message
    return message  # Возвращаем последнее сгенерированное сообщение, если все попытки не удались

# Обновление интерфейса для отображения сообщений
def generate_messages(description, advantages, config, *selected_values):
    generator = MessageGenerator(config)
    standard_prompt = generate_standard_prompt(description, advantages, *selected_values)

    results = {
        "prompt": standard_prompt,
        "gpt4o": generator.generate_message_gpt4o(standard_prompt),
        "gigachat_pro": generate_message_with_retry(generator, standard_prompt, 'GigaChat-Pro'),
        "gigachat_lite": generate_message_with_retry(generator, standard_prompt, 'GigaChat'),
        "gigachat_plus": generate_message_with_retry(generator, standard_prompt, 'GigaChat-Plus')
    }

    return results["prompt"], results["gpt4o"], results["gigachat_pro"], results["gigachat_lite"], results["gigachat_plus"]

# Функция для генерации персонализированного промпта
def generate_personalization_prompt(*selected_values):
    prompt = "Адаптируй, не превышая длину сообщения в 250 знаков с пробелами, текст с учетом следующих особенностей:\n"
    for i, feature in enumerate(features.keys()):
        if selected_values[i]:
            prompt += f"{features[feature][selected_values[i]]}\n"
    prompt += "Убедись, что в готовом тексте до 250 знаков с пробелами."
    return prompt.strip()

# Функция для выполнения персонализации на основе сгенерированного промпта и сообщения
def perform_personalization(standard_message, personalization_prompt, generator):
    full_prompt = f"{personalization_prompt}\n\nТекст для адаптации:\n{standard_message}"
    return generator.generate_message_gpt4o(full_prompt)

# Функция для выполнения персонализации с учетом повторных попыток
def perform_personalization_with_retry(standard_message, personalization_prompt, generator, model):
    for _ in range(10):
        message = perform_personalization(standard_message, personalization_prompt, generator)
        if len(message) <= 250:
            return message
    return message

# Обновление интерфейса для персонализации сообщений
def personalize_messages(gpt4o_message, gigachat_pro_message, gigachat_lite_message, gigachat_plus_message, config, *selected_values):
    generator = MessageGenerator(config)
    personalization_prompt = generate_personalization_prompt(*selected_values)

    personalized_gpt4o = perform_personalization_with_retry(gpt4o_message, personalization_prompt, generator, "GPT-4o")
    personalized_gigachat_pro = perform_personalization_with_retry(gigachat_pro_message, personalization_prompt, generator, "GigaChat-Pro")
    personalized_gigachat_lite = perform_personalization_with_retry(gigachat_lite_message, personalization_prompt, generator, "GigaChat")
    personalized_gigachat_plus = perform_personalization_with_retry(gigachat_plus_message, personalization_prompt, generator, "GigaChat-Plus")

    return personalization_prompt, personalized_gpt4o, personalized_gigachat_pro, personalized_gigachat_lite, personalized_gigachat_plus

# Функция для сохранения данных на GitHub
def save_to_github(personalized_message, model_name, comment, corrected_message, description, advantages, non_personalized_prompt, non_personalized_message, gender, generation, psychotype, business_stage, industry, legal_form):
    data_to_save = {
        "Модель": model_name,
        "Персонализированное сообщение": personalized_message,
        "Комментарий": comment,
        "Откорректированное сообщение": corrected_message,
        "Описание предложения": description,
        "Преимущества": advantages,
        "Неперсонализированный промпт": non_personalized_prompt,
        "Неперсонализированное сообщение": non_personalized_message,
        "Пол": gender,
        "Поколение": generation,
        "Психотип": psychotype,
        "Стадия бизнеса": business_stage,
        "Отрасль": industry,
        "ОПФ": legal_form
    }
    file_content_encoded = base64.b64encode(json.dumps(data_to_save).encode()).decode()
    repo = "fruitpicker01/Storage_1"
    path = f"file_{int(time.time())}.json"
    url = f"https://api.github.com/repos/{repo}/contents/{path}"
    headers = {
        "Authorization": f"token {token}",
        "Content-Type": "application/json"
    }
    data = {
        "message": f"Добавлен новый файл {path}",
        "content": file_content_encoded
    }
    requests.put(url, headers=headers, data=json.dumps(data))

# Создание интерфейса Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Генерация SMS-сообщений по заданным признакам")
    
    with gr.Row():
        with gr.Column(scale=1):
            temperature_slider = gr.Slider(label="Температура", minimum=0, maximum=2, step=0.01, value=1)

            description_input = gr.Textbox(
                label="Описание предложения", 
                lines=5,
                value="Необходимо предложить клиенту оформить дебетовую премиальную бизнес-карту Mastercard Preffered..."
            )
            advantages_input = gr.Textbox(
                label="Преимущества", 
                lines=5,
                value="Предложение по бесплатному обслуживанию — бессрочное..."
            )

            selections = []
            for feature in features.keys():
                selections.append(gr.Dropdown(choices=[None] + list(features[feature].keys()), label=f"Выберите {feature}"))

            submit_btn = gr.Button("Создать неперсонализированное сообщение")

        with gr.Column(scale=2):
            prompt_display = gr.Textbox(label="Неперсонализированный промпт", lines=10, interactive=False)
            output_text_gpt4o = gr.Textbox(label="Неперсонализированное сообщение GPT-4o", lines=3, interactive=False)
            output_text_gigachat_pro = gr.Textbox(label="Неперсонализированное сообщение GigaChat-Pro", lines=3, interactive=False)
            output_text_gigachat_lite = gr.Textbox(label="Неперсонализированное сообщение GigaChat-Lite", lines=3, interactive=False)
            output_text_gigachat_plus = gr.Textbox(label="Неперсонализированное сообщение GigaChat-Plus", lines=3, interactive=False)

            submit_btn.click(
                fn=generate_messages, 
                inputs=[description_input, advantages_input, gr.State(lambda: ModelConfig(temperature_slider.value))] + selections, 
                outputs=[prompt_display, output_text_gpt4o, output_text_gigachat_pro, output_text_gigachat_lite, output_text_gigachat_plus]
            )

    with gr.Row():
        personalize_btn = gr.Button("Выполнить персонализацию")

        with gr.Column(scale=2):
            personalization_prompt_display = gr.Textbox(label="Промпт для персонализации", lines=6, interactive=False)
            personalized_output_text_gpt4o = gr.Textbox(label="Персонализированное сообщение GPT-4o", lines=6, interactive=False)
            personalized_output_text_gigachat_pro = gr.Textbox(label="Персонализированное сообщение GigaChat-Pro", lines=6, interactive=False)
            personalized_output_text_gigachat_lite = gr.Textbox(label="Персонализированное сообщение GigaChat-Lite", lines=6, interactive=False)
            personalized_output_text_gigachat_plus = gr.Textbox(label="Персонализированное сообщение GigaChat-Plus", lines=6, interactive=False)

            personalize_btn.click(
                fn=personalize_messages,
                inputs=[output_text_gpt4o, output_text_gigachat_pro, output_text_gigachat_lite, output_text_gigachat_plus, gr.State(lambda: ModelConfig(temperature_slider.value))] + selections,
                outputs=[personalization_prompt_display, personalized_output_text_gpt4o, personalized_output_text_gigachat_pro, personalized_output_text_gigachat_lite, personalized_output_text_gigachat_plus]
            )

    with gr.Row():
        gr.Markdown("*Комментарий (опционально):*")
        comment_gpt4o = gr.Textbox(label="", lines=3)
        comment_gigachat_pro = gr.Textbox(label="", lines=3)
        comment_gigachat_lite = gr.Textbox(label="", lines=3)
        comment_gigachat_plus = gr.Textbox(label="", lines=3)

    with gr.Row():
        gr.Markdown("*Откорректированное сообщение (опционально):*")
        corrected_gpt4o = gr.Textbox(label="", lines=3)
        corrected_gigachat_pro = gr.Textbox(label="", lines=3)
        corrected_gigachat_lite = gr.Textbox(label="", lines=3)
        corrected_gigachat_plus = gr.Textbox(label="", lines=3)

    with gr.Row():
        gr.Button("Жми 👍 для сохранения удачного SMS в базу =>")
        save_gpt4o_btn = gr.Button("👍")
        save_gigachat_pro_btn = gr.Button("👍")
        save_gigachat_lite_btn = gr.Button("👍")
        save_gigachat_plus_btn = gr.Button("👍")

        save_gpt4o_btn.click(
            fn=lambda personalized_message, comment, corrected_message, description, advantages, non_personalized_prompt, non_personalized_message, gender, generation, psychotype, business_stage, industry, legal_form: 
                save_to_github(personalized_message, "GPT-4o", comment, corrected_message, description, advantages, non_personalized_prompt, non_personalized_message, gender, generation, psychotype, business_stage, industry, legal_form),
            inputs=[
                personalized_output_text_gpt4o,
                comment_gpt4o,
                corrected_gpt4o,
                description_input,
                advantages_input,
                prompt_display,
                output_text_gpt4o,
                selections[0],  # Пол
                selections[1],  # Поколение
                selections[2],  # Психотип
                selections[3],  # Стадия бизнеса
                selections[4],  # Отрасль
                selections[5]   # ОПФ
            ],
            outputs=None
        )

        save_gigachat_pro_btn.click(
            fn=lambda personalized_message, comment, corrected_message, description, advantages, non_personalized_prompt, non_personalized_message, gender, generation, psychotype, business_stage, industry, legal_form: 
                save_to_github(personalized_message, "GigaChat-Pro", comment, corrected_message, description, advantages, non_personalized_prompt, non_personalized_message, gender, generation, psychotype, business_stage, industry, legal_form),
            inputs=[
                personalized_output_text_gigachat_pro,
                comment_gigachat_pro,
                corrected_gigachat_pro,
                description_input,
                advantages_input,
                prompt_display,
                output_text_gigachat_pro,
                selections[0],  # Пол
                selections[1],  # Поколение
                selections[2],  # Психотип
                selections[3],  # Стадия бизнеса
                selections[4],  # Отрасль
                selections[5]   # ОПФ
            ],
            outputs=None
        )

        save_gigachat_lite_btn.click(
            fn=lambda personalized_message, comment, corrected_message, description, advantages, non_personalized_prompt, non_personalized_message, gender, generation, psychotype, business_stage, industry, legal_form: 
                save_to_github(personalized_message, "GigaChat-Lite", comment, corrected_message, description, advantages, non_personalized_prompt, non_personalized_message, gender, generation, psychotype, business_stage, industry, legal_form),
            inputs=[
                personalized_output_text_gigachat_lite,
                comment_gigachat_lite,
                corrected_gigachat_lite,
                description_input,
                advantages_input,
                prompt_display,
                output_text_gigachat_lite,
                selections[0],  # Пол
                selections[1],  # Поколение
                selections[2],  # Психотип
                selections[3],  # Стадия бизнеса
                selections[4],  # Отрасль
                selections[5]   # ОПФ
            ],
            outputs=None
        )

        save_gigachat_plus_btn.click(
            fn=lambda personalized_message, comment, corrected_message, description, advantages, non_personalized_prompt, non_personalized_message, gender, generation, psychotype, business_stage, industry, legal_form: 
                save_to_github(personalized_message, "GigaChat-Plus", comment, corrected_message, description, advantages, non_personalized_prompt, non_personalized_message, gender, generation, psychotype, business_stage, industry, legal_form),
            inputs=[
                personalized_output_text_gigachat_plus,
                comment_gigachat_plus,
                corrected_gigachat_plus,
                description_input,
                advantages_input,
                prompt_display,
                output_text_gigachat_plus,
                selections[0],  # Пол
                selections[1],  # Поколение
                selections[2],  # Психотип
                selections[3],  # Стадия бизнеса
                selections[4],  # Отрасль
                selections[5]   # ОПФ
            ],
            outputs=None
        )

    demo.launch()