app/gui_app.py

"""
Medical Transcription GUI Application
Полнофункциональное приложение для транскрибирования медицинских диктовок
с автоматической генерацией отчётов
"""

import sys
import logging
from pathlib import Path
from typing import Optional, Dict, Any
import threading
import traceback
from datetime import datetime
import os

from PyQt6.QtWidgets import (
    QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
    QLabel, QPushButton, QLineEdit, QTextEdit, QFileDialog,
    QComboBox, QSpinBox, QCheckBox, QProgressBar, QMessageBox,
    QTabWidget, QFormLayout, QGroupBox, QDialog, QScrollArea
)
from PyQt6.QtCore import Qt, pyqtSignal, QObject, QThread
from PyQt6.QtGui import QFont, QIcon, QColor
from PyQt6.QtCore import QTimer

# Import common utilities
from common import (
    get_logger,
    UIColors,
    UIDimensions,
    Messages,
    Placeholders,
    AudioFileException,
    TranscriptionException,
    ValidationException
)

# Import pipeline config for type hints
from pipeline import PipelineConfig

logger = get_logger(__name__)


class WorkerSignals(QObject):
    """Сигналы для воркера обработки"""
    progress = pyqtSignal(str)
    finished = pyqtSignal(dict)
    error = pyqtSignal(str)


class TranscriptionWorker(QThread):
    """Воркер для обработки аудио в отдельном потоке"""
    
    signals = WorkerSignals()
    
    def __init__(
        self,
        audio_path: str,
        config: PipelineConfig,
        patient_data: Dict[str, Any]
    ) -> None:
        super().__init__()
        self.audio_path: str = audio_path
        self.config: PipelineConfig = config
        self.patient_data: Dict[str, Any] = patient_data
    
    def run(self) -> None:
        try:
            # Импортируем здесь, чтобы избежать циклических зависимостей
            from pipeline.medical_pipeline import MedicalTranscriptionPipeline
            
            self.signals.progress.emit("Инициализация пайплайна...")
            pipeline = MedicalTranscriptionPipeline(self.config)
            
            self.signals.progress.emit("Запуск транскрибирования...")
            result = pipeline.process(
                audio_path=self.audio_path,
                patient_name=self.patient_data.get("patient_name"),
                patient_dob=self.patient_data.get("patient_dob"),
                study_area=self.patient_data.get("study_area"),
                study_number=self.patient_data.get("study_number"),
                study_date=self.patient_data.get("study_date"),
                doctor_name=self.patient_data.get("doctor_name"),
                generate_report=self.config.generate_report
            )
            
            self.signals.progress.emit("Обработка завершена!")
            self.signals.finished.emit(result)
            
        except Exception as e:
            logger.error(f"Error in transcription worker: {e}\n{traceback.format_exc()}")
            self.signals.error.emit(str(e))


class PatientDataDialog(QDialog):
    """Диалог для ввода данных пациента"""
    
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("Данные пациента")
        self.setGeometry(100, 100, 500, 400)
        self.init_ui()
        self.result = None
    
    def init_ui(self):
        layout = QFormLayout()
        
        self.patient_name = QLineEdit()
        self.patient_name.setPlaceholderText("Фамилия Имя Отчество")
        
        self.patient_dob = QLineEdit()
        self.patient_dob.setPlaceholderText("ДД.MM.YYYY")
        
        self.study_area = QLineEdit()
        self.study_area.setPlaceholderText("Область исследования (напр. МРТ головы)")
        
        self.study_number = QLineEdit()
        self.study_number.setPlaceholderText("Номер исследования")
        
        self.study_date = QLineEdit()
        self.study_date.setPlaceholderText("ДД.MM.YYYY")
        self.study_date.setText(datetime.now().strftime("%d.%m.%Y"))
        
        self.doctor_name = QLineEdit()
        self.doctor_name.setPlaceholderText("ФИО врача")
        
        layout.addRow("ФИО пациента:", self.patient_name)
        layout.addRow("Дата рождения:", self.patient_dob)
        layout.addRow("Область исследования:", self.study_area)
        layout.addRow("Номер исследования:", self.study_number)
        layout.addRow("Дата исследования:", self.study_date)
        layout.addRow("ФИО врача:", self.doctor_name)
        
        # Кнопки
        button_layout = QHBoxLayout()
        ok_btn = QPushButton("OK")
        cancel_btn = QPushButton("Отмена")
        
        ok_btn.clicked.connect(self.accept)
        cancel_btn.clicked.connect(self.reject)
        
        button_layout.addWidget(ok_btn)
        button_layout.addWidget(cancel_btn)
        
        layout.addRow(button_layout)
        self.setLayout(layout)
    
    def get_data(self):
        """Получить введённые данные"""
        return {
            "patient_name": self.patient_name.text(),
            "patient_dob": self.patient_dob.text(),
            "study_area": self.study_area.text(),
            "study_number": self.study_number.text(),
            "study_date": self.study_date.text(),
            "doctor_name": self.doctor_name.text()
        }


