import json
import time
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import streamlit as st

# ==========================================
# PAGE CONFIGURATION
# ==========================================
st.set_page_config(
    page_title="SmartClaim AI Platform",
    page_icon="🏥",
    layout="wide",
    initial_sidebar_state="expanded",
)

# Custom CSS for Enterprise Styling
st.markdown(
    """
<style>
    .main-header { font-size: 26px; font-weight: 700; color: #CC0000; margin-bottom: 0px; }
    .sub-header { font-size: 14px; color: #555555; margin-bottom: 20px; }
    .metric-card { background-color: #F8F9FA; border-left: 4px solid #CC0000; padding: 12px; border-radius: 4px; }
    .stButton>button { width: 100%; border-radius: 4px; height: 42px; font-weight: 600; }
    .badge-pass { background-color: #D4EDDA; color: #155724; padding: 4px 8px; border-radius: 4px; font-weight: 600; }
    .badge-review { background-color: #FFF3CD; color: #856404; padding: 4px 8px; border-radius: 4px; font-weight: 600; }
    .badge-deny { background-color: #F8D7DA; color: #721C24; padding: 4px 8px; border-radius: 4px; font-weight: 600; }
</style>
""",
    unsafe_allow_html=True,
)

# ==========================================
# SESSION STATE INITIALIZATION
# ==========================================
if "processed_claims" not in st.session_state:
    st.session_state.processed_claims = []

if "evaluated" not in st.session_state:
    st.session_state.evaluated = False

if "show_letter" not in st.session_state:
    st.session_state.show_letter = False

if "action_status" not in st.session_state:
    st.session_state.action_status = None

if "eval_results" not in st.session_state:
    st.session_state.eval_results = {}

# ==========================================
# KNOWLEDGE BASE: AETNA CPBs & SAMPLE CLAIMS
# ==========================================
AETNA_CPBS = {
    "CPB 0236 (Spine MRI)": {
        "title": "Magnetic Resonance Imaging (MRI) of the Spine",
        "code": "CPT 72148",
        "criteria": [
            "Clinical evidence of spinal stenosis or cauda equina compression",
            "Progressively severe symptoms despite conservative management",
            "Persistent back/neck pain with radiculopathy with failed 6+ weeks of conservative therapy (NSAIDs, physical therapy)",
            "Suspected spinal infection, fracture, or malignancy",
        ],
    },
    "CPB 0171 (Knee MRI)": {
        "title": "Magnetic Resonance Imaging (MRI) of the Extremities",
        "code": "CPT 73721",
        "criteria": [
            "Persistent knee swelling/instability not associated with injury, failed 3+ weeks conservative therapy",
            "True joint locking indicative of torn meniscus or loose body",
            "Suspected osteomyelitis or bone infection",
            "Multi-view X-rays performed to rule out fracture prior to advanced imaging",
        ],
    },
    "CPB 0157 (Bariatric Surgery)": {
        "title": "Obesity Surgery",
        "code": "CPT 43644",
        "criteria": [
            "Body Mass Index (BMI) >= 40 or BMI >= 35 with high-risk comorbidities (Type 2 Diabetes, Severe Sleep Apnea)",
            "Documented participation in medically supervised weight loss program for >= 6 months",
            "Psychosocial behavioral health evaluation and clearance completed",
            "Absence of active substance use disorder or uncontrolled psychiatric illness",
        ],
    },
}

