import torch
import whisper
from typing import Dict, Any, Optional
import os
from preprocessing import clean_audio
from diarization import diarize_segments
from medical_terms import load_turkish_medical_terms
from advanced_term_correction import MedicalTermCorrector
from config import settings

class MedicalASR:
    def __init__(self, config: Dict[str, Any]):
        """
        Tıbbi konuşma tanıma için ASR modülü
        
        Args:
            config: Yapılandırma parametreleri
                - language: Dil kodu (örn. "tr")
                - model: Kullanılacak model adı
                - domain: Alan adı (tıp için "medical")
        """
        self.config = config
        self.language = config.get("language", "tr")
        self.model_name = config.get("model", "whisper-large-v3")
        
        # CUDA kullanılabilirse tercih et
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        
        # Modeli yükle
        self.model = whisper.load_model(self.model_name, device=self.device)
        
        # Türkçe tıbbi terim sözlüğü
        self.medical_terms = load_turkish_medical_terms()
        
        # Gelişmiş terim düzeltici
        self.term_corrector = MedicalTermCorrector(self.medical_terms)
    
    def _load_medical_terms(self) -> Dict[str, str]:
        """Türkçe tıbbi terim sözlüğünü yükler"""
        # Örnek: Bu fonksiyon bir dosyadan ya da veritabanından tıbbi terimleri yükleyebilir
        return {
            "ateş": "ateş",
            "hipertansiyon": "hipertansiyon",
            "miyokard infarktüsü": "miyokard infarktüsü",
            # ... daha fazla tıbbi terim
        }
    
    def transcribe(self, audio_file: str, speaker_diarization: bool = True, enhance_audio: bool = True) -> Dict[str, Any]:
        """
        Ses dosyasını transkribe eder
        
        Args:
            audio_file: Ses dosyasının yolu
            speaker_diarization: Konuşmacı diyarizasyonu yapılsın mı
            enhance_audio: Ses iyileştirme yapılsın mı
            
        Returns:
            Transkripsiyon sonuçları
        """
        # 1. Ses iyileştirme (optional)
        if enhance_audio:
            processed_file = clean_audio(audio_file)
        else:
            processed_file = audio_file
        
        # 2. Transkripsiyon
        transcribe_options = {
            "language": self.language,
            "task": "transcribe",
            "beam_size": 5,  # Daha geniş arama ile alternatif kelimeleri değerlendir
            "best_of": 5,    # En iyi sonucu seç
            "temperature": 0, # Deterministik sonuçlar için
            "initial_prompt": "Çocuk doktoru vizitesi, pediatri uzmanı, çocuk hastalıkları, yenidoğan"  # Modele ipucu ver
        }
        
        result = self.model.transcribe(processed_file, **transcribe_options)
        
        # 3. Tıbbi terimleri düzelt
        corrected_text = self._correct_medical_terms(result["text"])
        result["text"] = corrected_text
        
        # 4. Diyarizasyon isteniyorsa ekle
        if speaker_diarization:
            segments = diarize_segments(processed_file)
            
            # Basit bir yaklaşım: Segmentlerin zaman etiketlerini kullanarak 
            # Whisper'ın döndürdüğü segmentleri konuşmacılara atayabiliriz
            if "segments" in result:
                whisper_segments = result["segments"]
                
                # Her bir diyarizasyon segmenti için
                for diar_seg in segments:
                    # Whisper segmentlerinden hangilerinin bu diyarizasyon segmentinde 
                    # olduğunu bulalım
                    for whisper_seg in whisper_segments:
                        w_start = whisper_seg["start"]
                        w_end = whisper_seg["end"]
                        
                        # Zaman aralıkları çakışıyorsa
                        if (diar_seg["start"] <= w_end and diar_seg["end"] >= w_start):
                            if not diar_seg["text"]:
                                diar_seg["text"] = whisper_seg["text"]
                            else:
                                diar_seg["text"] += " " + whisper_seg["text"]
            
            result["diarization"] = segments
        
        return result
    
    def _correct_medical_terms(self, text: str) -> str:
        """Transkribe edilmiş metindeki tıbbi terimleri düzeltir"""
        # Gelişmiş düzelticiyi kullan
        return self.term_corrector.correct_text(text)

# Basit dosya transkripsiyon fonksiyonu
def transcribe_file(file_path: str) -> str:
    # Bu fonksiyon eski API için burada bırakıldı
    _model = whisper.load_model(settings.ASR_MODEL)
    result = _model.transcribe(file_path)
    return result.get('text', '').strip()