import gradio as gr
import os
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 plotly.graph_objects as go
import random

# Авторизация в GigaChat Pro
gc_key = os.getenv('GC_KEY')
chat_pro = GigaChat(credentials=gc_key, model='GigaChat-Pro', max_tokens=68, temperature=1, verify_ssl_certs=False)

# Загрузка данных из Excel-файла
try:
    data = pd.read_excel('Признаки.xlsx', sheet_name=None)
except Exception as e:
    print(f"Ошибка при загрузке Excel-файла: {e}")
    data = {}

# Создание списка признаков и их значений
features = {}
for sheet_name, df in data.items():
    try:
        if sheet_name == "Пол Поколение Психотип":
            features[sheet_name] = df.set_index(['Пол', 'Поколение', 'Психотип'])['Инструкция'].to_dict()
        else:
            features[sheet_name] = df.set_index(df.columns[0]).to_dict()[df.columns[1]]
    except Exception as e:
        print(f"Ошибка при обработке данных листа {sheet_name}: {e}")
        features[sheet_name] = {}

# Функция для создания спидометра
def create_gauge(value):
    fig = go.Figure(go.Indicator(
        mode="gauge+number",
        value=value,
        gauge={
            'axis': {'range': [0, 100]},
            'bar': {'color': "black"},  # Цвет стрелки
            'steps': [
                {'range': [0, 40], 'color': "#55efc4"},  # Мягкий зеленый
                {'range': [40, 70], 'color': "#ffeaa7"},  # Желтый
                {'range': [70, 100], 'color': "#ff7675"}  # Мягкий красный
            ],
            'threshold': {
                'line': {'color': "black", 'width': 4},
                'thickness': 0.75,
                'value': value
            }
        },
        number={'font': {'size': 48}}  # Размер шрифта числа
    ))
    fig.update_layout(paper_bgcolor="#f8f9fa", font={'color': "#2d3436", 'family': "Arial"}, width=250, height=150)
    return fig

# Функция для генерации случайных значений спидометров
def generate_random_gauges():
    return create_gauge(random.randint(60, 90)), create_gauge(random.randint(60, 90)), create_gauge(random.randint(60, 90))

# Функция для смены вкладки
def change_tab(id):
    return gr.Tabs(selected=id)

# Вспомогательная функция для добавления префиксов и суффиксов
def add_prefix_suffix(prompt, prefix, suffix):
    return f"{prefix}\n{prompt}\n{suffix}"

# Функция для обрезки сообщения до последнего знака препинания
def clean_message(message):
    if not message.endswith(('.', '!', '?')):
        last_period = max(message.rfind('.'), message.rfind('!'), message.rfind('?'))
        if last_period != -1:
            message = message[:last_period + 1]
    return message

# Функция для генерации сообщения с GigaChat Pro
def generate_message_gigachat_pro(prompt):
    try:
        messages = [SystemMessage(content=prompt)]
        res = chat_pro(messages)
        cleaned_message = clean_message(res.content.strip())
        return cleaned_message
    except Exception as e:
        return f"Ошибка при обращении к GigaChat-Pro: {e}"

# Функция для повторной генерации сообщения, пока оно не станет короче 250 знаков
def generate_message_gigachat_pro_with_retry(prompt):
    for _ in range(10):
        message = generate_message_gigachat_pro(prompt)
        if len(message) <= 250:
            return message
    return message

