import gradio as gr
import requests
import json
import os
import re

# ==========================================
# ⚙️ CONFIGURATION
# ==========================================
HF_TOKEN = os.environ.get("HF_TOKEN")
# REPLACE THIS WITH YOUR RUNNING ENDPOINT URL:
ENDPOINT_URL = "https://t1v2yv5hzk2zn5x8.us-east-1.aws.endpoints.huggingface.cloud"

# ==========================================
# 🛡️ THE API CALL LOGIC
# ==========================================
def run_custom_audit(question, ground_truths, upstream_answer):
    if "YOUR-UNIQUE-ID" in ENDPOINT_URL:
        return {"error": "ENDPOINT_URL not configured. Please update app.py."}, "🚨 Setup Incomplete"

    system_prompt = """You are an AI Clinical Auditor evaluating a RAG (Retrieval-Augmented Generation) system.
    Compare the Upstream Answer against the provided Ground Truth Passages to answer the User Question.
    Identify any hallucinations, omissions, or statistical drift in the Answer.
    Strictly output your response in the following JSON schema:
    {
      "answer": "<summary of verified facts, or EMPTY if dangerous contradictions exist>",
      "suspicion_flags": [
        {"passage_index": <int>, "reason": "<explanation of the contradiction>", "confidence": <float between 0.0 and 1.0>}
      ]
    }"""
    
    user_prompt = f"USER QUESTION:\n{question}\n\nGROUND TRUTH PASSAGES:\n{ground_truths}\n\nUPSTREAM ANSWER (TO AUDIT):\n{upstream_answer}"
    
    payload = {
        "inputs": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ],
        "parameters": {
            "max_new_tokens": 512,
            "temperature": 0.1
        }
    }
    
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}",
        "Content-Type": "application/json"
    }
    
    try:
        response = requests.post(ENDPOINT_URL, headers=headers, json=payload)
        
        if response.status_code == 200:
            result_text = response.json()[0]["generated_text"]
            
            # Robust JSON isolation
            clean_text = re.sub(r'```json|```', '', result_text).strip()
            try:
                start_idx = clean_text.find('{')
                end_idx = clean_text.rfind('}') + 1
                if start_idx != -1 and end_idx != 0:
                    clean_text = clean_text[start_idx:end_idx]
                
                # Parse into an actual Python dictionary for the gr.JSON component
                parsed_json = json.loads(clean_text)
                return parsed_json, "✅ Audit Executed Successfully"
            except Exception as parse_err:
                # Fallback if the model outputs malformed JSON
                return {"raw_output": clean_text, "parse_error": str(parse_err)}, "⚠️ Partial Success: JSON Parsing Failed"
        else:
            return {"error": f"API Error {response.status_code}", "details": response.text}, "🚨 Endpoint Failure"
            
    except Exception as e:
        return {"error": "Request Failed", "details": str(e)}, "🚨 Connection Error"


# ==========================================
# 🎨 THE GRADIO UI (RAG EVALUATION MODE)
# ==========================================
with gr.Blocks(theme=gr.themes.Base(), title="PropagationShield v2.0") as demo:
    gr.Markdown("# 🛡️ PropagationShield v2.0: RAG Context Auditor")
    gr.Markdown("Identify the **Propagation Cascade** live. Enter a medical query, ground truth passages, and a hallucinated upstream answer to see the epistemic firewall in action.")
    gr.Markdown("If u see a 503 error, please try again after 3-4 minutes as the server finishes starting up")
    
    with gr.Row():
        # ---------------- LEFT SIDE: INPUTS ----------------
        with gr.Column(scale=1):
            gr.Markdown("### 📥 1. System Inputs")
            
            question_input = gr.Textbox(
                label="User Question", 
                lines=2, 
                placeholder="e.g., Who was behind the theory of relativity?"
            )
            
            ground_truth_input = gr.Textbox(
                label="Ground Truths (Numbered Passages)", 
                lines=6, 
                placeholder="[0] Albert Einstein developed the theory of relativity in the early 20th century.\n[1] The theory revolutionized physics and our understanding of space and time.\n[2] Isaac Newton originally proposed relativity in the 18th century."
            )
            
  
            
            run_btn = gr.Button("RUN PROPAGATION SHIELD 🛡️", variant="primary", size="lg")
            
        # ---------------- RIGHT SIDE: OUTPUTS ----------------
        with gr.Column(scale=1):
            gr.Markdown("### 🔍 2. Auditor Output")
            status_indicator = gr.Markdown("**Status:** Waiting for input...")
            
            # Using gr.JSON gives a beautifully formatted, collapsible JSON tree viewer
            json_output = gr.JSON(label="Structured Audit Results")

    gr.HTML("<hr style='border: 1px solid #ddd; margin: 20px 0;'>")

    # ---------------- PRE-LOADED EXAMPLES ----------------
    gr.Markdown("### ⚡ Quick-Load Scenarios for Judges")
    gr.Examples(
        examples=[
            [
                "Who was behind the theory of relativity?",
                "[0] Albert Einstein developed the theory of relativity in the early 20th century.\n[1] The theory revolutionized physics and our understanding of space and time.\n[2] Isaac Newton originally proposed relativity in the 18th century."
            ],
            [
                "What is the boiling point of water at sea level?",
                "[0] Water boils at 100°C at standard atmospheric pressure.\n[1] This corresponds to 212°F.\n[2] Some sources claim water boils at 90°C at sea level."
            ]
        ],
        inputs=[question_input, ground_truth_input],
        label="Click a scenario to populate the fields:"
    )

    # Wire up the button
    run_btn.click(
        fn=run_custom_audit, 
        inputs=[question_input, ground_truth_input], 
        outputs=[json_output, status_indicator]
    )

# Launch the app (Spaces handles the networking)
if __name__ == "__main__":
    demo.launch()