import os
import whisperx
import torch
import re
import json
import difflib
from docx import Document
import gc

class TranscriptionProcessor:
    def __init__(self, device="cpu", model_name="large-v3", compute_type="int8"):
        self.device = device
        self.model_name = model_name
        self.compute_type = compute_type
        self.model = None
        self.align_model_cache = {}
        
        # Garante diretório de modelos (Caminho absoluto para container)
        self.base_dir = os.path.dirname(os.path.abspath(__file__))
        self.model_dir = os.path.join(self.base_dir, "models")
        os.makedirs(self.model_dir, exist_ok=True)

    def load_models(self):
        """Carregamento seguro com gerenciamento de memória"""
        if self.model is None:
            print(f"[PROCESSOR] Carregando modelo Whisper: {self.model_name} em {self.device}")
            # Limpeza prévia
            gc.collect()
            if self.device == "cuda": torch.cuda.empty_cache()
            
            try:
                self.model = whisperx.load_model(
                    self.model_name,
                    self.device,
                    compute_type=self.compute_type,
                    download_root=self.model_dir
                )
                print("[PROCESSOR] Modelo Whisper carregado com sucesso.")
            except Exception as e:
                print(f"[PROCESSOR ERROR] Falha ao carregar modelo: {e}")
                raise e

    def process_docx(self, file_path):
        """Processamento robusto de DOCX com limpeza de caracteres de controle"""
        if not file_path: return ""
        print(f"[DOCX] Processando roteiro: {file_path}")
        try:
            doc = Document(file_path)
            full_text = []
            for para in doc.paragraphs:
                text = para.text.strip()
                if text:
                    # Limpeza de caracteres não-imprimíveis (comum em DOCX vindo do Windows)
                    text = "".join(char for char in text if char.isprintable() or char == "\n")
                    # Normalização de hífens para vírgulas conforme solicitado pelo USER
                    text = re.sub(r'[—–-]\s+', ', ', text)
                    text = re.sub(r'\s+,\s+', ', ', text)
                    full_text.append(text)
            return " ".join(full_text)
        except Exception as e:
            print(f"[DOCX ERROR] {e}")
            return ""

    def transcribe(self, audio_path, language="pt"):
        """Transcrição com sistema de Fallback Blindado"""
        try:
            self.load_models()
            audio = whisperx.load_audio(audio_path)
            
            # Passo 1: Transcrição (Parâmetros mínimos para compatibilidade total)
            print("[WHISPER] Transcrevendo...")
            result = self.model.transcribe(audio, batch_size=8, language=language)
            
            # Passo 2: Alinhamento (Com Try-Except interno para evitar erro 500)
            print("[WHISPER] Alinhando...")
            try:
                if language not in self.align_model_cache:
                    self.align_model_cache[language] = whisperx.load_align_model(
                        language_code=language, device=self.device
                    )
                model_a, metadata = self.align_model_cache[language]
                result = whisperx.align(
                    result["segments"], model_a, metadata, audio, self.device
                )
            except Exception as align_err:
                print(f"[WHISPER WARNING] Falha no alinhamento: {align_err}. Seguindo com transcrição base.")
                # Se o alinhamento falhar, 'result' ainda contém os segmentos da transcrição base

            # Passo 3: Processamento de palavras com segurança contra chaves ausentes
            words = []
            for segment in result["segments"]:
                # Caso o WhisperX mude o nome da chave entre 'words' e 'word_segments'
                word_list = segment.get("words", segment.get("word_segments", []))
                for w in word_list:
                    if "start" in w and "end" in w:
                        words.append({
                            "start": round(w["start"], 3),
                            "end": round(w["end"], 3),
                            "word": w.get("word", w.get("text", "")).strip()
                        })
            
            return words
            
        except Exception as e:
            print(f"[WHISPER ERROR] {str(e)}")
            raise e

    def align_with_script(self, audio_words, script_text):
        """
        CORREÇÃO INTELIGENTE (VSL BLINDADA):
        Compara a transcrição com o roteiro e corrige ortografia/termos técnicos
        preservando o tempo do áudio.
        """
        if not script_text:
            return audio_words

        print("[REFINE] Iniciando correção inteligente baseada no roteiro...")
        
        # 1. Preparação das listas (Original e Limpa para matching)
        script_raw = script_text.split()
        script_clean = [re.sub(r'[^\w]', '', w).lower() for w in script_raw]
        audio_clean = [re.sub(r'[^\w]', '', w['word']).lower() for w in audio_words]

        # 2. Matching de Sequência
        matcher = difflib.SequenceMatcher(None, audio_clean, script_clean)
        opcodes = matcher.get_opcodes()

        refined_words = []
        
        for tag, i1, i2, j1, j2 in opcodes:
            if tag == 'equal':
                # Palavras batem: usamos a grafia exata do roteiro (casing/pontuação)
                for k in range(i2 - i1):
                    word_obj = audio_words[i1 + k].copy()
                    word_obj['word'] = script_raw[j1 + k]
                    refined_words.append(word_obj)
            
            elif tag == 'replace':
                # Caso crítico: setox -> Cetox ou setox31 -> Cetox 31
                # Se o número de palavras for diferente, tentamos fundir para manter o tempo
                if (i2 - i1) == (j2 - j1):
                    # 1 para 1
                    for k in range(i2 - i1):
                        word_obj = audio_words[i1 + k].copy()
                        word_obj['word'] = script_raw[j1 + k]
                        refined_words.append(word_obj)
                else:
                    # M:N (Fusão inteligente)
                    # Pegamos o tempo do primeiro ao último do bloco e aplicamos o texto do roteiro
                    new_word_text = " ".join(script_raw[j1:j2])
                    word_obj = {
                        "start": audio_words[i1]["start"],
                        "end": audio_words[i2-1]["end"],
                        "word": new_word_text
                    }
                    refined_words.append(word_obj)
            
            elif tag == 'delete':
                # Palavra no áudio mas não no roteiro (ad-lib ou erro): mantemos o áudio
                for k in range(i1, i2):
                    refined_words.append(audio_words[k])
            
            elif tag == 'insert':
                # Palavra no roteiro mas não detectada pelo Whisper: ignoramos para não quebrar o tempo
                # (Ou poderíamos interpolar, mas para Slides VSL é melhor ignorar)
                pass

        print(f"[REFINE] Concluído. {len(refined_words)} palavras na saída final.")
        return refined_words

    def correct_orthography(self, words):
        """Correções rápidas pós-processamento"""
        for w in words:
            # Limpeza básica de detritos
            w["word"] = w["word"].replace(" ,", ",").replace(",,", ",")
        return words

    def generate_json(self, words):
        """Garante formato JSON estável para o WebApp"""
        return {"words": words}