MOCK_CLAIMS = {
    "CLM-90821 (Lumbar MRI - Compliant)": {
        "claim_id": "CLM-90821",
        "patient": "Jane Doe (DOB: 11/04/1978)",
        "cpb": "CPB 0236 (Spine MRI)",
        "cpt": "72148 - MRI Lumbar Spine w/o Contrast",
        "diagnosis": "M54.16 - Radiculopathy, lumbar region",
        "clinical_notes": "46yo female with severe low back pain radiating to L5 distribution. Completed 8 weeks of physical therapy and trials of Meloxicam without relief. Straight leg raise test positive on left. X-ray showed mild disc space narrowing at L4-L5, no acute fracture.",
        "expected_status": "APPROVED",
        "confidence": 0.97,
        "citation": "Meets CPB 0236: Persistent back pain with radiculopathy + documented >6 weeks conservative therapy (8 wks PT + NSAIDs) + prior X-ray.",
    },
    "CLM-44319 (Spine MRI - Insufficient PT)": {
        "claim_id": "CLM-44319",
        "patient": "Robert Smith (DOB: 03/15/1985)",
        "cpb": "CPB 0236 (Spine MRI)",
        "cpt": "72148 - MRI Lumbar Spine w/o Contrast",
        "diagnosis": "M54.50 - Low back pain, unspecified",
        "clinical_notes": "41yo male presenting with acute low back pain following lifting heavy box 10 days ago. No numbness, tingling, or bowel/bladder dysfunction. Patient requests MRI today. Took OTC Ibuprofen twice with minor improvement.",
        "expected_status": "MANUAL_REVIEW",
        "confidence": 0.72,
        "citation": "Fails CPB 0236: Only 10 days of symptoms. Required 6 weeks of conservative therapy (PT/NSAIDs) not completed. Red flag symptoms absent.",
    },
    "CLM-11920 (Bariatric Surgery - Missing Psych)": {
        "claim_id": "CLM-11920",
        "patient": "Maria Garcia (DOB: 08/22/1981)",
        "cpb": "CPB 0157 (Bariatric Surgery)",
        "cpt": "43644 - Laparoscopic Roux-en-Y Gastric Bypass",
        "diagnosis": "E66.01 - Morbid severe obesity",
        "clinical_notes": "44yo female, BMI 42.1 with poorly controlled Type 2 Diabetes (HbA1c 8.4%). Documented 6-month physician-monitored diet program completed in June 2026. Behavioral health evaluation pending scheduled visit next month.",
        "expected_status": "PEND_DOCS",
        "confidence": 0.81,
        "citation": "Meets BMI criteria (>40) and 6-mo weight program, but missing mandatory Behavioral Health / Psych Clearance per CPB 0157.",
    },
}

# ==========================================
# SIDEBAR CONTROL PANEL
# ==========================================
with st.sidebar:
    st.image(
        "https://upload.wikimedia.org/wikipedia/commons/f/f3/Health_logo.svg",
        width=180,
    )
    st.markdown("### **AI Claims Platform Engine**")
    st.caption("Agentic Decision-Support Framework")
    st.divider()

    st.subheader("⚙️ Platform Governance")
    auto_adj_threshold = st.slider(
        "Auto-Adjudication Confidence Threshold",
        min_value=0.70,
        max_value=0.99,
        value=0.90,
        step=0.01,
        help="Claims with model confidence above this threshold bypass manual review.",
    )

    st.subheader("🔑 Live LLM Config (Optional)")
    api_key = st.text_input(
        "OpenAI API Key",
        type="password",
        placeholder="sk-...",
        help="Leave blank to use internal deterministic Agent Engine",
    )

    st.divider()
    st.info(
        "**Prototype**\n\nFocus Area: Agentic Workflow Automation, Policy RAG, & Human-in-the-Loop Safeguards."
    )

# ==========================================
# HEADER SECTION
# ==========================================
st.markdown(
    "<div class='main-header'> SmartClaim AI | Enterprise Claims Decision-Support</div>",
    unsafe_allow_html=True,
)
st.markdown(
    "<div class='sub-header'>Powered by Policy Retrieval-Augmented Generation (RAG) & Agentic Reasoning Engines</div>",
    unsafe_allow_html=True,
)

tab1, tab2, tab3 = st.tabs(
    [
        "📋 Agentic Claims Copilot (Ops)",
        "⚡ Auto-Adjudication & Guardrails",
        "📊 Executive Observability (AVP)",
    ]
)

