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

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

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

# Access the secret securely from HF Settings
api_key = os.getenv("Cost_Plus_Integrated")

# 2. Mock Data Engine (Enhanced with Acquisition & PBM Cost Data for Strategy View)
mock_meds = {
    "Zepbound": {
        "tier": "Specialty", "copay": 150, "true_cost": 112, "coupon": 25, "notes": "GLP-1 Agonist",
        "traditional_pbm_cost": 1050, "cost_plus_acq": 800, "markup_fee": 40, "dispensing_fee": 10
    },
    "Humira": {
        "tier": "Specialty", "copay": 250, "true_cost": 190, "coupon": 5, "notes": "Immunology",
        "traditional_pbm_cost": 6800, "cost_plus_acq": 5200, "markup_fee": 260, "dispensing_fee": 10
    },
    "Atorvastatin": {
        "tier": "Generic", "copay": 10, "true_cost": 4, "coupon": None, "notes": "Cholesterol",
        "traditional_pbm_cost": 15, "cost_plus_acq": 2, "markup_fee": 0.30, "dispensing_fee": 1.70
    },
    "Stelara": {
        "tier": "Specialty", "copay": 300, "true_cost": 210, "coupon": 10, "notes": "Biosimilar available",
        "traditional_pbm_cost": 12500, "cost_plus_acq": 9800, "markup_fee": 490, "dispensing_fee": 10
    }
}

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

# 3. Header Section
st.title("🔴 Rx - ClearPath Price Engine")
st.subheader("Commercial Strategy Portfolio & Member Transparency Prototype")

# Navigation Tabs to split Member Concierge and Executive Strategy
tab1, tab2 = st.tabs(["👤 Member Concierge View", "💼 Plan Sponsor (Employer) Strategy View"])

with tab1:
    # 4. 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...",
                                     key="member_drug_select")

    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)

    # 5. 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)

        # 6. Agentic Explanation (LangChain Core syntax)
        st.write("---")
        st.write("### 🤖 Concierge Explanation")
        
        if st.button("Click here for AI Analysis", key="member_ai_btn"):
            if not api_key:
                st.error("API Key not found in Space Secrets. Please check Settings.")
            else:
                with st.spinner("Analyzing benefit design..."):
                    try:
                        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. 
                        Do not provide numbers without the '$' prefix.
                        """
                        prompt = PromptTemplate.from_template(template)
                        chain = prompt | llm
                    
                        explanation = chain.invoke({
                            "drug": selected_drug,
                            "copay": data['copay'],
                            "true_cost": data['true_cost'],
                            "coupon": data['coupon'] or "N/A",
                            "met": 450, "total": 3000
                        })
                        st.success(explanation)
                    except Exception as e:
                        st.error(f"Analysis failed: {e}")

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

with tab2:
    st.write("### 🏢 Aetna Commercial Plan Sponsor Value Assessment")
    st.write("Demonstrating the macro value of moving a client portfolio from traditional PBM Spread/Rebate pricing models to transparent **CVS CostVantage** architectures.")
    
    # Strategy input controls
    col_strat1, col_strat2 = st.columns(2)
    with col_strat1:
        account_size = st.selectbox("Select Target Employer Group Size:", ["Mid-Market (500-5000 lives)", "National Accounts (5000+ lives)"])
        selected_drug_strat = st.selectbox("Select Medication for Financial Impact Analysis:", options=list(mock_meds.keys()), key="strat_drug_select")
    
    with col_strat2:
        employer_subsidy = st.slider("Employer Plan Share of Drug Cost (%)", min_value=50, max_value=100, value=80)
    
    if selected_drug_strat:
        s_data = mock_meds[selected_drug_strat]
        
        # Financial Calculations
        # CostVantage Formula = Acquisition + Markup + Dispensing
        total_cost_plus = s_data['cost_plus_acq'] + s_data['markup_fee'] + s_data['dispensing_fee']
        employer_traditional_spend = s_data['traditional_pbm_cost'] * (employer_subsidy / 100)
        employer_cost_plus_spend = total_cost_plus * (employer_subsidy / 100)
        net_savings_per_fill = employer_traditional_spend - employer_cost_plus_spend
        
        st.write("#### 📊 Transactional Financial Decomposition")
        
        sc1, sc2, sc3 = st.columns(3)
        with sc1:
            st.metric(label="Traditional PBM Billed Cost", value=f"${s_data['traditional_pbm_cost']:,}")
        with sc2:
            st.metric(label="CVS CostVantage True Cost", value=f"${total_cost_plus:,}", help="Acquisition + Fixed Markup + Flat Dispensing Fee")
        with sc3:
            st.metric(label="Net Employer Savings / Fill", value=f"${net_savings_per_fill:,.2f}", delta=f"{((s_data['traditional_pbm_cost'] - total_cost_plus)/s_data['traditional_pbm_cost'])*100:.1f}% Reduction")
        
        # Transparent breakdown expander
        with st.expander("🔍 View Transparent CostVantage Formula Components"):
            st.json({
                "Drug Acquisition Cost (AAC)": f"${s_data['cost_plus_acq']:,}",
                "CVS Defined Markup Fee": f"${s_data['markup_fee']:,}",
                "Flat Dispensing Fee": f"${s_data['dispensing_fee']:,}",
                "Total Formulary Cost Structure": f"${total_cost_plus:,}"
            })
            
        # Strategy Narrative Hook for Panel
        st.write("---")
        st.write("### 🎯 Executive Strategy Alignment Notes")
        st.markdown(f"""
        * **Formulary Differentiation:** By leveraging **CVS CostVantage** for *{selected_drug_strat}*, Aetna Commercial Sales can approach plan sponsors with guaranteed transparent pass-through pricing, stripping out historical opaque PBM spreads.
        * **Client Retention Playbook:** For a plan sponsor in the **{account_size}** segment, modeling this transparency directly mitigates consultant-driven RFP pressures by aligning interests across the plan design.
        """)