import gradio as gr
import requests
import os
import base64
from PIL import Image
import numpy as np
import io

# Функция для обработки изображения и текста и отправки запроса к OpenAI
def generate_text(image, prompt):
    # Конвертируем NumPy массив в объект Image
    image_pil = Image.fromarray(image.astype('uint8'), 'RGB')
    
    # Конвертируем изображение для отправки через API
    image_bytes = io.BytesIO()
    image_pil.save(image_bytes, format='PNG')
    image_base64 = base64.b64encode(image_bytes.getvalue()).decode('utf-8')

    # API ключ для OpenAI
    api_key = os.getenv("API_KEY")

    # Заголовки для запроса
    headers = {
        'Authorization': f'Bearer {api_key}',
        'Content-Type': 'application/json',
    }

    # Данные для запроса
    data = {
        "model": "gpt-4-vision-preview",
        "prompt": prompt,
        "n": 1,
        "temperature": 0.5,
        "top_p": 1,
        "frequency_penalty": 0,
        "presence_penalty": 0,
        "stop": ["\n"],
        "image": {
            "data": image_base64,
            "mime_type": "image/png"
        }
    }

    # URL для запроса к модели gpt-4-vision-preview
    url = 'https://api.openai.com/v1/completions'

    # Отправляем запрос к OpenAI
    response = requests.post(url, headers=headers, json=data)
    
    # Проверяем ответ и возвращаем результат
    if response.status_code == 200:
        response_data = response.json()
        return response_data['choices'][0]['text'].strip()
    else:
        return f"Error: {response.status_code} - {response.text}"

# Создаем интерфейс с помощью Gradio
with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(label="Загрузите изображение", type="numpy")
            text_input = gr.Textbox(label="Введите текст")
            submit_button = gr.Button("Решить")
        with gr.Column():
            output_text = gr.Textbox(label="Ответ", interactive=True, lines=10)
            output_markdown = gr.Textbox(label="Ответ в Markdown", interactive=True, lines=10, visible=False)

    # Функция для обновления Markdown поля при получении ответа
    def update_markdown(answer):
        output_markdown.update(f"```\n{answer}\n```")
        output_markdown.change_visibility(True)

    # Привязываем функции к кнопке
    submit_button.click(fn=generate_text, inputs=[image_input, text_input], outputs=[output_text])
    output_text.change(fn=update_markdown, inputs=[output_text], outputs=[output_markdown])

demo.launch()