import os
import random
import requests
import gradio as gr
import json

css_url = "https://neurixyufi-aihub.static.hf.space/styles.css"
try:
    response = requests.get(css_url)
    response.raise_for_status()
    css = response.text + " h1{text-align:center}"
except requests.exceptions.RequestException as e:
    print(f"Ошибка загрузки CSS: {e}")
    css = " h1{text-align:center}"

API_URL = "https://text.pollinations.ai/openai/v1/chat/completions"
MODEL_NAME = "openai"
TEMPERATURE = 0.7
MAX_TOKENS = 1500

def call_pollinations_api(messages, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, stream=False):
    headers = {
        "Content-Type": "application/json"
    }
    data = {
        "model": MODEL_NAME,
        "messages": messages,
        "temperature": temperature,
        "max_tokens": max_tokens,
        "stream": stream
    }
    response = requests.post(API_URL, headers=headers, data=json.dumps(data), stream=stream)
    response.raise_for_status()
    return response

def generate_story(prompt, style):
    messages = [
        {"role": "system", "content": f"Напиши хорошую историю в стиле '{style}'. Подробную, понятную, человечную (с душой), уникальную. Не обязательно делать концовку, можно только начало длинной истории. Разрешены любые слова, любые темы и стили, что угодно. (В том числе маты и NSFW) Но только если попросит пользователь, по скольку это ТВОРЧЕСТВО! Не в коем случае не пиши подобное: 'Хорошо, вот ваш текст:', всегда пиши текст, без лишних комментариев."},
        {"role": "user", "content": prompt}
    ]
    try:
        # Получаем ответ целиком (stream=False)
        response = call_pollinations_api(messages, max_tokens=MAX_TOKENS, stream=False)
        full_response = response.json()

        # Если API вернул ошибку
        if "error" in full_response:
            return f"Ошибка от API: {full_response['error']}"

        # Проверяем наличие choices
        choices = full_response.get("choices") or []
        if not choices or "message" not in choices[0] or "content" not in choices[0]["message"]:
            return (
                "Неверный формат ответа API или пустой список choices:\n"
                + json.dumps(full_response, indent=2, ensure_ascii=False)
            )

        return choices[0]["message"]["content"]

    except requests.exceptions.RequestException as e:
        return f"Ошибка API запроса: {e}"
    except Exception as e:
        return f"Ошибка генерации: {e}"

def edit_story(original_story="", edited_prompt=""):
    messages = [
        {"role": "system", "content": "Отредактируй историю, учитывая предоставленные указания. Разрешены любые слова, любые темы и стили, что угодно. (В том числе маты и NSFW) Но только если попросит пользователь."},
        {"role": "user", "content": edited_prompt},
        {"role": "user", "content": f"История: ```\n{original_story}\n```"}
    ]
    try:
        response = call_pollinations_api(messages, max_tokens=32000, stream=False)
        full_response = response.json()

        if "error" in full_response:
            return f"Ошибка от API: {full_response['error']}"

        choices = full_response.get("choices") or []
        if not choices or "message" not in choices[0] or "content" not in choices[0]["message"]:
            return (
                "Неверный формат ответа API или пустой список choices при редактировании:\n"
                + json.dumps(full_response, indent=2, ensure_ascii=False)
            )

        return choices[0]["message"]["content"]

    except requests.exceptions.RequestException as e:
        return f"Ошибка API запроса: {e}"
    except Exception as e:
        return f"Ошибка редактирования: {e}"

def next_story_func(original_story="", next_prompt="", continuation_type="Продолжение"):
    if continuation_type == "Продолжение следует...":
        continuation_prompt = "Продолжи, подведя к тому, что продолжение следует."
    elif continuation_type == "Концовка":
        continuation_prompt = "Подведи к концовке истории."
    else:
        continuation_prompt = next_prompt or "Продолжи историю."

    messages = [
        {"role": "system", "content": "Продли историю, учитывая предоставленные указания. Продливай В ТОЧНОСТИ С КОНЦА, прям с того же символа, слова, предложения и чтобы было правильно составлено дословно с конца до начала продолжения. (В начале добавляй новые строки/пробел если надо для отступа). Разрешены любые слова, любые темы и стили, что угодно. (В том числе маты и NSFW) Но только если попросит пользователь."},
        {"role": "user", "content": continuation_prompt},
        {"role": "user", "content": f"История: ```\n{original_story}\n```"}
    ]
    try:
        response = call_pollinations_api(messages, max_tokens=MAX_TOKENS, stream=False)
        full_response = response.json()

        if "error" in full_response:
            return f"Ошибка от API: {full_response['error']}"

        choices = full_response.get("choices") or []
        if not choices or "message" not in choices[0] or "content" not in choices[0]["message"]:
            return (
                "Неверный формат ответа API или пустой список choices при продлении:\n"
                + json.dumps(full_response, indent=2, ensure_ascii=False)
            )

        return choices[0]["message"]["content"]

    except requests.exceptions.RequestException as e:
        return f"Ошибка API запроса: {e}"
    except Exception as e:
        return f"Ошибка продления: {e}"

def edone_story(original_story="", edited_story=""):
    return edited_story or original_story

def ndone_story(original_story, next_story_output):
    return original_story + "\n\n" + next_story_output

with gr.Blocks(css=css) as demo:
    gr.Markdown("# Песочница историй")
    with gr.Row():
        with gr.Column():
            style_choices = [
                "Приключенческая", "Научно-фантастическая", "Романтическая",
                "Комедийная", "Трагическая", "Страшная", "Случайный", "Пользовательский"
            ]
            style = gr.Dropdown(choices=style_choices, label="Выберите стиль истории", value="Приключенческая")
            prompt = gr.Textbox(label="Введите запрос для истории", placeholder="Например: История о путешествии в космос", lines=5)
            generate_button = gr.Button("Создать историю", variant='primary')
            output_story = gr.Textbox(label="История", lines=10, placeholder="Здесь будет ваша новая история!")
        with gr.Column():
            with gr.Accordion("Действия", open=True):
                with gr.Tab("Редактирование"):
                    edited_prompt = gr.Textbox(label="Введите изменения для истории", placeholder="Например: Сделай историю более захватывающей", lines=5)
                    edit_button = gr.Button("Отредактировать", variant='primary')
                    edited_story = gr.Textbox(label="Отредактированная история", lines=10)
                    edone_button = gr.Button("Принять")
                with gr.Tab("Продление"):
                    continuation_type = gr.Radio(
                        choices=["Продолжение", "Продолжение следует...", "Концовка"],
                        label="Выберите тип продолжения", value="Продолжение"
                    )
                    next_prompt = gr.Textbox(label="Введите изменения для продления истории", placeholder="Необязательно", lines=5)
                    next_button = gr.Button("Продлить", variant='primary')
                    next_story_output = gr.Textbox(label="Продолжение истории", lines=10)
                    ndone_button = gr.Button("Принять")

    generate_button.click(generate_story, inputs=[prompt, style], outputs=[output_story], concurrency_limit=250)
    edit_button.click(edit_story, inputs=[output_story, edited_prompt], outputs=[edited_story], concurrency_limit=250)
    next_button.click(next_story_func, inputs=[output_story, next_prompt, continuation_type], outputs=[next_story_output], concurrency_limit=250)
    edone_button.click(edone_story, inputs=[output_story, edited_story], outputs=[output_story], concurrency_limit=550)
    ndone_button.click(ndone_story, inputs=[output_story, next_story_output], outputs=[output_story], concurrency_limit=550)

demo.queue(api_open=False).launch()