# Измененная функция для генерации всех сообщений
def generate_all_messages(desc, benefits, key_message):
    # Варианты предложений для начала и конца
    prefixes = [
        "Начни сообщение с призыва к действию с продуктом.",
        "Начни сообщение с указания на пользу продукта. Используй глагол в побудительном наклонении.",
        "Начни сообщение с вопроса, который указывает на пользу продукта для клиента."
    ]
    suffixes = [
        "Убедись, что готовый текст начинается с призыва к действию с продуктом.",
        "Убедись, что готовый текст начинается с указания на пользу продукта и использования глагола в побудительном наклонении.",
        "Убедись, что готовый текст начинается с вопроса, который указывает на пользу продукта для клиента."
    ]
    
    non_personalized_messages = []
    personalized_messages = []
    
    # Формируем общий неперсонализированный текст для использования в сообщениях
    non_personalized_prompt = f"{desc}\n\nПреимущества:\n{benefits}\n\nКлючевое сообщение: {key_message}"
    personalized_prompt = f"Персонализированное сообщение для клиента: {desc}\n\n{benefits}\n\nКлючевое сообщение: {key_message}"

    # Выводим задание для копирайтера
    yield non_personalized_prompt, None, None, None, None, None, None, None
    
    # Небольшая пауза для демонстрации постепенной генерации
    time.sleep(1)
    
    # Выводим задание для редактора
    yield non_personalized_prompt, personalized_prompt, None, None, None, None, None, None
    
    # Небольшая пауза для демонстрации постепенной генерации
    time.sleep(1)
    
    # Генерация трех неперсонализированных сообщений
    for i in range(3):
        prompt = add_prefix_suffix(non_personalized_prompt, prefixes[i], suffixes[i])
        message = generate_message_gigachat_pro_with_retry(prompt)
        non_personalized_messages.append(message)
    
    # Генерация трех персонализированных сообщений
    for i in range(3):
        full_personalized_prompt = f"{personalized_prompt}\n\nТекст для адаптации: {non_personalized_messages[i]}"
        prompt = add_prefix_suffix(full_personalized_prompt, prefixes[i], suffixes[i])
        message = generate_message_gigachat_pro_with_retry(prompt)
        personalized_messages.append(message)
    
    # Выводим все сгенерированные сообщения
    yield (
        non_personalized_prompt, personalized_prompt,  # Поля для задания копирайтера и редактора
        non_personalized_messages[0], personalized_messages[0],  # Первое неперсонализированное и персонализированное сообщение
        non_personalized_messages[1], personalized_messages[1],  # Второе неперсонализированное и персонализированное сообщение
        non_personalized_messages[2], personalized_messages[2]   # Третье неперсонализированное и персонализированное сообщение
    )


