import os
import numpy as np
import pandas as pd
import gradio as gr

import faiss
from sentence_transformers import SentenceTransformer
from settings import OUTPUT_DIR, DEVICE
os.environ['WANDB_DISABLED'] = 'true'

    
fine_tuned_model = SentenceTransformer(OUTPUT_DIR, device=DEVICE)
passages         = pd.read_parquet('data/processed/corpus_data.parquet')['text'].tolist()
legal_index      = faiss.read_index('data/retrieval/legal_faiss.index')

def retrieval(emb_model, query, index, top_k=10):
    q_emb = emb_model.encode(
        query, 
        convert_to_numpy=True, 
        normalize_embeddings=True,
    ).astype(np.float32).reshape(1, -1)
    
    scores, indices = index.search(q_emb, top_k)  # shape: (1, top_k)
    
    cand_idxs   = indices[0]
    cand_scores = scores[0]
    cand_texts  = [passages[i] for i in cand_idxs]

    results = [{
        'index': int(cand_idxs[i]),
        'score': float(cand_scores[i]),
        'text': cand_texts[i]
    } for i in range(len(cand_idxs))]
    
    return results

def get_results(query, top_k):
    hits = retrieval(fine_tuned_model, query, legal_index, top_k=top_k)
    
    result = ""
    for rank, h in enumerate(hits, start=1):
        result += f"[Kết quả {rank} - Độ tin cậy={h['score']:.4f}]\n\n{h['text']}\n{'-'*100}\n"
    return result

    
demo = gr.Interface(
    'huggingface/YuITC/bert-base-multilingual-cased-finetuned-VNLegalDocs',
    fn=get_results,
    inputs=[
        gr.Textbox(lines=2, placeholder='Nhập câu hỏi pháp lý của bạn...', label='Câu hỏi'),
        gr.Slider(minimum=5, maximum=20, value=10, step=1, label='Số lượng kết quả'),
    ],
    outputs=gr.Textbox(lines=20, label='Kết quả'),
    title='Vietnamese Legal Document Retrieval System',
    description='🔍 Nhập câu hỏi pháp lý của bạn bằng tiếng Việt để nhận các đoạn văn bản pháp luật liên quan.',
    examples=[
        ['Tội xúc phạm danh dự?'],
        ['Quyền lợi của người lao động?'],
        ['Thủ tục đăng ký kết hôn?'],
    ],
    flagging_mode='never'
)

if __name__ == '__main__':
    demo.launch(share=True)