import os
import streamlit as st
import pandas as pd
from langchain_openai import OpenAI
from langchain_core.prompts import PromptTemplate
from langchain_classic.chains import LLMChain
from langchain_google_genai import GoogleGenerativeAI
from langchain_core.prompts import PromptTemplate

# 1. Page Configuration & Styling
st.set_page_config(page_title="ClearPath Engine", layout="wide")

st.markdown("""
    <style>
    /* Unify fonts across the app */
    html, body, [class*="css"], .stMarkdown, .stMetric, .stAlert {
        font-family: 'Source Sans Pro', sans-serif !important;
    }
    .main { background-color: #f5f7f9; }
    
    /* Make metric values bold and clear */
    [data-testid="stMetricValue"] {
        font-size: 1.8rem !important;
        font-weight: 700 !important;
    }
    </style>
    """, unsafe_allow_html=True)

# 1.1. Access the secret securely
# os.getenv looks for the secret you just saved in HF Settings
api_key = os.getenv("Cost_Plus_Key")

# 2. Mock Data Engine (Simulating FHIR/Benefit Data)
# In a real scenario, this would be a RAG retrieval or API call
mock_meds = {
    "Zepbound": {"tier": "Specialty", "copay": 150, "true_cost": 112, "coupon": 25, "notes": "GLP-1 Agonist"},
    "Humira": {"tier": "Specialty", "copay": 250, "true_cost": 190, "coupon": 5, "notes": "Immunology"},
    "Atorvastatin": {"tier": "Generic", "copay": 10, "true_cost": 4, "coupon": None, "notes": "Cholesterol"},
    "Stelara": {"tier": "Specialty", "copay": 300, "true_cost": 210, "coupon": 10, "notes": "Biosimilar available"}
}

member_profile = {
    "name": "Member",
    "deductible_met": 450,
    "deductible_total": 3000,
    "plan_type": "Aetna Choice POS II"
}

# 3. Sidebar - App Info & API Key
# with st.sidebar:
#   st.title("⚙️ Admin Settings")
#    api_key = st.text_input("Enter OpenAI API Key", type="password")
#    st.info("This prototype demonstrates clinical-financial transparency for CVS Aetna.")

# 4. Header Section
st.title("🔴 Rx - Price My Meds")
st.subheader("Personalized Pharmacy Pricing Transparency")

# 5. User Input Section
col_search, col_info = st.columns([2, 1])
with col_search:
    selected_drug = st.selectbox("Select or Search for a Medication:", 
                                 options=list(mock_meds.keys()),
                                 index=None, 
                                 placeholder="Choose a medication to begin...")

with col_info:
    progress = member_profile['deductible_met'] / member_profile['deductible_total']
    st.write(f"**Deductible Progress:** ${member_profile['deductible_met']} / ${member_profile['deductible_total']}")
    st.progress(progress)

# 6. Pricing Logic & Display
if selected_drug:
    data = mock_meds[selected_drug]
    
    st.write("### Comparison Across Pricing Lanes")
    c1, c2, c3 = st.columns(3)
    
    with c1:
        st.metric(label="Standard Plan Copay", value=f"${data['copay']}", help="Based on your Aetna Benefit Design")
    with c2:
        st.metric(label="Your plan TrueCost", value=f"${data['true_cost']}", delta="-25% vs Plan", delta_color="normal")
    with c3:
        st.metric(label="Manufacturer Direct", value=f"${data['coupon'] if data['coupon'] else 'N/A'}", delta="Best Value" if data['coupon'] else None)

    # 7. Agentic Explanation (LangChain)
    st.write("---")
    st.write("### 🤖 Concierge Explanation")
    
    # 3. The "Click Here" Logic
    if st.button("Click here for AI Analysis"):
        if not api_key:
            st.error("API Key not found in Space Secrets. Please check Settings.")
        else:
            with st.spinner("Analyzing benefit design..."):
                try:
                    # Use the secret key automatically
                    llm = GoogleGenerativeAI(model="gemini-3.1-flash-lite", google_api_key=api_key)
                
                    template = """
                    You are a HealthPlan such as CVS or Unitedhealthcare Member Concierge. A member is looking at {drug}.
                    Plan: ${copay}, TrueCost: ${true_cost}, Coupon: ${coupon}.
                    Deductible: ${met}/${total}.
                    Explain the best financial path in 2 sentences.
                    CRITICAL INSTRUCTION: You MUST use the '$' sign before every single numerical amount (e.g., $5.00, $450.00). 
                    Do not provide numbers without the '$' prefix.
                    """
                    prompt = PromptTemplate.from_template(template)
                    chain = prompt | llm
                
                    explanation = chain.invoke({
                        "drug": selected_drug,
                        "copay": mock_meds[selected_drug]['copay'],
                        "true_cost": mock_meds[selected_drug]['true_cost'],
                        "coupon": mock_meds[selected_drug]['coupon'] or "N/A",
                        "met": 450, "total": 3000 # Example profile data
                    })
                    st.success(explanation)
                except Exception as e:
                    st.error(f"Analysis failed: {e}")

    # 8. Clinical Validation Note
    st.info(f"**Clinical Note:** This pricing reflects the {data['tier']} tier status. Always consult with your provider regarding therapeutic interchanges.")
else:
    # This shows when the app is "Empty"
    st.info("Please select a medication from the dropdown above to view pricing and AI insights.")