class MedicalTranscriptionApp(QMainWindow):
    """Главное окно приложения"""
    
    def __init__(self):
        super().__init__()
        self.setWindowTitle(Messages.APP_TITLE)
        self.setGeometry(100, 100, UIDimensions.WINDOW_WIDTH, UIDimensions.WINDOW_HEIGHT)
        
        # Переменные
        self.audio_path = None
        self.model_path = Path(__file__).parent.parent
        self.worker = None
        self.patient_data = {}
        
        self.init_ui()
        self.setup_logging()
        
        # Установка стилей
        self.apply_styles()
    
    def setup_logging(self):
        """Настройка логирования"""
        from common import configure_logging
        configure_logging()
    
    def init_ui(self):
        """Инициализация интерфейса"""
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        
        # Создание табов
        tabs = QTabWidget()
        
        # Таб 1: Транскрибирование
        transcription_tab = self.create_transcription_tab()
        tabs.addTab(transcription_tab, "Транскрибирование")
        
        # Таб 2: Настройки
        settings_tab = self.create_settings_tab()
        tabs.addTab(settings_tab, "Настройки")
        
        # Главный layout
        main_layout = QVBoxLayout()
        main_layout.addWidget(tabs)
        main_widget.setLayout(main_layout)
    
    def create_transcription_tab(self):
        """Создание вкладки транскрибирования"""
        widget = QWidget()
        layout = QVBoxLayout()
        
        # --- Выбор аудиофайла ---
        file_group = QGroupBox("1. Выбор аудиофайла")
        file_layout = QHBoxLayout()
        
        self.file_path_label = QLineEdit()
        self.file_path_label.setReadOnly(True)
        self.file_path_label.setPlaceholderText("Аудиофайл не выбран")
        
        browse_btn = QPushButton("Обзор...")
        browse_btn.clicked.connect(self.browse_audio_file)
        
        file_layout.addWidget(QLabel("Файл:"))
        file_layout.addWidget(self.file_path_label, 1)
        file_layout.addWidget(browse_btn)
        
        file_group.setLayout(file_layout)
        layout.addWidget(file_group)
        
        # --- Данные пациента ---
        patient_group = QGroupBox("2. Данные пациента")
        patient_layout = QVBoxLayout()
        
        self.patient_info_label = QLabel("Данные пациента не заполнены")
        patient_info_font = QFont()
        patient_info_font.setItalic(True)
        self.patient_info_label.setFont(patient_info_font)
        
        patient_btn = QPushButton("Заполнить данные пациента...")
        patient_btn.clicked.connect(self.open_patient_dialog)
        
        patient_layout.addWidget(self.patient_info_label)
        patient_layout.addWidget(patient_btn)
        patient_group.setLayout(patient_layout)
        layout.addWidget(patient_group)
        
        # --- Опции обработки ---
        options_group = QGroupBox("3. Опции обработки")
        options_layout = QFormLayout()
        
        self.llm_checkbox = QCheckBox("Использовать LLM-коррекцию")
        self.llm_checkbox.setChecked(True)
        
        self.report_checkbox = QCheckBox("Автоматически создать отчёт")
        self.report_checkbox.setChecked(True)
        
        self.save_original_checkbox = QCheckBox("Сохранить оригинальную транскрипцию")
        self.save_original_checkbox.setChecked(True)
        
        options_layout.addRow(self.llm_checkbox)
        options_layout.addRow(self.report_checkbox)
        options_layout.addRow(self.save_original_checkbox)
        
        options_group.setLayout(options_layout)
        layout.addWidget(options_group)
        
        # --- Прогресс ---
        progress_group = QGroupBox("4. Статус обработки")
        progress_layout = QVBoxLayout()
        
        self.progress_label = QLabel("Готов к обработке")
        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)
        self.progress_bar.setVisible(False)
        
        progress_layout.addWidget(self.progress_label)
        progress_layout.addWidget(self.progress_bar)
        
        progress_group.setLayout(progress_layout)
        layout.addWidget(progress_group)
        
        # --- Результаты ---
        results_group = QGroupBox("5. Результаты")
        results_layout = QVBoxLayout()
        
        self.results_text = QTextEdit()
        self.results_text.setReadOnly(True)
        self.results_text.setPlaceholderText("Результаты обработки появятся здесь")
        self.results_text.setMinimumHeight(200)
        
        results_layout.addWidget(self.results_text)
        results_group.setLayout(results_layout)
        layout.addWidget(results_group)
        
        # --- Кнопки управления ---
        button_layout = QHBoxLayout()
        
        self.start_btn = QPushButton("▶ Начать транскрибирование")
        self.start_btn.setStyleSheet("""
            QPushButton {
                background-color: #4CAF50;
                color: white;
                font-weight: bold;
                padding: 10px;
                border-radius: 5px;
            }
            QPushButton:hover {
                background-color: #45a049;
            }
            QPushButton:disabled {
                background-color: #cccccc;
            }
        """)
        self.start_btn.clicked.connect(self.start_transcription)
        
        clear_btn = QPushButton("🗑 Очистить результаты")
        clear_btn.clicked.connect(lambda: self.results_text.clear())
        
        button_layout.addWidget(self.start_btn, 1)
        button_layout.addWidget(clear_btn)
        
        layout.addLayout(button_layout)
        
        widget.setLayout(layout)
        return widget
    
    def create_settings_tab(self):
        """Создание вкладки настроек"""
        widget = QWidget()
        layout = QVBoxLayout()
        
        # --- Модель Whisper ---
        model_group = QGroupBox("Модель Whisper")
        model_layout = QFormLayout()
        
        self.model_path_input = QLineEdit()
        self.model_path_input.setText(str(self.model_path))
        
        browse_model_btn = QPushButton("Обзор...")
        browse_model_btn.clicked.connect(self.browse_model_path)
        
        model_path_layout = QHBoxLayout()
        model_path_layout.addWidget(self.model_path_input, 1)
        model_path_layout.addWidget(browse_model_btn)
        
        model_layout.addRow("Путь к модели:", model_path_layout)
        
        self.device_combo = QComboBox()
        self.device_combo.addItems(["auto", "cuda", "cpu"])
        model_layout.addRow("Устройство:", self.device_combo)
        
        self.dtype_combo = QComboBox()
        self.dtype_combo.addItems(["float32", "float16", "bfloat16"])
        model_layout.addRow("Тип данных:", self.dtype_combo)
        
        model_group.setLayout(model_layout)
        layout.addWidget(model_group)
        
        # --- OpenRouter API ---
        api_group = QGroupBox("OpenRouter API (для LLM-коррекции)")
        api_layout = QFormLayout()
        
        self.api_key_input = QLineEdit()
        self.api_key_input.setEchoMode(QLineEdit.EchoMode.Password)
        self.api_key_input.setPlaceholderText("Введите ваш API ключ OpenRouter")
        api_layout.addRow("API Ключ:", self.api_key_input)
        
        self.model_combo = QComboBox()
        self.model_combo.addItems([
            "gpt-4o",
            "claude-3-opus",
            "gemini-pro",
            "gpt-4-turbo"
        ])
        api_layout.addRow("Модель LLM:", self.model_combo)
        
        api_group.setLayout(api_layout)
        layout.addWidget(api_group)
        
        # --- Медицинские термины ---
        terms_group = QGroupBox("База медицинских терминов")
        terms_layout = QFormLayout()
        
        self.terms_path_input = QLineEdit()
        self.terms_path_input.setText(str(Path(__file__).parent.parent / "medical_terms.txt"))
        
        browse_terms_btn = QPushButton("Обзор...")
        browse_terms_btn.clicked.connect(self.browse_terms_path)
        
        terms_path_layout = QHBoxLayout()
        terms_path_layout.addWidget(self.terms_path_input, 1)
        terms_path_layout.addWidget(browse_terms_btn)
        
        terms_layout.addRow("Путь к файлу терминов:", terms_path_layout)
        
        terms_group.setLayout(terms_layout)
        layout.addWidget(terms_group)
        
        layout.addStretch()
        
        # Кнопка сохранения
        save_settings_btn = QPushButton("💾 Сохранить настройки")
        save_settings_btn.clicked.connect(self.save_settings)
        layout.addWidget(save_settings_btn)
        
        widget.setLayout(layout)
        return widget
    
    def apply_styles(self):
        """Применение стилей к приложению"""
        style = f"""
            QMainWindow {{
                background-color: {UIColors.BACKGROUND};
            }}
            QPushButton {{
                background-color: {UIColors.PRIMARY};
                color: white;
                border: none;
                padding: 8px 16px;
                border-radius: 4px;
                font-size: 14px;
            }}
            QPushButton:hover {{
                background-color: {UIColors.PRIMARY_HOVER};
            }}
            QPushButton:disabled {{
                background-color: {UIColors.DISABLED};
                color: {UIColors.TEXT_DISABLED};
            }}
            QProgressBar {{
                border: 1px solid {UIColors.BORDER};
                border-radius: 5px;
                text-align: center;
            }}
            QProgressBar::chunk {{
                background-color: {UIColors.SUCCESS};
            }}
            QGroupBox {{
                font-weight: bold;
                border: 1px solid {UIColors.BORDER};
                border-radius: 5px;
                margin-top: 10px;
                padding-top: 10px;
            }}
            QGroupBox::title {{
                subcontrol-origin: margin;
                left: 10px;
                padding: 0 3px 0 3px;
            }}
            QLineEdit, QTextEdit, QComboBox, QSpinBox {{
                border: 1px solid {UIColors.BORDER};
                border-radius: 4px;
                padding: 5px;
                background-color: white;
            }}
            QLabel {{
                color: {UIColors.TEXT};
            }}
        """
        self.setStyleSheet(style)
    
    def browse_audio_file(self):
        """Выбор аудиофайла"""
        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "Выберите аудиофайл",
            "",
            "Audio Files (*.wav *.mp3 *.m4a);;All Files (*)"
        )
        if file_path:
            self.audio_path = file_path
            self.file_path_label.setText(file_path)
    
    def browse_model_path(self):
        """Выбор пути к модели"""
        path = QFileDialog.getExistingDirectory(
            self,
            "Выберите папку с моделью Whisper"
        )
        if path:
            self.model_path_input.setText(path)
    
    def browse_terms_path(self):
        """Выбор пути к файлу терминов"""
        file_path, _ = QFileDialog.getOpenFileName(
            self,
            "Выберите файл с медицинскими терминами",
            "",
            "Text Files (*.txt);;All Files (*)"
        )
        if file_path:
            self.terms_path_input.setText(file_path)
    
    def open_patient_dialog(self):
        """Открытие диалога ввода данных пациента"""
        dialog = PatientDataDialog(self)
        if dialog.exec() == QDialog.DialogCode.Accepted:
            self.patient_data = dialog.get_data()
            self.update_patient_info_label()
    
    def update_patient_info_label(self):
        """Обновление метки с информацией о пациенте"""
        if self.patient_data:
            text = f"Пациент: {self.patient_data.get('patient_name', 'Не указано')}"
            self.patient_info_label.setText(text)
            self.patient_info_label.setStyleSheet("color: #4CAF50; font-weight: bold;")
        else:
            self.patient_info_label.setText("Данные пациента не заполнены")
            self.patient_info_label.setStyleSheet("color: #ff9800; font-style: italic;")
    
    def save_settings(self):
        """Сохранение настроек"""
        try:
            # Здесь можно добавить сохранение настроек в конфиг файл
            QMessageBox.information(
                self,
                "Успешно",
                "Настройки сохранены!"
            )
        except Exception as e:
            QMessageBox.critical(
                self,
                "Ошибка",
                f"Ошибка при сохранении настроек: {e}"
            )
    
    def start_transcription(self):
        """Запуск транскрибирования"""
        # Проверка выбран ли файл
        if not self.audio_path:
            QMessageBox.warning(
                self,
                "Ошибка",
                "Пожалуйста, выберите аудиофайл!"
            )
            return
        
        # Проверка наличие файла
        if not Path(self.audio_path).exists():
            QMessageBox.critical(
                self,
                "Ошибка",
                f"Файл не найден: {self.audio_path}"
            )
            return
        
        # Проверка данных пациента если нужен отчёт
        if self.report_checkbox.isChecked() and not self.patient_data:
            QMessageBox.warning(
                self,
                "Ошибка",
                "Для создания отчёта необходимо заполнить данные пациента!"
            )
            return
        
        # Отключение кнопки запуска
        self.start_btn.setEnabled(False)
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(0)
        
        # Создание конфига пайплайна
        try:
            from pipeline.pipeline_config import PipelineConfig
            
            config = PipelineConfig(
                model_path=Path(self.model_path_input.text()),
                device=self.device_combo.currentText(),
                dtype=self.dtype_combo.currentText(),
                medical_terms_file=Path(self.terms_path_input.text()),
                openai_api_key=self.api_key_input.text() or None,
                openai_model=self.model_combo.currentText(),
                correction_enabled=self.llm_checkbox.isChecked(),
                save_original=self.save_original_checkbox.isChecked(),
                save_corrected=True,
                generate_report=self.report_checkbox.isChecked()
            )
        except Exception as e:
            QMessageBox.critical(
                self,
                "Ошибка конфигурации",
                f"Ошибка при создании конфига: {e}"
            )
            self.start_btn.setEnabled(True)
            self.progress_bar.setVisible(False)
            return
        
        # Запуск воркера
        self.worker = TranscriptionWorker(
            self.audio_path,
            config,
            self.patient_data
        )
        
        self.worker.signals.progress.connect(self.on_progress)
        self.worker.signals.finished.connect(self.on_finished)
        self.worker.signals.error.connect(self.on_error)
        
        self.worker.start()
    
    def on_progress(self, message: str):
        """Обновление прогресса"""
        self.progress_label.setText(message)
        self.progress_bar.setValue(min(self.progress_bar.value() + 20, 90))
    
    def on_finished(self, result: dict):
        """Завершение обработки"""
        self.progress_bar.setValue(100)
        self.start_btn.setEnabled(True)
        
        # Вывод результатов
        output = "=" * 60 + "\n"
        output += "РЕЗУЛЬТАТЫ ОБРАБОТКИ\n"
        output += "=" * 60 + "\n\n"
        
        if "transcription_original" in result:
            output += "ОРИГИНАЛЬНАЯ ТРАНСКРИПЦИЯ:\n"
            output += "-" * 40 + "\n"
            output += result["transcription_original"] + "\n\n"
        
        if "transcription_corrected" in result:
            output += "СКОРРЕКТИРОВАННАЯ ТРАНСКРИПЦИЯ:\n"
            output += "-" * 40 + "\n"
            output += result["transcription_corrected"] + "\n\n"
        
        if "report_path" in result:
            output += "✓ Отчёт успешно создан:\n"
            output += f"  {result['report_path']}\n\n"
        
        output += "=" * 60 + "\n"
        output += "Обработка завершена успешно!"
        
        self.results_text.setText(output)
        
        QMessageBox.information(
            self,
            "Успешно",
            "Транскрибирование завершено!"
        )
    
    def on_error(self, error_message: str):
        """Обработка ошибки"""
        self.progress_bar.setVisible(False)
        self.start_btn.setEnabled(True)
        
        self.results_text.setText(f"ОШИБКА:\n{error_message}")
        
        QMessageBox.critical(
            self,
            "Ошибка обработки",
            f"Произошла ошибка:\n{error_message}"
        )


def main():
    """Запуск приложения"""
    from PyQt6.QtWidgets import QApplication
    
    app = QApplication(sys.argv)
    window = MedicalTranscriptionApp()
    window.show()
    sys.exit(app.exec())


if __name__ == "__main__":
    # Базовое логирование
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s'
    )
    main()