import os
import gradio as gr
import asyncio
import json
from datetime import datetime
from dialogue.manager import DialogueManager
from dotenv import load_dotenv

# Загружаем переменные окружения
load_dotenv()

def create_interface():
    with gr.Blocks(theme=gr.themes.Soft()) as interface:
        gr.Markdown("# AI Dialogue Lab")
        
        state = gr.State({"is_processing": False})
        
        with gr.Row():
            with gr.Column(scale=3):
                topic = gr.TextArea(
                    label="Тема для обсуждения",
                    placeholder="Опишите тему максимально подробно...",
                    lines=4
                )
            with gr.Column(scale=1):
                iterations = gr.Slider(
                    minimum=1,
                    maximum=5,
                    step=1,
                    value=3,
                    label="Количество итераций"
                )
        
        with gr.Row():
            start_btn = gr.Button("Начать", variant="primary")
            save_btn = gr.Button("Сохранить")
            clear_btn = gr.Button("Очистить")
        
        status = gr.Markdown("")
        chatbot = gr.Chatbot(label="Диалог", height=600)
        dialogue_history = gr.State([])
        
        def save_dialogue(history):
            if not history:
                return "❌ Нет данных для сохранения"
            
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"dialogue_{timestamp}.json"
            
            data = {
                "timestamp": datetime.now().isoformat(),
                "messages": history
            }
            
            with open(filename, "w", encoding="utf-8") as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            
            return f"✅ Диалог сохранен в файл: {filename}"

        def update_ui_state(is_processing: bool):
            return {
                start_btn: gr.Button(interactive=not is_processing),
                topic: gr.TextArea(interactive=not is_processing),
                iterations: gr.Slider(interactive=not is_processing),
                save_btn: gr.Button(interactive=not is_processing and bool(dialogue_history))
            }
        
        async def start_dialogue(topic_text, num_iterations):
            # Деактивируем элементы интерфейса
            start_btn.interactive = False
            topic.interactive = False
            iterations.interactive = False
            save_btn.interactive = False
            
            chat_history = []
            dialogue_history = []

            try:
                manager = DialogueManager(topic_text, num_iterations)
                
                # Оптимизация промпта
                yield "🔄 Оптимизация запроса...", chat_history
                
                optimized = await manager.optimize_prompt(topic_text)
                chat_history = [("🔧 System", f"Оптимизированная тема:\n{optimized}")]
                dialogue_history = chat_history.copy()
                yield "✨ Начинаем диалог...", chat_history
                
                current_prompt = optimized
                for i in range(num_iterations):
                    # Ход Claude
                    yield "🤔 Claude обдумывает ответ...", chat_history
                    
                    claude_response = await manager.claude.generate_response(current_prompt)
                    chat_history.append(("🔵 Claude", claude_response))
                    dialogue_history = chat_history.copy()
                    yield f"✨ Ход {i+1}/{num_iterations}: Claude ответил", chat_history
                    
                    # Ход Grok
                    yield "🤔 Grok обдумывает ответ...", chat_history
                    
                    grok_response = await manager.grok.generate_response(claude_response)
                    chat_history.append(("🟢 Grok", grok_response))
                    dialogue_history = chat_history.copy()
                    yield f"✨ Ход {i+1}/{num_iterations}: Grok ответил", chat_history
                    
                    current_prompt = grok_response
                
                # Активируем элементы интерфейса
                start_btn.interactive = True
                topic.interactive = True
                iterations.interactive = True
                save_btn.interactive = True
                
                yield "✅ Диалог завершен", chat_history
                
            except Exception as e:
                # В случае ошибки тоже активируем элементы
                start_btn.interactive = True
                topic.interactive = True
                iterations.interactive = True
                save_btn.interactive = False
                yield f"❌ Ошибка: {str(e)}", chat_history

        def clear_chat():
            return {
                status: "",
                chatbot: [],
                dialogue_history: [],
                start_btn: gr.Button(interactive=True),
                topic: gr.TextArea(interactive=True),
                iterations: gr.Slider(interactive=True),
                save_btn: gr.Button(interactive=False)
            }

        start_btn.click(
            fn=start_dialogue,
            inputs=[topic, iterations],
            outputs=[status, chatbot],
            api_name="start_dialogue",
            show_progress="full"
        ).then(
            lambda: gr.Button(interactive=True),
            None,
            [start_btn]
        )
        
        save_btn.click(
            fn=save_dialogue,
            inputs=[dialogue_history],
            outputs=[status],
            api_name="save_dialogue"
        )
        
        clear_btn.click(
            fn=clear_chat,
            outputs=[status, chatbot, dialogue_history, start_btn, topic, iterations, save_btn],
            api_name="clear_chat"
        )

    return interface

if __name__ == "__main__":
    interface = create_interface()
    interface.queue().launch()