import gradio as gr
import os
import re
import numpy as np
from metrology_rag import MetrologyRAGSystem

class CalibrationAssistant:
    def __init__(self):
        # Inicializa com documentos pré-carregados
        self.rag_system = MetrologyRAGSystem()
        self.rag_system.initialize_system("rag_documents")
        
    def technical_query(self, query):
        try:
            response = self.rag_system.generate_technical_response(query)
            return self._clean_response(response)
        except Exception as e:
            return f"Erro na consulta: {str(e)}"

    def analyze_report(self, file):
        try:
            analysis = self.rag_system.analyze_metrology_report(file.name)
            return self._clean_response(analysis)
        except Exception as e:
            return f"Erro na análise: {str(e)}"

    def _clean_response(self, text):
        return re.sub(r'\x1B\[[0-?]*[ -/]*[@-~]', '', text)

assistant = CalibrationAssistant()

css = """
footer {visibility: hidden}
.markdown {max-width: 95% !important}
"""

with gr.Blocks(css=css, title="Assistente de Metrologia") as demo:
    gr.Markdown("# 🔧 Assistente de Metrologia Certificada")
    
    with gr.Tab("Consulta Técnica"):
        gr.Markdown("### Sistema RAG com Documentação Pré-Carregada")
        query_input = gr.Textbox(
            label="Faça sua consulta técnica",
            placeholder="Ex: Qual o procedimento para calibração de micrômetros?..."
        )
        query_btn = gr.Button("Enviar")
        response_output = gr.Markdown(label="Resposta Técnica Certificada")
        
    with gr.Tab("Análise de Relatório"):
        gr.Markdown("### Verificação de Conformidade de Relatórios")
        report_input = gr.File(
            label="Carregue seu relatório de calibração (PDF)",
            file_types=[".pdf"]
        )
        analyze_btn = gr.Button("Analisar")
        analysis_output = gr.Markdown(label="Resultado da Análise")
    
    query_btn.click(
        assistant.technical_query,
        inputs=query_input,
        outputs=response_output
    )
    
    analyze_btn.click(
        assistant.analyze_report,
        inputs=report_input,
        outputs=analysis_output
    )

demo.launch()