#!/usr/bin/env python3
"""
Оновлений MAI-DX Gradio Interface з повним логуванням розмов агентів 
(ВЕРСІЯ З ВИПРАВЛЕНИМ СТАТУС-БАРОМ І СТРУКТУРОЮ)
"""
import os
import sys
import json
import time
import uuid
import pandas as pd
import gradio as gr
from datetime import datetime
from typing import Dict, List, Tuple, Optional, Any
from dataclasses import dataclass, asdict
import warnings

# Налаштування середовища
os.environ.update({
    "SWARMS_VERBOSITY": "ERROR", "RICH_TRACEBACK": "0", 
    "SWARMS_SHOW_PANEL": "true", "SWARMS_AUTO_PRINT": "true"
})
warnings.filterwarnings("ignore")

from dotenv import load_dotenv
load_dotenv()

try:
    from mai_dx_wrapper import SafeMaiDxOrchestrator as MaiDxOrchestrator, MAI_DX_AVAILABLE
    from enhanced_mai_dx_logger import MAIDxConversationLogger, DiagnosisSession, AgentConversation, EnhancedOutputCapture
except ImportError as e:
    MAI_DX_AVAILABLE = False
    IMPORT_ERROR = str(e)

class UpdatedMAIDXInterface:
    """
    Клас, що інкапсулює логіку та побудову Gradio інтерфейсу.
    Це запобігає дублюванню обробників подій.
    """
    def __init__(self):
        self.conversation_logger = MAIDxConversationLogger("mai_dx_logs")
        self.sessions_history = []
        self.sample_cases = {
            "🫀 Кардіологічний (Гострий MI)": {
                "info": "Пацієнт: 58-річний чоловік, менеджер, гіпертонія в анамнезі\nСкарги: Гострий роздираючий біль у грудях 3 години, іррадіація в ліву руку\nОгляд: Блідий, пітливий, АТ 160/90, ЧСС 95\nЕКГ: ST-підйоми у відведеннях II, III, aVF (нижня стінка)\nТропонін I: 8.5 нг/мл (норма <0.04)\nАнамнез: Куріння 30 років, дислипідемія, сімейний анамнез ІХС",
                "expected": "Acute inferior wall myocardial infarction (STEMI)"
            },
            "🧠 Неврологічний (Гострий інсульт)": {
                "info": "Пацієнтка: 67-річна жінка, раптові неврологічні симптоми\nПрезентація: Раптова слабкість правої сторони 2 години тому\nОгляд: Свідома, дезорієнтована у часі, правостороннє опущення обличчя\nНеврологія: Правостороння геміплегія, афазія, девіація очей вліво\nКТ голови: Гострого крововиливу немає, рання ішемія у лівій МСА\nNIHSS: 15 балів",
                "expected": "Acute ischemic stroke in the left middle cerebral artery territory"
            },
            "🦠 Інфекційний (Сепсис)": {
                "info": "Пацієнт: 45-річний чоловік з прогресуючою лихоманкою\nСкарги: Висока температура 39.5°C, озноб, загальна слабкість 3 дні\nОгляд: Гарячий, тахікардія 120/хв, гіпотензія 85/50\nЛабораторно: Лейкоцити 18000, С-реактивний білок 180, прокальцитонін 5.2\nПосів крові: Pending, lactate 4.2 ммоль/л\nАнамнез: Нещодавня стоматологічна процедура",
                "expected": "Sepsis with a possible odontogenic source"
            }
        }
        # Створюємо інтерфейс один раз при ініціалізації
        self.demo = self.build_interface()

    def diagnose_with_full_logging(
        self, case_name: str, patient_info: str, mode: str, budget: int, max_iterations: int,
        model_name: str, expected_diagnosis: str = "",
        progress=gr.Progress(track_tqdm=True)
    ):
        if not MAI_DX_AVAILABLE: return self._format_error(f"❌ MAI-DX недоступний: {IMPORT_ERROR}")
        if not patient_info.strip(): return self._format_error("❌ Введіть інформацію про пацієнта")
        
        session = self.conversation_logger.create_session(case_name, patient_info, mode, budget)
        
        try:
            progress(0.2, desc="🤖 Створення AI-панелі...")
            orchestrator = MaiDxOrchestrator(model_name, max_iterations, budget, mode)
            
            progress(0.3, desc="🔍 Запуск діагностики...")
            start_time = time.time()
            
            with EnhancedOutputCapture() as capture:
                result = orchestrator.run(patient_info, patient_info, expected_diagnosis or "Unknown")
            raw_output = capture.get_value()

            duration = time.time() - start_time
            progress(0.8, desc="📊 Обробка результатів...")
            
            session = self.conversation_logger.finalize_and_save_session(session, result, raw_output, duration)
            self.sessions_history.append(session)
            
            progress(1.0, desc="✅ Готово!")
            
            main_result_str = self._format_main_result(session)
            detailed_analysis_str = self._format_detailed_analysis(session)
            recommendations_str = self._generate_enhanced_recommendations(result, expected_diagnosis)
            agent_report_str = self._create_agent_activity_report(session.conversations)
            structured_conv_str = self._format_structured_conversations(session.conversations)
            
            return main_result_str, detailed_analysis_str, recommendations_str, agent_report_str, structured_conv_str

        except Exception as e:
            import traceback
            error_msg = f"❌ Критична помилка: {traceback.format_exc()}"
            return self._format_error(error_msg)

    def _format_structured_conversations(self, conversations: List[AgentConversation]) -> str:
        if not conversations: return "### 🗣️ Розмови агентів\n\n*Розмови не знайдено в логах. Можливо, діагноз було поставлено миттєво.*"
        output = "### 🗣️ Розмови агентів\n\n"
        for conv in conversations:
            output += f"#### 🔄 Раунд {conv.round_number}\n"
            for msg in conv.messages:
                output += f"<details><summary><strong>{msg.agent_name}</strong> <code>[{msg.message_type}]</code></summary>\n\n"
                output += f"```\n{msg.content.strip()}\n```\n\n</details>\n"
        return output
    
    def _create_agent_activity_report(self, conversations: List[AgentConversation]) -> str:
        if not conversations:
            return "### 🤖 Активність агентів\n\n*Дані відсутні.*"
        
        agent_activity = {}
        for conv in conversations:
            for msg in conv.messages:
                agent_activity[msg.agent_name] = agent_activity.get(msg.agent_name, 0) + 1
        
        if not agent_activity:
            return "### 🤖 Активність агентів\n\n*Повідомлення від агентів не знайдено.*"

        report_md = "### 🤖 Активність агентів\n\n"
        sorted_activity = sorted(agent_activity.items(), key=lambda item: item[1], reverse=True)
        
        for agent, count in sorted_activity:
            report_md += f"- **{agent}**: {count} повідомлень\n"
            
        return report_md
    
    def _format_main_result(self, session: DiagnosisSession) -> str:
        efficiency = ((session.budget - session.cost) / session.budget * 100) if session.budget > 0 else 0
        confidence_emoji = "🎯" if session.confidence >= 4.0 else "👍" if session.confidence >= 3.0 else "⚠️"
        
        return f"""
## 🏥 Результати MAI-DX Діагностики
### 📋 Основна інформація
- **🗂️ Випадок**: {session.case_name}
- **🆔 ID сесії**: {session.case_id}
- **⏰ Час**: {session.timestamp}  
- **🔧 Режим**: {session.mode}
### {confidence_emoji} Діагностичний висновок
**{session.diagnosis}**
### 📊 Показники якості
- **Точність**: {session.confidence:.1f}/5.0 ⭐
- **Статус**: {session.status}
- **Ітерації**: {session.iterations} циклів
### 💰 Економічні показники  
- **Витрачено**: ${session.cost:,.2f}
- **Бюджет**: ${session.budget:,}
- **Ефективність**: {efficiency:.1f}% бюджету збережено
- **Швидкість**: {session.duration:.1f} секунд
"""
    
    def _format_detailed_analysis(self, session: DiagnosisSession) -> str:
        return f"""
## 🔬 Детальний клінічний аналіз
### 💭 Медичне обґрунтування
{session.reasoning}
### 📂 Файли логів
- JSON: `mai_dx_logs/{session.case_id}.json`
- Raw text: `mai_dx_logs/{session.case_id}_raw.txt`
- HTML звіт: `mai_dx_logs/{session.case_id}_conversation.html`
"""
    
    def _generate_enhanced_recommendations(self, result: Any, expected: str) -> str:
        comparison = ""
        if expected:
            is_match = expected.lower() in result.final_diagnosis.lower()
            comparison = f"""
### 🎯 Порівняння з очікуваним
**Очікувався**: {expected}
**Отримано**: {result.final_diagnosis}
**Збіг**: {'✅ Так' if is_match else '❌ Ні'}
"""
        return f"""
## 💡 Клінічні рекомендації
- 🔍 Верифікувати діагноз з лікарем-спеціалістом.
- 📊 Переглянути детальний лог розмов агентів.
{comparison}
### ⚠️ Важливе застереження
🔴 **Цей діагноз згенеровано ШІ і НЕ замінює професійну медичну консультацію.**
"""
    
    def _format_error(self, msg: str) -> Tuple[str, str, str, str, str]:
        return (msg, "", "", "", "")
    
    def load_sample_case(self, key: str) -> Tuple[str, str, str]:
        case = self.sample_cases.get(key, {})
        return case.get("info", ""), key, case.get("expected", "")

    def build_interface(self):
        """
        Метод, що будує інтерфейс Gradio. Викликається один раз.
        """
        with gr.Blocks(title="🏥 MAI-DX Enhanced Platform", theme=gr.themes.Soft(primary_hue="blue")) as demo:
            gr.HTML("""<div style='background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; padding: 2rem; border-radius: 15px; margin-bottom: 2rem; text-align: center;'><h1>🏥 MAI-DX Enhanced Research Platform</h1></div>""")
            with gr.Tabs():
                with gr.Tab("🩺 Діагностика"):
                    with gr.Row():
                        with gr.Column(scale=1):
                            gr.Markdown("### 📝 Клінічний випадок")
                            case_name = gr.Textbox(label="🏷️ Назва випадку", placeholder="Наприклад: Кардіологічний випадок №1")
                            sample_selector = gr.Dropdown(choices=list(self.sample_cases.keys()), label="🎯 Готові випадки", interactive=True)
                            patient_info = gr.Textbox(label="👤 Інформація про пацієнта", lines=12)
                            expected_diagnosis = gr.Textbox(label="🎯 Очікуваний діагноз (англ.)", lines=2)
                        with gr.Column(scale=1):
                            gr.Markdown("### ⚙️ Налаштування")
                            mode = gr.Radio(["instant", "question_only", "budgeted", "no_budget", "ensemble"], label="🔧 Режим", value="budgeted")
                            budget = gr.Slider(500, 10000, 3000, step=500, label="💵 Бюджет ($)")
                            max_iterations = gr.Slider(1, 15, 8, step=1, label="🔄 Макс. ітерацій")
                            model_name = gr.Dropdown(["gemini/gemini-2.5-flash", "gpt-4o", "claude-3-5-sonnet", "gpt-4", "gpt-4-turbo"], label="🤖 LLM Модель", value="gemini/gemini-2.5-flash")
                            diagnose_btn = gr.Button("🚀 Запустити діагностику", variant="primary", size="lg")
                    
                    gr.Markdown("---")
                    with gr.Row():
                        with gr.Column(scale=2):
                            main_result = gr.Markdown()
                            detailed_analysis = gr.Markdown()
                        with gr.Column(scale=1):
                            recommendations = gr.Markdown()
                    
                    gr.Markdown("---")
                    with gr.Row():
                        with gr.Column(scale=1):
                            agent_activity_report = gr.Markdown(label="🤖 Активність агентів")
                        with gr.Column(scale=2):
                            structured_conversations = gr.Markdown(label="🗣️ Структуровані розмови")
                    
                    # Прив'язуємо обробники подій до методів класу
                    sample_selector.change(self.load_sample_case, inputs=sample_selector, outputs=[patient_info, case_name, expected_diagnosis])
                    
                    diagnose_btn.click(
                        self.diagnose_with_full_logging,
                        inputs=[case_name, patient_info, mode, budget, max_iterations, model_name, expected_diagnosis],
                        outputs=[main_result, detailed_analysis, recommendations, agent_activity_report, structured_conversations],
                        show_progress="full"
                    )
                
                with gr.Tab("📊 Аналіз логів"):
                    gr.Markdown("### 🔍 Аналіз збережених діагностичних сесій")

        return demo

    def launch(self, *args, **kwargs):
        """Метод для запуску Gradio-додатка."""
        self.demo.launch(*args, **kwargs)

# Ця функція тепер є обгорткою для створення екземпляра класу
def create_updated_gradio_interface():
    app = UpdatedMAIDXInterface()
    return app.demo

if __name__ == "__main__":
    # Якщо файл запускається напряму, створюємо екземпляр і запускаємо
    app = UpdatedMAIDXInterface()
    app.launch(server_name="0.0.0.0", server_port=7860, share=False)