import gradio as gr
import torch
from sentence_transformers import SentenceTransformer, util
from transformers import AutoTokenizer, AutoModelForCausalLM
from pypdf import PdfReader
import os
import re

# --- 1. Carregamento dos Modelos ---
print("A carregar o modelo de recuperação (Sentence Transformer)...")
retriever_model = SentenceTransformer('all-MiniLM-L6-v2')

print("A carregar o modelo de geração (DeepSeek)...")
generator_tokenizer = AutoTokenizer.from_pretrained(
    'deepseek-ai/deepseek-coder-1.3b-instruct', 
    trust_remote_code=True
)
generator_model = AutoModelForCausalLM.from_pretrained(
    'deepseek-ai/deepseek-coder-1.3b-instruct', 
    trust_remote_code=True
)
print("Modelos carregados com sucesso!")

# --- 2. Função de Processamento de Ficheiros (Chunking Estruturado) ---
def process_files(files):
    if not files:
        return None, "Por favor, envie um ou mais ficheiros."
    
    knowledge_text = ""
    for file in files:
        file_path = file.name
        if file_path.endswith(".pdf"):
            try:
                reader = PdfReader(file_path)
                for page in reader.pages:
                    page_text = page.extract_text()
                    if page_text:
                        knowledge_text += page_text + "\n\n"
            except Exception as e:
                return None, f"Erro ao ler o ficheiro PDF {os.path.basename(file_path)}: {e}"
        elif file_path.endswith(".txt"):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    knowledge_text += f.read() + "\n\n"
            except Exception as e:
                return None, f"Erro ao ler o ficheiro TXT {os.path.basename(file_path)}: {e}"

    if not knowledge_text.strip():
        return None, "Não foi possível extrair texto dos ficheiros fornecidos."

    chunk_pattern = r"(?m)(^\d+\..*)"
    text_chunks = [chunk.strip() for chunk in re.split(chunk_pattern, knowledge_text) if chunk.strip()]
    
    structured_chunks = []
    i = 0
    while i < len(text_chunks):
        if re.match(chunk_pattern, text_chunks[i]) and i + 1 < len(text_chunks):
            structured_chunks.append(text_chunks[i] + "\n" + text_chunks[i+1])
            i += 2
        else:
            structured_chunks.append(text_chunks[i])
            i += 1

    if not structured_chunks:
        return None, "O texto extraído não continha blocos de texto válidos para processamento."

    print(f"A processar {len(structured_chunks)} chunks estruturados dos ficheiros...")
    knowledge_base_embeddings = retriever_model.encode(structured_chunks, convert_to_tensor=True, show_progress_bar=True)
    print("Base de conhecimento criada a partir dos ficheiros.")
    
    return (structured_chunks, knowledge_base_embeddings), f"✅ Sucesso! {len(files)} ficheiro(s) processado(s), gerando {len(structured_chunks)} chunks estruturados."


