import os
import fitz  # PyMuPDF
from pathlib import Path
import gradio as gr
import google.generativeai as genai
import time
import tempfile
from PIL import Image

class PDFTranscriberGradio:
    def __init__(self):
        self.temp_files = []
        
    def get_selected_model(self, model_choice):
        """Retorna o ID do modelo selecionado"""
        return model_choice.split(" - ")[0].strip()
    
    def validate_pdf_file(self, file_path):
        """Valida se o arquivo é um PDF usando múltiplos métodos"""
        try:
            # Método 1: Verificar se o arquivo existe
            if not os.path.exists(file_path):
                return False, "Arquivo não encontrado"
            
            # Método 2: Verificar tamanho
            file_size = os.path.getsize(file_path)
            if file_size == 0:
                return False, "Arquivo está vazio"
            
            # Método 3: Tentar abrir com PyMuPDF
            try:
                doc = fitz.open(file_path)
                page_count = len(doc)
                doc.close()
                if page_count > 0:
                    return True, f"PDF válido com {page_count} página(s)"
                else:
                    return False, "PDF não contém páginas"
            except Exception as e:
                pass
            
            # Método 4: Verificar header do arquivo
            try:
                with open(file_path, 'rb') as f:
                    header = f.read(8)
                    if header.startswith(b'%PDF-'):
                        return True, "PDF detectado por header"
            except Exception as e:
                pass
            
            return False, "Arquivo não é um PDF válido"
            
        except Exception as e:
            return False, f"Erro na validação: {str(e)}"
    
    def convert_pdf_page_to_image(self, pdf_document, page_number):
        """Converte uma página do PDF em imagem e salva temporariamente"""
        try:
            page = pdf_document[page_number]
            pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
            img_data = pix.tobytes("png")
            
            # Criar arquivo temporário com nome único baseado no timestamp
            temp_filename = f'page_{page_number}_{int(time.time()*1000)}_{os.getpid()}.png'
            temp_path = os.path.join(tempfile.gettempdir(), temp_filename)
            
            with open(temp_path, 'wb') as temp_file:
                temp_file.write(img_data)
            
            # Adicionar à lista de arquivos temporários
            self.temp_files.append(temp_path)
            return temp_path
        except Exception as e:
            raise Exception(f"Erro ao converter página {page_number}: {str(e)}")
    
    def analyze_page_with_gemini(self, image_path, page_num, model_name, temperature, top_p, top_k):
        """Analisa uma página usando o modelo Gemini selecionado"""
        try:
            # Carregar a imagem
            img = Image.open(image_path)
            
            # Configurar o modelo Gemini
            model = genai.GenerativeModel(model_name)
            
            # Prompt otimizado para transcrição
            prompt = """Analise esta página do documento PDF em detalhes e forneça uma transcrição completa. Inclua:

1. **TEXTO COMPLETO**: Transcreva todo o texto visível, mantendo a formatação original
2. **ESTRUTURA**: Descreva títulos, subtítulos, parágrafos, listas
3. **ELEMENTOS VISUAIS**: Descreva tabelas, gráficos, imagens, logotipos
4. **FORMATAÇÃO**: Mencione textos em negrito, itálico, cores especiais
5. **LAYOUT**: Como o conteúdo está organizado na página
6. **INFORMAÇÕES IMPORTANTES**: Destaque dados relevantes, números, datas, certificações

Seja preciso e detalhado na transcrição."""
            
            # Gerar resposta
            response = model.generate_content([prompt, img], generation_config={
                'temperature': temperature,
                'max_output_tokens': 8192,
                'top_p': top_p,
                'top_k': int(top_k),
            })
            
            # Fechar a imagem explicitamente
            img.close()
            
            return response.text
            
        except Exception as e:
            return f"Erro ao analisar página {page_num}: {str(e)}"
    
    def cleanup_temp_files(self):
        """Limpa todos os arquivos temporários"""
        cleaned_count = 0
        for temp_file in self.temp_files:
            try:
                if os.path.exists(temp_file):
                    # Tentar várias vezes em caso de arquivo em uso
                    for attempt in range(3):
                        try:
                            os.unlink(temp_file)
                            cleaned_count += 1
                            break
                        except PermissionError:
                            time.sleep(0.5)
                        except Exception as e:
                            if attempt == 2:  # Última tentativa
                                print(f"Aviso: Não foi possível remover {temp_file}: {e}")
                                break
            except Exception as e:
                print(f"Erro ao limpar arquivo temporário {temp_file}: {e}")
        
        self.temp_files = []
        print(f"Limpeza concluída: {cleaned_count} arquivos removidos")
    
    def extract_text_from_pdf(self, pdf_path):
        """Extrai texto diretamente do PDF"""
        try:
            doc = fitz.open(pdf_path)
            text_content = ""
            
            for page_num in range(len(doc)):
                page = doc[page_num]
                page_text = page.get_text()
                if page_text.strip():
                    text_content += f"\n=== PÁGINA {page_num + 1} - TEXTO EXTRAÍDO ===\n"
                    text_content += page_text
                    text_content += "\n" + "="*50 + "\n"
            
            doc.close()
            return text_content if text_content.strip() else "Nenhum texto extraível encontrado no PDF"
            
        except Exception as e:
            raise Exception(f"Erro ao extrair texto do PDF: {str(e)}")
    
    def process_pdf(self, pdf_file, api_key, model_choice, temperature, top_p, top_k):
        """Função principal para processar o PDF"""
        status_updates = ["🔄 Iniciando processamento..."]
        yield "\n".join(status_updates), None, None
        
        try:
            # Verificar se o arquivo foi fornecido
            if pdf_file is None:
                raise Exception("❌ Nenhum arquivo enviado")
            
            # Obter caminho do arquivo
            file_path = pdf_file.name if hasattr(pdf_file, 'name') else pdf_file
            
            status_updates.append(f"📁 Arquivo: {os.path.basename(file_path)}")
            yield "\n".join(status_updates), None, None
            
            # Validar arquivo PDF
            is_valid, validation_msg = self.validate_pdf_file(file_path)
            if not is_valid:
                raise Exception(f"❌ {validation_msg}")
            
            status_updates.append(f"✅ {validation_msg}")
            yield "\n".join(status_updates), None, None
            
            # Verificar API Key
            if not api_key or api_key.strip() == "":
                raise Exception("❌ API Key do Gemini não fornecida")
            
            # Configurar API Gemini
            genai.configure(api_key=api_key.strip())
            
            # Obter modelo selecionado
            model_name = self.get_selected_model(model_choice)
            status_updates.append(f"🤖 Modelo: {model_name}")
            yield "\n".join(status_updates), None, None
            
            # Abrir PDF
            doc = fitz.open(file_path)
            total_pages = len(doc)
            
            status_updates.append(f"📄 PDF: {total_pages} página(s)")
            yield "\n".join(status_updates), None, None
            
            # Extrair texto tradicional
            status_updates.append("📝 Extraindo texto...")
            yield "\n".join(status_updates), None, None
            
            extracted_text = self.extract_text_from_pdf(file_path)
            
            # Analisar páginas com IA
            status_updates.append("🔍 Analisando com IA...")
            yield "\n".join(status_updates), None, None
            
            page_analyses = []
            
            for page_num in range(total_pages):
                status_updates.append(f"🖼️ Página {page_num + 1}/{total_pages}...")
                yield "\n".join(status_updates), None, None
                
                try:
                    image_path = self.convert_pdf_page_to_image(doc, page_num)
                    analysis = self.analyze_page_with_gemini(
                        image_path, 
                        page_num + 1, 
                        model_name, 
                        temperature, 
                        top_p, 
                        top_k
                    )
                    page_analyses.append(f"\n{'='*60}\n📄 ANÁLISE DA PÁGINA {page_num + 1}\n{'='*60}\n{analysis}")
                    
                except Exception as e:
                    error_msg = f"❌ Erro na página {page_num + 1}: {str(e)}"
                    page_analyses.append(f"\n{'='*60}\n📄 PÁGINA {page_num + 1} - ERRO\n{'='*60}\n{error_msg}")
            
            # Fechar documento
            doc.close()
            
            # Compilar resultado final
            status_updates.append("📋 Compilando resultado...")
            yield "\n".join(status_updates), None, None
            
            final_output = f"""{'='*80}
🔍 TRANSCRIÇÃO COMPLETA DO PDF
{'='*80}

📄 ARQUIVO: {os.path.basename(file_path)}
🤖 MODELO: {model_name}
📊 PÁGINAS: {total_pages}
⏰ PROCESSADO: {time.strftime('%Y-%m-%d %H:%M:%S')}

{'='*80}
📝 TEXTO EXTRAÍDO DIRETAMENTE
{'='*80}
{extracted_text}

{'='*80}
🔍 ANÁLISE VISUAL DETALHADA
{'='*80}
{chr(10).join(page_analyses)}

{'='*80}
✅ PROCESSAMENTO CONCLUÍDO
{'='*80}
"""
            
            # Salvar resultado
            output_filename = f"transcricao_pdf_{int(time.time())}.txt"
            output_path = os.path.join(tempfile.gettempdir(), output_filename)
            
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(final_output)
            
            status_updates.append("✅ Processamento concluído!")
            status_updates.append(f"💾 Arquivo: {output_filename}")
            
            yield "\n".join(status_updates), final_output, output_path
            
        except Exception as e:
            error_msg = f"❌ ERRO: {str(e)}"
            status_updates.append(error_msg)
            yield "\n".join(status_updates), None, None
            
        finally:
            # Limpar arquivos temporários
            self.cleanup_temp_files()
    
    def create_interface(self):
        """Cria interface Gradio"""
        with gr.Blocks(
            title="Transcritor de PDF com Gemini AI",
            theme=gr.themes.Soft()
        ) as interface:
            
            gr.Markdown("""
            # 🔍 Transcritor de PDF com Gemini AI
            
            **Ferramenta para transcrição e análise de documentos PDF**
            
            ✨ **Recursos:**
            - 📝 Extração de texto nativo
            - 🖼️ Análise visual com IA
            - 📊 Descrição de elementos visuais
            - 💾 Download do resultado
            """)
            
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### ⚙️ Configurações")
                    
                    api_key = gr.Textbox(
                        label="🔑 API Key Gemini", 
                        type="password",
                        placeholder="Insira sua API key do Google Gemini"
                    )
                    
                    # CORREÇÃO PRINCIPAL: Remover validação restritiva
                    pdf_file = gr.File(
                        label="📄 Arquivo PDF",
                        file_types=None,  # Aceita qualquer tipo
                        file_count="single"
                    )
                    
                    model_choice = gr.Dropdown(
                        label="🤖 Modelo", 
                        choices=[
                            "gemini-2.0-flash-exp - Última geração (Experimental)",
                            "gemini-1.5-flash - Rápido e versátil",
                            "gemini-1.5-flash-8b - Alto volume",
                            "gemini-1.5-pro - Raciocínio complexo"
                        ],
                        value="gemini-1.5-pro - Raciocínio complexo"
                    )
                    
                    with gr.Row():
                        temperature = gr.Slider(
                            label="🌡️ Temperatura", 
                            minimum=0.0, 
                            maximum=1.0, 
                            step=0.1, 
                            value=0.1
                        )
                    
                    with gr.Row():
                        top_p = gr.Slider(
                            label="🎯 Top P", 
                            minimum=0.0, 
                            maximum=1.0, 
                            step=0.1, 
                            value=0.9
                        )
                        
                        top_k = gr.Slider(
                            label="🔢 Top K", 
                            minimum=1, 
                            maximum=40, 
                            step=1, 
                            value=40
                        )
                    
                    process_button = gr.Button(
                        "🚀 Transcrever PDF", 
                        variant="primary"
                    )
                
                with gr.Column(scale=2):
                    gr.Markdown("### 📊 Resultados")
                    
                    status_output = gr.Textbox(
                        label="📋 Status", 
                        lines=8,
                        interactive=False
                    )
                    
                    result_output = gr.Textbox(
                        label="📄 Transcrição", 
                        lines=20,
                        show_copy_button=True
                    )
                    
                    file_output = gr.File(
                        label="💾 Download"
                    )
            
            # Conectar função
            process_button.click(
                fn=self.process_pdf,
                inputs=[pdf_file, api_key, model_choice, temperature, top_p, top_k],
                outputs=[status_output, result_output, file_output],
                show_progress=True
            )
            
            gr.Markdown("""
            ### 💡 Dicas
            - Obtenha API Key gratuita em: https://ai.google.dev/
            - Use temperatura baixa (0.1-0.3) para maior precisão
            - Gemini 1.5 Pro é melhor para análise complexa
            """)
            
        return interface
    
    def run(self):
        """Inicia a aplicação"""
        interface = self.create_interface()
        interface.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=False
        )

if __name__ == "__main__":
    app = PDFTranscriberGradio()
    app.run()