import gradio as gr
import requests
import os
import pandas as pd
import time
from langchain.schema import SystemMessage
from langchain.chat_models.gigachat import GigaChat
from PIL import Image, ImageDraw, ImageFont

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

# Авторизация в сервисе GigaChat
chat_pro = GigaChat(credentials=gc_key, model='GigaChat-Pro', verify_ssl_certs=False)
chat_lite = GigaChat(credentials=gc_key, model='GigaChat', 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:
        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 generate_prompt(*selected_values):
    prompt = "Напиши типичное рекламное SMS-сообщение от ПАО Сбербанк, но с учетом следующих правил (длина сообщения не должна превышать 250 знаков с учетом пробелов):\n"
    for i, feature in enumerate(features.keys()):
        if selected_values[i]:
            try:
                prompt += f"{features[feature][selected_values[i]]}\n"
            except KeyError:
                return f"Ошибка: выбранное значение {selected_values[i]} не найдено в данных."
    return prompt.strip()

# Функция для генерации сообщения с помощью GPT-4o
def generate_message_gpt4o(prompt):
    try:
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {openai_api_key}"
        }
        data = {
            "model": "gpt-4o",
            "messages": [{"role": "system", "content": prompt}],
            "max_tokens": 700
        }
        response = requests.post("https://api.openai.com/v1/chat/completions", json=data, headers=headers)
        response_data = response.json()
        return response_data["choices"][0]["message"]["content"].strip()
    except Exception as e:
        return f"Ошибка при обращении к GPT-4o: {e}"

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

# Функция для генерации сообщения с помощью GigaChat
def generate_message_gigachat_lite(prompt):
    try:
        time.sleep(2)  # Задержка 2 секунды перед запросом к GigaChat
        messages = [SystemMessage(content=prompt)]
        res = chat_lite(messages)
        return res.content.strip()
    except Exception as e:
        return f"Ошибка при обращении к GigaChat-Lite: {e}"

# Функция для добавления текста на изображение
def add_text_to_image(image_path, prompt, gpt4o_message, gigachat_pro_message, gigachat_lite_message):
    img = Image.open(image_path)
    draw = ImageDraw.Draw(img)
    font = ImageFont.truetype("arial.ttf", 20)
    
    # Определите позиции, где будут отображаться ваши сообщения
    prompt_position = (50, 50)
    gpt4o_position = (50, 150)
    gigachat_pro_position = (50, 250)
    gigachat_lite_position = (50, 350)
    
    draw.text(prompt_position, prompt, font=font, fill="black")
    draw.text(gpt4o_position, gpt4o_message, font=font, fill="black")
    draw.text(gigachat_pro_position, gigachat_pro_message, font=font, fill="black")
    draw.text(gigachat_lite_position, gigachat_lite_message, font=font, fill="black")
    
    img.save("output_pipeline.png")
    return "output_pipeline.png"

# Функция для обработки нажатия кнопки Submit
def generate_messages(*selected_values):
    prompt = generate_prompt(*selected_values)
    if "Ошибка" in prompt:
        return prompt, "", "", "", None
    message_gpt4o = generate_message_gpt4o(prompt)
    message_gigachat_pro = generate_message_gigachat_pro(prompt)
    message_gigachat_lite = generate_message_gigachat_lite(prompt)
    
    # Добавление текста на изображение
    image_path = add_text_to_image('pipeline.png', prompt, message_gpt4o, message_gigachat_pro, message_gigachat_lite)
    
    return prompt, message_gpt4o, message_gigachat_pro, message_gigachat_lite, image_path

# Создание интерфейса Gradio
with gr.Blocks() as demo:
    gr.Markdown("# Генерация SMS-сообщений по заданным признакам")
    
    with gr.Row():
        with gr.Column(scale=1):
            selections = []
            for feature in features.keys():
                selections.append(gr.Dropdown(choices=[None] + list(features[feature].keys()), label=f"Выберите {feature}"))

            submit_btn = gr.Button("Submit")  # Оранжевая кнопка