import gradio as gr
from g4f.client import Client
import re
import json

client = Client()

def get_task(prompt):
    return f'''Ты генерируешь скрипт продолжения прямой линии с Владимиром Владимировичем Путиным. Вопросы телезрителей передаёт Киселёв - он максимально подобострастен, поддерживает Путина и Россию. Пиши в духе пропаганды пророссийской. Каждая реплика - не менее 150-200 символов. Английские названия при необходимости всегда транскрибируй кириллицей. Числительные также записывай кириллицей. Пиши в стиле кремлёвской пропаганды. Пиши в стереотипном стиле. Киселёв задаёт вопросы, обрашается к Путину как "Владимир Владимирович" или "Господин Президент". Путин отвечает. Всё максимально аутентично. Фразы Киселёва: {{"Киселёв":"фраза"}} Фразы Путина: {{"Путин":"фраза"}} Ответ дай в формате JSON без дополнительных символов: [{{"Киселёв":"фраза"}}, {{"Путин":"фраза"}} . . . ].
Вопрос от пользователя поступил: "{prompt}"'''

def validate_response(response):
    try:
        # Attempt to parse the response as JSON
        data = json.loads(response)
        # Check if the response is a list and follows the required pattern
        if isinstance(data, list) and all(isinstance(item, dict) and len(item) == 1 for item in data):
            return True
    except json.JSONDecodeError:
        pass
    return False

def generate_text(prompt):
    max_retries = 4
    for attempt in range(max_retries):
        response = client.chat.completions.create(
            model="llama-3.3-70b",
            messages=[{"role": "user", "content": get_task(prompt)}],
            web_search=False
        )
        response_text = response.choices[0].message.content
        if validate_response(response_text):
            return response_text
    # If all retries fail, return a placeholder
    return '[{"Киселёв":"К сожалению, не удалось расслышать вопрос. Пожалуйста, попробуйте еще раз."}, {"Путин":"Мы работаем над улучшением системы. Спасибо за понимание."}]'

examples = [
    "Почему такие высокие налоги?",
    "Какие цели СВО?",
    "Когда развалится Америка?"
]

with gr.Blocks() as demo:
    gr.Markdown("# Kisel TV")
    
    with gr.Row():
        prompt_input = gr.Textbox(
            label="Input Prompt",
            placeholder="Enter your text here...",
            lines=3
        )
    
    generate_btn = gr.Button("Generate", variant="primary")
    output = gr.Textbox(label="Generated Text", interactive=False)
    
    gr.Examples(
        examples=examples,
        inputs=prompt_input,
        outputs=output,
        fn=generate_text,
        cache_examples=False
    )

    generate_btn.click(
        fn=generate_text,
        inputs=prompt_input,
        outputs=output
    )

if __name__ == "__main__":
    demo.launch()