import gradio as gr
import requests
import time
import json
import os
from PIL import Image
import io
import base64
from datetime import datetime
import re
import time
import random
from bs4 import BeautifulSoup
from contextlib import closing
from websocket import create_connection


def flip_text(prompt, key):

    key_os = os.getenv("key_os")

    url_api1 = os.getenv("url_api1")
    url_api2 = os.getenv("url_api2")

    model_gr = os.getenv("model_gr")
    headers_a = os.getenv("headers_gr")
    a = eval(str(headers_a))
    api_gr = os.getenv("api_gr")
    
    if str(key).lower() == str(key_os).lower():
        try:
            with closing(create_connection(f"{api_gr}", timeout=60, header=a)) as conn:
                conn.send('{"fn_index":46,"session_hash":""}')
                conn.send(f'{{"data":["{model_gr}",null,4096,"gpt-3.5-turbo","{prompt}","",1,1,[],null,"","1024x1024"],"event_data":null,"fn_index":46,"session_hash":""}}')
                while True:
                    conn_s = conn.recv()
                    print(conn_s)
                    status = json.loads(conn_s)['msg']
                    print(status)
                    if status == 'estimation':
                        continue
                    if status == 'process_starts':
                        continue
                    if status == 'process_generating':
                        continue
                    if status == 'process_completed':
                        s = BeautifulSoup(str(conn_s), 'html.parser')
                        return str(s.find('img').get('src')).replace(r'\"', '')
                        break
        except:
            try:
                print(prompt)
                headers = os.getenv("headers")
                url_api = os.getenv("url_api")
                
                headers = eval(str(headers))
                json_data = {'prompt': prompt, 'n': 1,'size': '1024x1024','response_format': 'b64_json','model': 'dall-e-3','quality': 'hd',}
                response = requests.post(f"{url_api}", headers=headers, json=json_data)
                js = response.json()
                if "content_policy_violation" in str(js):
                    return "https://myneuralnetworks.ru/static/img/zp.jpg"
                else:
                    photo = response.json()['data'][0]['b64_json']
                    photo = Image.open(io.BytesIO(base64.decodebytes(bytes(photo, "utf-8"))))
                return photo
            except:
                send_time = str(datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
                url_id1 = os.getenv("url_id1")
                r_m = requests.post("https://myneuralnetworks.ru/get_key_dalle/", data={"key": url_id1})
                key_a = r_m.json()['result']
                
                headers = {'authorization': key_a}
                json_data = {'apiSource': 'gpt-all-tools'}
                response = requests.post(f'{url_api1}', headers=headers, json=json_data)
                chatUuid = response.json()['data']['chatUuid']
                
                json_data2 = {
                    'prompt': f'Нарисуй {prompt}',
                    'chatUuid': chatUuid,
                    'sendTime': send_time,
                    'attachments': [],
                }
                
                response2 = requests.post(f'{url_api2}', headers=headers, json=json_data2)
                url_pattern = r'http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+'
                urls = re.findall(url_pattern, str(response2.text))
                return urls[0]
            
    else:
        return "https://myneuralnetworks.ru/static/img/zp2.jpg"


css = """
#generate {
    width: 100%;
    background: #e253dd !important;
    border: none;
    border-radius: 50px;
    outline: none !important;
    color: white;
}
#generate:hover {
    background: #de6bda !important;
    outline: none !important;
    color: #fff;
    }
footer {visibility: hidden !important;}
#image_output {
height: 100% !important;
}
"""

with gr.Blocks(css=css) as demo:

    with gr.Row():
        prompt = gr.Textbox(placeholder="Введите описание изображения...", show_label=True, label='Описание изображения:', lines=3)
    with gr.Row():
        tg = gr.gradio.HTML("Ключ доступа можно найти в моём <a href='https://t.me/myneuralnetworks' target='_blank'>Telegram-канале</a>")
    with gr.Row():
        key = gr.Textbox(placeholder="Введите ключ доступа...", show_label=True, label='Ключ доступа:', lines=1)
      
    with gr.Row():
        text_button = gr.Button("Сгенерировать изображение", variant='primary', elem_id="generate")
    with gr.Row():
        image_output = gr.Image(show_download_button=True, interactive=False, type="numpy", label='Результат:', elem_id='image_output')
        text_button.click(flip_text, inputs=[prompt, key], outputs=image_output, concurrency_limit=12)
        
demo.launch()