import streamlit as st
import chromadb
import time
from llama_cpp import Llama
from sentence_transformers import SentenceTransformer
import os

# Initialize Vector DB
chroma_client = chromadb.PersistentClient(path="./clause_index")
collection = chroma_client.get_or_create_collection("legal_clauses")

# Load Embedding Model
embed_model = SentenceTransformer("BAAI/bge-small-en-v1.5")

# Load Local LLM (GGUF)
llm = Llama(
    model_path="mistral-7b-instruct-v0.2.Q6_K.gguf",
    n_ctx=2300,
    n_gpu_layers=38,
    n_batch=512
)

# Generate answer using RAG
def generate_answer(query, context, max_tokens=512):
    prompt = f"""
You are an AI legal assistant trained to extract and summarize specific contract clauses from corporate legal documents.

Your primary role is to answer legal queries based strictly on the provided context or retrieved clauses. Ensure that all responses are accurate, clearly worded, and aligned with standard legal interpretation, but avoid offering legal advice.

Query: {query}

Context:
{context}

Answer:
"""
    output = llm(
        prompt=prompt.strip(),
        max_tokens=max_tokens,
        temperature=0.0,
        top_p=0.95,
        top_k=40,
        stop=["</s>"]
    )
    return output["choices"][0]["text"].strip()

# Evaluation Prompts
groundedness_prompt = """You are tasked with rating..."""  # replace with your full prompt
relevance_prompt = """You are tasked with rating..."""     # replace with your full prompt

user_template = """
###Question
{question}

###Context
{context}

###Answer
{answer}
"""

# Evaluation function
def evaluate_output(question, context, answer):
    g_prompt = f"""[INST]{groundedness_prompt.strip()}\n{user_template.format(question=question, context=context, answer=answer)}[/INST]"""
    r_prompt = f"""[INST]{relevance_prompt.strip()}\n{user_template.format(question=question, context=context, answer=answer)}[/INST]"""
    
    with st.spinner("Evaluating groundedness..."):
        g_eval = llm(prompt=g_prompt, max_tokens=256, stop=["</s>"])["choices"][0]["text"].strip()
    with st.spinner("Evaluating relevance..."):
        r_eval = llm(prompt=r_prompt, max_tokens=256, stop=["</s>"])["choices"][0]["text"].strip()
    
    return g_eval, r_eval

# Streamlit UI
st.set_page_config(page_title="GL_LegalMind - Contract Clause Retriever", layout="centered")

st.title("GL_LegalMind - Summarized Contract Clause Retrieval Using RAG and LLMs")

query = st.text_area("Enter your legal question:", height=100)

if st.button("Run RAG Retrieval"):
    if not query.strip():
        st.warning("Please enter a query.")
    else:
        with st.spinner("Searching clauses and generating answer..."):
            query_embedding = embed_model.encode([query])[0]
            results = collection.query(query_embeddings=[query_embedding], n_results=3)
            context = "\n\n".join(results["documents"][0])
            answer = generate_answer(query, context)
            st.success("Answer Generated!")
            st.markdown("### Answer")
            st.write(answer)
            st.markdown("### Context Used")
            st.code(context)
            
            if st.button("Evaluate Answer (Groundedness & Relevance)"):
                g_eval, r_eval = evaluate_output(query, context, answer)
                st.markdown("### Groundedness Evaluation")
                st.write(g_eval)
                st.markdown("### Relevance Evaluation")
                st.write(r_eval)