import os
import sys

# --- 1. CHROMA DB FIX FOR HUGGING FACE ---
# ChromaDB requires a newer version of sqlite3 than the one pre-installed on Linux
try:
    __import__('pysqlite3')
    sys.modules['sqlite3'] = sys.modules.pop('pysqlite3')
except ImportError:
    pass # Pass if running locally or if not available

import gradio as gr
from langchain_community.document_loaders import PyPDFLoader
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
# STABLE IMPORT (Matches langchain==0.1.20)
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
# STABLE IMPORT
from langchain.chains import RetrievalQA
from datasets import Dataset
from ragas import evaluate
from ragas.metrics import faithfulness, answer_relevancy

# --- 2. KEY LOADER ---
api_key = os.getenv("OPENAI_API_KEY")

if api_key:
    key_status = "✅ ACTIVE (Loaded from Secrets)"
    os.environ["OPENAI_API_KEY"] = api_key
else:
    key_status = "❌ MISSING (Check Settings -> Secrets)"

def audit_rag(pdf_file, user_question):
    if not api_key:
        return "ERROR: API Key is missing.", "ERROR", "0", "0"

    if not pdf_file or not user_question:
        return "Please upload a PDF and ask a question.", "Waiting...", "0.00", "0.00"

    try:
        # Load & Split
        loader = PyPDFLoader(pdf_file.name)
        documents = loader.load()
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        texts = text_splitter.split_documents(documents)
        
        # RAG Engine
        embeddings = OpenAIEmbeddings(openai_api_key=api_key)
        db = Chroma.from_documents(texts, embeddings)
        retriever = db.as_retriever(search_kwargs={"k": 3})
        
        llm = ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0, openai_api_key=api_key)
        
        qa_chain = RetrievalQA.from_chain_type(
            llm=llm, 
            chain_type="stuff", 
            retriever=retriever,
            return_source_documents=True
        )
        
        # Answer
        result = qa_chain.invoke({"query": user_question})
        generated_answer = result['result']
        source_docs = [doc.page_content for doc in result['source_documents']]
        
        # Ragas Audit
        data = {
            'question': [user_question],
            'answer': [generated_answer],
            'contexts': [source_docs],
            'ground_truth': [""] 
        }
        dataset = Dataset.from_dict(data)
        
        score = evaluate(
            dataset=dataset,
            metrics=[faithfulness, answer_relevancy],
            llm=llm,       
            embeddings=embeddings
        )
        
        audit_results = score.to_pandas()
        faith_score = audit_results.iloc[0]['faithfulness']
        relevancy_score = audit_results.iloc[0]['answer_relevancy']

        verdict = "✅ PASS" if faith_score > 0.8 else "❌ FAIL (Hallucination Detected)"
        
        return generated_answer, verdict, f"{faith_score:.2f}", f"{relevancy_score:.2f}"

    except Exception as e:
        return f"System Error: {str(e)}", "ERROR", "0", "0"

# UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# ⚖️ Veritas: AI Hallucination Auditor")
    gr.Markdown(f"**System Status:** {key_status}")
    gr.Markdown("Upload a document (e.g., Financial Report) and ask a question.")
    
    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="Upload PDF Evidence", file_types=[".pdf"])
            question_input = gr.Textbox(label="Cross-Examination Question", placeholder="e.g., What was the net profit in Q3?")
            submit_btn = gr.Button("Run Audit", variant="primary")
        
        with gr.Column():
            answer_output = gr.Textbox(label="AI Witness Testimony (Answer)")
            with gr.Row():
                verdict_output = gr.Textbox(label="Verdict")
                faith_output = gr.Textbox(label="Faithfulness Score (0-1)")
            relevance_output = gr.Textbox(label="Relevancy Score")

    submit_btn.click(
        audit_rag, 
        inputs=[file_input, question_input], 
        outputs=[answer_output, verdict_output, faith_output, relevance_output]
    )

if __name__ == "__main__":
    demo.launch()