# ==========================================
# TAB 1: AGENTIC CLAIMS COPILOT
# ==========================================
with tab1:
    st.markdown("### 1. Select or Input Claim Information")

    col_input1, col_input2 = st.columns([1, 1])

    with col_input1:
        selected_sample = st.selectbox(
            "Select Pre-loaded Test Scenario:", list(MOCK_CLAIMS.keys())
        )
        claim_data = MOCK_CLAIMS[selected_sample]

        claim_id = st.text_input("Claim ID", claim_data["claim_id"])
        patient_info = st.text_input("Patient Info", claim_data["patient"])
        cpb_selection = st.selectbox(
            "Target Policy Bulletin (CPB)",
            list(AETNA_CPBS.keys()),
            index=list(AETNA_CPBS.keys()).index(claim_data["cpb"]),
        )

    with col_input2:
        cpt_code = st.text_input("Procedure / CPT Code", claim_data["cpt"])
        diag_code = st.text_input("Diagnosis Code", claim_data["diagnosis"])
        clinical_notes = st.text_area(
            "Clinical Progress Notes & Unstructured EHR Data:",
            claim_data["clinical_notes"],
            height=120,
        )

    run_btn = st.button("🚀 Run Agentic Claim Evaluation", type="primary")

    if run_btn:
        with st.spinner("Agent retrieving Aetna CPB guidelines and executing semantic policy matching..."):
            time.sleep(1.0)

            # Store in Session State
            st.session_state.evaluated = True
            st.session_state.show_letter = False
            st.session_state.action_status = None
            st.session_state.eval_results = {
                "claim_id": claim_id,
                "patient": patient_info,
                "cpb": cpb_selection,
                "cpt": cpt_code,
                "confidence": claim_data["confidence"],
                "status": claim_data["expected_status"],
                "citation": claim_data["citation"],
            }

            # Add to persistent audit log
            st.session_state.processed_claims.append(
                {
                    "timestamp": time.strftime("%H:%M:%S"),
                    "claim_id": claim_id,
                    "cpb": cpb_selection,
                    "status": claim_data["expected_status"],
                    "confidence": claim_data["confidence"],
                    "route": "Straight-Through (Auto)"
                    if claim_data["confidence"] >= auto_adj_threshold
                    else "HITL Specialist Queue",
                }
            )

    # Render results if evaluation has occurred
    if st.session_state.evaluated and st.session_state.eval_results:
        res = st.session_state.eval_results
        confidence = res["confidence"]
        status = res["status"]

        st.divider()
        st.markdown("### 2. Agentic Reasoning & Policy Matching Results")

        res_col1, res_col2, res_col3 = st.columns([1, 1, 1])

        with res_col1:
            st.markdown("**System Recommendation:**")
            if status == "APPROVED" and confidence >= auto_adj_threshold:
                st.markdown(
                    "<span class='badge-pass'>AUTO-APPROVE (STP)</span>",
                    unsafe_allow_html=True,
                )
            elif status == "PEND_DOCS":
                st.markdown(
                    "<span class='badge-review'>PEND FOR ADDITIONAL DOCUMENTS</span>",
                    unsafe_allow_html=True,
                )
            else:
                st.markdown(
                    "<span class='badge-review'>REFER TO HUMAN SPECIALIST</span>",
                    unsafe_allow_html=True,
                )

        with res_col2:
            st.markdown("**Model Confidence Score:**")
            st.metric(
                label="Confidence",
                value=f"{int(confidence*100)}%",
                delta=f"{'+' if confidence >= auto_adj_threshold else '-'}{abs(round((confidence - auto_adj_threshold)*100, 1))}% vs Threshold",
            )

        with res_col3:
            st.markdown("**Workflow Routing:**")
            if confidence >= auto_adj_threshold:
                st.success("✅ Straight-Through Processing (Zero Human Touch)")
            else:
                st.warning("⚠️ Routed to Human Specialist Queue (Below Risk Threshold)")

        st.markdown("#### 📜 Policy Line-Item Evidence Citations")
        st.info(f"**CPB Citation Analysis:** {res['citation']}")

        # Human-in-the-Loop (HITL) Action Panel
        st.markdown("#### 🛠️ Specialist Human-in-the-Loop Actions")
        action_col1, action_col2, action_col3 = st.columns(3)

        with action_col1:
            if st.button("✅ Confirm & Approve Claim"):
                st.session_state.action_status = f"✅ Claim {res['claim_id']} approved by Specialist. Adjudication logged."
                st.session_state.show_letter = False

        with action_col2:
            if st.button("✉️ Draft Pre-filled Request Letter"):
                st.session_state.show_letter = True
                st.session_state.action_status = None

        with action_col3:
            if st.button("❌ Issue Prior Auth Denial Notice"):
                st.session_state.action_status = f"❌ Denial Notice initiated for {res['claim_id']} per {res['cpb']} non-compliance."
                st.session_state.show_letter = False

        # Display persistent action statuses or pre-filled letter
        if st.session_state.action_status:
            st.info(st.session_state.action_status)

        if st.session_state.show_letter:
            st.markdown("##### ✉️ Pre-filled Provider Outreach Letter")
            letter_text = f"Dear Provider,\n\nRegarding Claim {res['claim_id']} for {res['cpt']}, our automated review against Aetna {res['cpb']} indicates missing required documentation:\n- {res['citation']}\n\nPlease submit clinical records within 14 days.\n\nSincerely,\nAetna Clinical Operations"
            st.text_area("Generated Outreach Draft:", value=letter_text, height=150)