# --- 3. O CÉREBRO DA ANÁLISE: generate_compliance_report ---
def generate_compliance_report(task, knowledge_state, progress=gr.Progress()):
    """
    Esta função orquestra o processo de RAG com melhor feedback e tratamento de erros.
    """
    try:
        if not task:
            return "Por favor, forneça uma tarefa de análise."
        if not knowledge_state or not knowledge_state[0] or knowledge_state[1] is None:
            return "⚠️ A base de conhecimento está vazia. Por favor, processe alguns ficheiros primeiro."

        knowledge_base, knowledge_base_embeddings = knowledge_state
        
        progress(0, desc="A iniciar análise...")
        
        # Passo 1: Recuperação Ampla do Contexto
        progress(0.1, desc="A recuperar contexto relevante do documento...")
        search_query = "Informações completas do certificado de calibração"
        question_embedding = retriever_model.encode(search_query, convert_to_tensor=True)
        cosine_scores = util.cos_sim(question_embedding, knowledge_base_embeddings)
        top_k = min(15, len(knowledge_base))
        top_results = torch.topk(cosine_scores, k=top_k, dim=-1)
        retrieved_context = "\n\n---\n\n".join([knowledge_base[i] for i in top_results.indices[0]])

        progress(0.4, desc="Contexto recuperado. A gerar o relatório com o modelo de IA (pode demorar)...")

        # Passo 2: Geração com "Super-Prompt"
        final_prompt = f"""### Instruction:
Você é um auditor de metrologia a preencher um relatório de conformidade. Com base no 'Contexto do Documento' fornecido, preencha cada item da 'Checklist de Análise' abaixo. Se uma informação não for encontrada no contexto, escreva 'Não encontrado'.

**Contexto do Documento:**
{retrieved_context}

**Checklist de Análise:**
# Relatório de Análise de Conformidade

## 1. Incerteza de Medição
   - **1.1 a 1.2 Casas Decimais / Compatibilidade:**
   - **1.3 Nível de Confiança, Fator de Abrangência (k) e Graus de Liberdade:**
   - **1.4 Declaração de Rastreabilidade dos Resultados:**

## 2. Resultados da Calibração
   - **2.1 a 2.3 Unidades SI, Casas Decimais e Algarismos Significativos:**

## 3. Conformidade da Faixa
   - **3.1 Faixa e Especificações Solicitadas:**

## 4. Condições Ambientais
   - **4.1 Registro das Condições e Incerteza Associada:**

## 5. Identificação do Item
   - **5.1 Descrição e Identificação do Item:**

## 6. Identificação do Método
   - **6.1 Método/Procedimento Utilizado:**

## 7. Identificação do Cliente
   - **7.1 Nome e Endereço do Cliente:**

## 8. Identificação do Laboratório
   - **8.1 Nome e Endereço do Laboratório:**

## 9. Identificação do Certificado
   - **9.1 Número do Certificado:**

## 10. Autorização
   - **10.1 Pessoas Autorizadas:**

### Response:
"""
        input_ids = generator_tokenizer(final_prompt, return_tensors="pt").input_ids
        
        # Parâmetros de geração otimizados
        outputs = generator_model.generate(
            input_ids, 
            max_new_tokens=800,  # Reduzido para maior eficiência
            do_sample=False, 
            pad_token_id=generator_tokenizer.eos_token_id
        )
        
        progress(0.9, desc="A formatar o relatório final...")
        final_report = generator_tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Limpa a resposta para remover o prompt inicial
        if "### Response:" in final_report:
            final_report = final_report.split("### Response:")[1].strip()
            
        progress(1, desc="Análise concluída.")
        return final_report

    except Exception as e:
        # Tratamento de erros para retornar uma mensagem clara ao utilizador
        print(f"Ocorreu um erro durante a geração do relatório: {e}")
        return f"### ⚠️ Ocorreu um erro durante a análise.\n\n**Causa provável:** O modelo de IA pode ter excedido os limites de memória ou tempo. Por favor, tente novamente com um documento mais simples ou verifique os logs para mais detalhes.\n\n**Detalhes do Erro:** {str(e)}"


# --- 4. Interface Gráfica (Simplificada para o novo fluxo) ---
with gr.Blocks(theme=gr.themes.Soft()) as interface:
    knowledge_state = gr.State()
    gr.Markdown(
        """
        # 🤖 Agente de Análise de Conformidade Metrológica (v14 - Robusto)
        **1. Carregue um documento**: Envie um certificado de calibração (`.pdf` ou `.txt`).
        **2. Processe o documento**: Clique no botão para criar a base de conhecimento.
        **3. Inicie a Análise**: Dê uma tarefa ao agente (ex: "Analisar conformidade deste certificado") e clique em "Iniciar Análise".
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            file_uploader = gr.File(label="Carregar Documento (.pdf, .txt)", file_count="multiple", file_types=[".pdf", ".txt"])
            process_button = gr.Button("Processar Documento", variant="primary")
            status_box = gr.Textbox(label="Status do Processamento", interactive=False)
        
        with gr.Column(scale=2):
            task_box = gr.Textbox(label="Tarefa de Análise", placeholder='Ex: Avaliar a conformidade deste certificado de calibração.')
            submit_button = gr.Button("Iniciar Análise", variant="primary")

    with gr.Row():
        report_box = gr.Markdown(label="Relatório Final de Análise")

    # A chamada de click continua a mesma, pois o gr.Progress é gerido dentro da função
    process_button.click(fn=process_files, inputs=[file_uploader], outputs=[knowledge_state, status_box])
    submit_button.click(fn=generate_compliance_report, inputs=[task_box, knowledge_state], outputs=[report_box])

# --- 5. Lançamento do App ---
if __name__ == "__main__":
    interface.launch()