# Интерфейс Gradio
with gr.Blocks() as demo:
    with gr.Tabs() as tabs:
        
        # Вкладка 1: Исходные данные
        with gr.TabItem("Исходные данные", id=0):
            with gr.Row():
                with gr.Column():
                    desc = gr.Textbox(
                        label="Описание предложения (предзаполненный пример можно поменять на свой)", 
                        lines=7,
                        value=(
                            "Необходимо предложить клиенту оформить дебетовую премиальную бизнес-карту Mastercard Preffered. "
                            "Обслуживание карты стоит 700 рублей в месяц, но клиент может пользоваться ей бесплатно. "
                            "Что необходимо сделать, чтобы воспользоваться предложением:\n"
                            "1. Оформить премиальную бизнес-карту в офисе банка или онлайн в интернет-банке СберБизнес.\n"
                            "2. Забрать карту.\n"
                            "3. В течение календарного месяца совершить по ней покупки на сумму от 100 000 рублей.\n"
                            "4. В течение следующего месяца пользоваться ей бесплатно."
                        )
                    )
                    benefits = gr.Textbox(
                        label="Преимущества (предзаполненный пример можно поменять на свой)", 
                        lines=5,
                        value=(
                            "Предложение по бесплатному обслуживанию — бессрочное.\n"
                            "Оплата покупок без отчётов и платёжных поручений.\n"
                            "Платёжные документы без комиссии.\n"
                            "Лимиты на расходы сотрудников.\n"
                            "Мгновенные переводы на карты любых банков."
                        )
                    )
                    
                    key_message = gr.Textbox(
                        label="Ключевое сообщение (предзаполненный пример можно поменять на свой)",
                        lines=3,
                        value="Бесплатное обслуживание при покупках от 100 000 рублей в месяц."
                    )

                with gr.Column():
                    gender = gr.Dropdown(label="Пол", choices=[None] + list(features.get('Пол', {}).keys()))
                    generation = gr.Dropdown(label="Поколение", choices=[None] + list(features.get('Поколение', {}).keys()))
                    psychotype = gr.Dropdown(label="Психотип", choices=[None] + list(features.get('Психотип', {}).keys()))
                    business_stage = gr.Dropdown(label="Стадия бизнеса", choices=[None] + list(features.get('Стадия бизнеса', {}).keys()))
                    industry = gr.Dropdown(label="Отрасль", choices=[None] + list(features.get('Отрасль', {}).keys()))
                    opf = gr.Dropdown(label="ОПФ", choices=[None] + list(features.get('ОПФ', {}).keys()))
            btn_to_prompts = gr.Button("Создать")
            
        # Вкладка 2: Промпты
        with gr.TabItem("Ассистент", id=1):
            with gr.Row():
                with gr.Column():
                    non_personalized_prompt = gr.Textbox(
                        label="Задание для копирайтера", 
                        lines=25,
                        interactive=False)
                with gr.Column():
                    personalized_prompt = gr.Textbox(label="Задание для редактора", lines=25)  # Увеличенная высота
                    
        # Вкладка 3: Сообщения
        with gr.TabItem("Сообщения", id=2):
            with gr.Row():
                gr.Markdown("### Копирайтер")
                gr.Markdown("### Редактор")
                
            with gr.Row():
                non_personalized_1 = gr.Textbox(label="Стандартное сообщение 1", lines=4, interactive=False)
                personalized_1 = gr.Textbox(label="Персонализированное сообщение 1", lines=4, interactive=False)

            with gr.Row():
                non_personalized_2 = gr.Textbox(label="Стандартное сообщение 2", lines=4, interactive=False)
                personalized_2 = gr.Textbox(label="Персонализированное сообщение 2", lines=4, interactive=False)

            with gr.Row():
                non_personalized_3 = gr.Textbox(label="Стандартное сообщение 3", lines=4, interactive=False)
                personalized_3 = gr.Textbox(label="Персонализированное сообщение 3", lines=4, interactive=False)

            # Четвертый ряд
            with gr.Row():
                btn_check = gr.Button("Проверить", elem_id="check3")
                btn_check.click(fn=change_tab, inputs=[gr.Number(value=3, visible=False)], outputs=tabs)

            # Сначала переключаем вкладку, потом запускаем генерацию сообщений
            btn_to_prompts.click(
                fn=change_tab, 
                inputs=[gr.Number(value=1, visible=False)],  # Переключение на вкладку "Сообщения" (id=1)
                outputs=tabs  # Обновляем вкладку
            ).then(
                fn=generate_all_messages, 
                inputs=[desc, benefits, key_message],  # Входные текстовые поля
                outputs=[
                    non_personalized_prompt, personalized_prompt,  # Поля для задания копирайтера и редактора (на вкладке "Ассистент")
                    non_personalized_1, personalized_1,  # Сообщения на вкладке "Сообщения"
                    non_personalized_2, personalized_2,
                    non_personalized_3, personalized_3
                ]
            )

        # Вкладка 4: Проверка
        with gr.TabItem("Проверка", id=3):
            with gr.Row():
                gr.Markdown("### Редактор")
                gr.Markdown("### Корректор")
                gr.Markdown("### Аналитик")
                
            with gr.Row():
                personalized_message_1 = gr.Textbox(label="Персонализированное сообщение 1", lines=5)
                check_message_1 = gr.Textbox(label="Проверка сообщения 1", lines=5)
                with gr.Column():
                    gr.HTML("<div style='display:flex; justify-content:center; width:100%;'>")
                    success_forecast_1 = gr.Plot(label="Прогноз успешности сообщения 1")
                    gr.HTML("</div>")
            
            with gr.Row():
                personalized_message_2 = gr.Textbox(label="Персонализированное сообщение 2", lines=5)
                check_message_2 = gr.Textbox(label="Проверка сообщения 2", lines=5)
                with gr.Column():
                    gr.HTML("<div style='display:flex; justify-content:center; width:100%;'>")
                    success_forecast_2 = gr.Plot(label="Прогноз успешности сообщения 2")
                    gr.HTML("</div>")
            
            with gr.Row():
                personalized_message_3 = gr.Textbox(label="Персонализированное сообщение 3", lines=5)
                check_message_3 = gr.Textbox(label="Проверка сообщения 3", lines=5)              
                with gr.Column():
                    gr.HTML("<div style='display:flex; justify-content:center; width:100%;'>")
                    success_forecast_3 = gr.Plot(label="Прогноз успешности сообщения 3")
                    gr.HTML("</div>")
                
            btn_check.click(fn=generate_random_gauges, inputs=[], outputs=[success_forecast_1, success_forecast_2, success_forecast_3])

demo.launch()