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

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

class SpaceInterface:
    def __init__(self):
        self.rag_system = None
        self.current_response = ""
        
    def initialize_system(self, files):
        try:
            temp_dir = "/tmp/metrology_pdfs"
            os.makedirs(temp_dir, exist_ok=True)
            
            for file in files:
                file_path = os.path.join(temp_dir, os.path.basename(file.name))
                with open(file_path, "wb") as f:
                    f.write(file.read())
            
            self.rag_system = MetrologyRAGSystem()
            self.rag_system.initialize_system(temp_dir)
            return f"✅ Sistema inicializado com {len(self.rag_system.documents)} documentos!"
        except Exception as e:
            return f"❌ Erro: {str(e)}"

    def process_query(self, query):
        if not self.rag_system:
            return "⚠️ Inicialize o sistema primeiro com documentos PDF!"
        
        try:
            response = self.rag_system.generate_technical_response(query)
            self.current_response = response
            return self._clean_response(response)
        except Exception as e:
            return f"Erro: {str(e)}"

    def generate_report(self, report_title):
        if not self.rag_system or not self.current_response:
            return None, "⚠️ Execute uma consulta primeiro!"
        
        try:
            report_path = "/tmp/relatorio_tecnico.md"
            self.rag_system.generate_report(report_title, self.current_response, report_path)
            return report_path, "✅ Relatório gerado com sucesso!"
        except Exception as e:
            return None, f"❌ Falha: {str(e)}"

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

interface = SpaceInterface()

with gr.Blocks(css=css, title="Metrology Assistant") as demo:
    gr.Markdown("# 🔧 Assistente de Metrologia Inteligente")
    
    with gr.Row():
        with gr.Column(scale=2):
            file_input = gr.File(
                file_count="multiple",
                file_types=[".pdf"],
                label="Upload de Documentos Técnicos"
            )
            init_btn = gr.Button("Inicializar Sistema")
            init_status = gr.Textbox(label="Status", interactive=False)
        
        with gr.Column(scale=3):
            query_input = gr.Textbox(
                label="Consulta Técnica",
                placeholder="Digite sua pergunta sobre metrologia..."
            )
            query_btn = gr.Button("Enviar Consulta")
            response_output = gr.Markdown(label="Resposta Técnica")
            
            report_input = gr.Textbox(
                label="Título do Relatório",
                placeholder="Digite o título do relatório..."
            )
            report_btn = gr.Button("Gerar Relatório")
            report_download = gr.File(label="Download do Relatório")
            report_status = gr.Markdown()

    init_btn.click(
        interface.initialize_system,
        inputs=file_input,
        outputs=init_status
    )
    
    query_btn.click(
        interface.process_query,
        inputs=query_input,
        outputs=response_output
    )
    
    report_btn.click(
        interface.generate_report,
        inputs=report_input,
        outputs=[report_download, report_status]
    )

demo.launch(debug=True)