import os
import random
import requests
from deep_translator import GoogleTranslator
from langdetect import detect
from huggingface_hub import InferenceClient
import gradio as gr

def get_random_api_key():
    keys = os.getenv("KEYS", "").split(",")
    if keys and keys[0]:  # Check if KEYS is set and not empty
        return random.choice(keys).strip()
    else:
        raise ValueError("API keys not found. Please set the KEYS environment variable.")

# Ссылка на файл CSS
css_url = "https://neurixyufi-aihub.static.hf.space/style.css"

# Получение CSS по ссылке
try:
    response = requests.get(css_url)
    response.raise_for_status()
    css = response.text + " h1{text-align:center} .container { max-width: 800px; margin: 0 auto; }"
except requests.exceptions.RequestException as e:
    print(f"Ошибка при загрузке CSS: {e}")
    css = " h1{text-align:center} .container { max-width: 800px; margin: 0 auto; }"


def generate_story(prompt, style, api_key):
    try:
        client = InferenceClient(api_key=api_key, model_id="Qwen/QwQ-32B-Preview")
        
        language = detect(prompt)
        if language != 'en':
            prompt = GoogleTranslator(source=language, target='en').translate(prompt)
        
        messages = [
            {"role": "system", "content": f"Напиши историю в стиле {style}."},
            {"role": "user", "content": prompt}
        ]

        completion = client.chat.completions.create(messages=messages, max_tokens=1000)
        story = completion.choices[0].message.content
        return story
    except Exception as e:
        return f"Ошибка генерации: {e}"


def edit_story(original_story, edited_prompt, api_key):
    try:
        client = InferenceClient(api_key=api_key, model_id="Qwen/QwQ-32B-Preview")
        
        messages = [
            {"role": "system", "content": "Отредактируй историю, учитывая предоставленные указания."},
            {"role": "user", "content": edited_prompt},
            {"role": "assistant", "content": original_story}
        ]
        
        completion = client.chat.completions.create(messages=messages, max_tokens=1000)
        edited_story = completion.choices[0].message.content
        return edited_story
    except Exception as e:
        return f"Ошибка редактирования: {e}"

with gr.Blocks(css=css) as demo:
    with gr.Row():
        style_choices = ["Приключенческая", "Научно-фантастическая", "Романтическая", "Комедийная", "Трагическая", "Введите свой стиль:"]
        style = gr.Dropdown(choices=style_choices, label="Выберите стиль истории", value="Приключенческая")
    with gr.Row():
        prompt = gr.Textbox(label="Введите запрос для истории", placeholder="Например: История о путешествии в космос", lines=5)
    with gr.Row():
        generate_button = gr.Button("Сгенерировать историю", variant='primary')
        
    with gr.Row():
        output_story = gr.Textbox(label="История", lines=10)
    
    with gr.Tab("Редактирование"):
        edited_prompt = gr.Textbox(label="Введите изменения для истории", placeholder="Например: Сделай историю более захватывающей", lines=5)
        edited_story = gr.Textbox(label="Отредактированная история", lines=10)
        edit_button = gr.Button("Отредактировать", variant='primary')
    
    api_key = get_random_api_key()
    generate_button.click(generate_story, inputs=[prompt, style], outputs=[output_story], api_key=api_key)
    edit_button.click(edit_story, inputs=[output_story, edited_prompt], outputs=[edited_story], api_key=api_key)

demo.launch(show_api=False, share=False,concurrency_limit=250)