# ==========================================
# TAB 2: AUTO-ADJUDICATION & GUARDRAILS
# ==========================================
with tab2:
    st.markdown("### Dynamic Risk & Threshold Impact Simulator")
    st.caption("Evaluate how adjusting confidence guardrails impacts operational throughput vs audit risk.")

    sim_col1, sim_col2 = st.columns([1, 2])

    with sim_col1:
        st.markdown("#### Simulation Control")
        total_daily_volume = st.number_input(
            "Daily Claim Volume:", value=25000, step=1000
        )
        avg_cost_per_manual = st.number_input(
            "Cost per Manual Claim Review ($):", value=14.50, step=0.50
        )

        stp_rate = max(0.20, min(0.85, 1.25 - (auto_adj_threshold * 0.8)))
        auto_volume = int(total_daily_volume * stp_rate)
        manual_volume = total_daily_volume - auto_volume
        daily_savings = auto_volume * avg_cost_per_manual

        st.metric(
            "Simulated Auto-Adjudication (STP) Rate", f"{round(stp_rate*100, 1)}%"
        )
        st.metric(
            "Projected Annual Operating Savings",
            f"${daily_savings * 260:,.0f}",
        )

    with sim_col2:
        st.markdown("#### Daily Volume Distribution Forecast")
        fig_pie = px.pie(
            values=[auto_volume, manual_volume],
            names=["Auto-Adjudicated (AI)", "Manual Specialist Review (HITL)"],
            color_discrete_sequence=["#28A745", "#FFC107"],
            hole=0.4,
        )
        fig_pie.update_layout(margin=dict(t=20, b=20, l=20, r=20), height=300)
        st.plotly_chart(fig_pie, use_container_width=True)

    st.divider()
    st.markdown("### Platform Guardrails & Safety Controls")

    g_col1, g_col2, g_col3 = st.columns(3)
    with g_col1:
        st.markdown("#### 🔒 Anti-Hallucination")
        st.caption("Pydantic strict schema parsing enforces structured JSON output with mandated CPB paragraph citations.")
    with g_col2:
        st.markdown("#### ⚖️ Compliance & HIPAA")
        st.caption("De-identification pipelines strip PHI before prompt embedding; audit logging records all LLM inference seeds.")
    with g_col3:
        st.markdown("#### 🔄 Model Drift Monitoring")
        st.caption("Continuous monitoring flags discrepancies between specialist override patterns and agent recommendations.")

# ==========================================
# TAB 3: EXECUTIVE OBSERVABILITY (AVP VIEW)
# ==========================================
with tab3:
    st.markdown("### Executive Dashboard | Analytics & Behavior Change (A&BC)")

    kpi1, kpi2, kpi3, kpi4 = st.columns(4)
    kpi1.metric("Current STP Rate", "68.4%", "+14.2% YoY")
    kpi2.metric("Average Handle Time (AHT)", "2.3 min", "-11.8 min")
    kpi3.metric("First-Pass Accuracy", "99.1%", "+1.8%")
    kpi4.metric("Specialist Override Rate", "3.2%", "-0.8%")

    st.divider()

    st.markdown("### Platform Performance Trends")
    chart_col1, chart_col2 = st.columns(2)

    with chart_col1:
        st.markdown("#### Weekly Processing Volume vs Manual Touch Points")
        weeks = [f"Week {i}" for i in range(1, 9)]
        df_trends = pd.DataFrame(
            {
                "Week": weeks,
                "Auto-Adjudicated": [12000, 13500, 14200, 15800, 16500, 17200, 18100, 19000],
                "Specialist Review": [8000, 7200, 6800, 5900, 5200, 4800, 4200, 3800],
            }
        )
        fig_bar = px.bar(
            df_trends,
            x="Week",
            y=["Auto-Adjudicated", "Specialist Review"],
            color_discrete_map={"Auto-Adjudicated": "#002B49", "Specialist Review": "#CC0000"},
            barmode="stack",
        )
        fig_bar.update_layout(height=320, margin=dict(t=20, b=20, l=20, r=20))
        st.plotly_chart(fig_bar, use_container_width=True)

    with chart_col2:
        st.markdown("#### Agent vs Specialist Agreement Rate (By Clinical Category)")
        df_agree = pd.DataFrame(
            {
                "Category": ["Radiology (MRI/CT)", "Bariatric / Surgery", "Oncology", "Orthopedics", "Cardiology"],
                "Agreement Rate (%)": [98.2, 94.5, 99.1, 96.4, 97.8],
            }
        )
        fig_gauge = px.bar(
            df_agree,
            x="Agreement Rate (%)",
            y="Category",
            orientation="h",
            color="Agreement Rate (%)",
            color_continuous_scale="Reds",
        )
        fig_gauge.update_layout(height=320, margin=dict(t=20, b=20, l=20, r=20))
        st.plotly_chart(fig_gauge, use_container_width=True)

    st.markdown("### Real-time Session Audit Log")
    if st.session_state.processed_claims:
        st.dataframe(pd.DataFrame(st.session_state.processed_claims), use_container_width=True)
    else:
        st.info("No claims processed in current session. Run a claim evaluation in Tab 1 to see real-time audit logging.")