import streamlit as st
from datetime import date
import math
import random

from llm_utils import explain_savings_plan


# ---------------------- PAGE CONFIG ----------------------
st.set_page_config(
    page_title="Pistol Pete SmartAgent",
    page_icon="🏠",
    layout="centered",
)


# ---------------------- STATE INIT -----------------------
def init_state():
    if "step" not in st.session_state:
        st.session_state.step = 1

    # From your original JS userData object
    st.session_state.setdefault("home_type", "")
    st.session_state.setdefault("location", "")
    st.session_state.setdefault("timeline_bucket", 3)   # 1–5 slider
    st.session_state.setdefault("budget", 350000)

    st.session_state.setdefault("annual_income", "75000")
    st.session_state.setdefault("monthly_debt", "350")
    st.session_state.setdefault("current_savings", "12000")

    # Controls on Screen 5
    st.session_state.setdefault("timeline_years_plan", 4)  # 2–7 years
    st.session_state.setdefault("down_pct", 20)            # 10–20

    # LLM output
    st.session_state.setdefault("llm_explanation", None)


def to_int(value, default):
    try:
        if value is None:
            return default
        if isinstance(value, (int, float)):
            return int(value)
        v = str(value).replace(",", "").strip()
        if v == "":
            return default
        return int(v)
    except Exception:
        return default


init_state()


# ---------------------- GOVERNANCE SIDEBAR ----------------------
with st.sidebar:
    st.markdown("### 🔐 Governance & Fairness")
    st.caption(
        "- No protected attributes (race, gender, etc.) are used in calculations.\n"
        "- The LLM **only** explains numbers already computed.\n"
        "- Users can adjust timeline & budget and keep full control.\n"
        "- This is a prototype; no real accounts or PII are stored."
    )


# ---------------------- HELPERS ----------------------
def compute_savings_plan(budget, down_pct, current_savings, timeline_years):
    """
    Port of your JS updatePlanCalculations() logic:
    - Down payment = budget * %
    - Remaining need
    - Months = years * 12
    - Closing costs = 5000 / months
    - Interest earnings via 1.5% APY on (savings + contributions)
    - Recommended monthly = base + closing - interestEarnings
    """
    down_payment_amount = round(budget * (down_pct / 100))
    remaining_need = max(0, down_payment_amount - current_savings)
    months = max(1, timeline_years * 12)

    monthly_base = remaining_need / months
    closing_costs = 5000 / months

    # interestEarnings (same idea as your JS)
    annual_rate = 0.015
    balance = current_savings
    total_interest = 0.0
    for _ in range(timeline_years):
        balance += monthly_base * 12
        annual_interest = balance * annual_rate
        total_interest += annual_interest
        balance += annual_interest
    interest_earnings = total_interest / months

    monthly_savings = monthly_base + closing_costs - interest_earnings

    today = date.today()
    target_year = today.year + timeline_years
    # Same month/day as today
    target_date = date(target_year, today.month, today.day)

    return {
        "down_payment_amount": down_payment_amount,
        "remaining_need": round(remaining_need),
        "months": months,
        "monthly_base": monthly_base,
        "closing_costs": closing_costs,
        "interest_earnings": interest_earnings,
        "monthly_savings": monthly_savings,
        "target_date": target_date,
    }


def compute_estimated_home_value(budget, timeline_years, location, down_pct):
    """
    Port of estimateValue() from your JS:
    - Start with budget
    - timelineAdjPct = ((years - 4) * 0.02)
    - random adjustment ± 15k
    - Location tweak for Pittsburgh / Lawrenceville
    - Confidence from randAdj + timeline
    """
    base_budget = budget
    est = base_budget

    rand_adj = random.randint(-15000, 15000)
    timeline_adj_pct = (timeline_years - 4) * 0.02
    est = round(est * (1 + timeline_adj_pct) + rand_adj)

    loc = (location or "").lower()
    if "pittsburgh" in loc or "lawrenceville" in loc:
        est = round(est * 0.98)

    if abs(rand_adj) < 6000 and timeline_years <= 5:
        confidence = "High"
    elif abs(rand_adj) < 12000:
        confidence = "Moderate"
    else:
        confidence = "Low"

    return est, confidence


def next_step():
    st.session_state.step = min(7, st.session_state.step + 1)


def prev_step():
    st.session_state.step = max(1, st.session_state.step - 1)


# ---------------------- PROGRESS BAR ----------------------
st.markdown(
    "<h2 style='text-align:center;'>🏠 Pistol Pete SmartAgent</h2>",
    unsafe_allow_html=True,
)
st.caption("Prototype 2025 – Trusted, explainable first-home savings coach.")
st.progress(st.session_state.step / 7.0)
st.divider()

step = st.session_state.step


# ---------------------- SCREEN 1 ----------------------
if step == 1:
    st.header("Ready to Plan for Your Home?")
    st.write(
        "Our SmartAgent helps you create a personalized savings roadmap. "
        "You're in control every step of the way with transparent, explainable recommendations."
    )

    with st.expander("How it works", expanded=True):
        st.markdown(
            "- We ask about your **home goals** and **finances**.\n"
            "- We create a **down-payment savings plan**.\n"
            "- You can see the **exact math** and adjust your own timeline and budget.\n"
            "- An **LLM explainer** turns the numbers into plain-English explanations."
        )

    st.button("Start Your Personalized Plan ➜", on_click=next_step)


# ---------------------- SCREEN 2 ----------------------
elif step == 2:
    st.header("Let's Build Your Home Plan Together")
    st.write("Tell us about your dream home — this helps us create a plan that's right for you.")

    st.session_state.home_type = st.radio(
        "What type of home are you dreaming of?",
        ["Single-Family Home", "Townhouse", "Condominium", "I'm not sure yet"],
        index=3 if st.session_state.home_type == "" else
        ["Single-Family Home", "Townhouse", "Condominium", "I'm not sure yet"].index(
            st.session_state.home_type
        ),
    )

    st.session_state.location = st.text_input(
        "Where are you thinking of buying?",
        value=st.session_state.location,
        placeholder="e.g., Pittsburgh, PA – Lawrenceville",
    )

    col1, col2 = st.columns(2)
    with col1:
        st.button("⬅ Back", on_click=prev_step, key="back_2")
    with col2:
        st.button("Continue ➜", on_click=next_step, key="next_2")


# ---------------------- SCREEN 3 ----------------------
elif step == 3:
    st.header("Your Timeline & Budget")
    st.write("Now let's talk about when you'd like to buy and your comfortable budget range.")

    st.session_state.timeline_bucket = st.slider(
        "What's your ideal timeline to buy?",
        min_value=1,
        max_value=5,
        value=st.session_state.timeline_bucket,
    )

    def bucket_label(b):
        return {
            1: "< 1 year",
            2: "1–2 years",
            3: "3–5 years",
            4: "5–7 years",
            5: "7+ years",
        }[b]

    st.caption(f"Selected range: **{bucket_label(st.session_state.timeline_bucket)}**")

    # Raw budget input
    home_budget_raw = st.text_input(
        "What is your comfortable estimated home budget?",
        value=str(st.session_state.budget),
        placeholder="e.g., 350000",
    )
    st.session_state.budget = to_int(home_budget_raw, st.session_state.budget)

    if st.session_state.location and any(
        k in st.session_state.location.lower() for k in ["pittsburgh", "lawrenceville"]
    ):
        st.info("We see average prices in Lawrenceville are ~ $350,000 (illustrative).")

    st.checkbox(
        "Let our AI calculate a budget for me based on my finances (concept only)",
        value=False,
        help="For this prototype, the slider/budget is still user-controlled.",
    )

    col1, col2 = st.columns(2)
    with col1:
        st.button("⬅ Back", on_click=prev_step, key="back_3")
    with col2:
        st.button("Continue ➜", on_click=next_step, key="next_3")


# ---------------------- SCREEN 4 ----------------------
elif step == 4:
    st.header("Your Financial Snapshot")
    st.write(
        "To provide the most accurate plan, we’ll use the financial info you've shared. "
        "Verify or update any of the details below."
    )

    st.info("We value your privacy. This information is used only to create your plan and is not sold.")

    st.session_state.annual_income = st.text_input(
        "Annual Gross Income",
        value=st.session_state.annual_income,
        help="Example: sourced from bank deposits",
    )

    st.session_state.monthly_debt = st.text_input(
        "Monthly Debt Payments",
        value=st.session_state.monthly_debt,
        help="Example: credit cards, loans",
    )

    st.session_state.current_savings = st.text_input(
        "Current Down Payment Savings",
        value=st.session_state.current_savings,
        help="Example: savings accounts",
    )

    col1, col2 = st.columns(2)
    with col1:
        st.button("⬅ Back", on_click=prev_step, key="back_4")
    with col2:
        st.button("Generate My Home Plan ➜", on_click=next_step, key="next_4")


# ---------------------- SCREEN 5 ----------------------
elif step == 5:
    st.header("Your Personalized Home Plan")

    location_label = st.session_state.location or "your target area"
    st.write(
        f"Based on your goals for a home in **{location_label}**, "
        "here's your tailored savings roadmap."
    )

    # Collect numeric values
    budget = st.session_state.budget
    income = to_int(st.session_state.annual_income, 75000)
    debt = to_int(st.session_state.monthly_debt, 350)
    savings = to_int(st.session_state.current_savings, 12000)

    # -------- Adjust Your Plan (Control Panel) --------
    st.subheader("Adjust Your Plan")

    colA, colB = st.columns(2)
    with colA:
        st.session_state.timeline_years_plan = st.slider(
            "Move Your Timeline (years)",
            min_value=2,
            max_value=7,
            value=st.session_state.timeline_years_plan,
        )
    with colB:
        st.session_state.budget = st.slider(
            "Adjust Home Budget ($)",
            min_value=250000,
            max_value=600000,
            step=50000,
            value=st.session_state.budget,
        )

    st.session_state.down_pct = st.slider(
        "Change Down Payment %",
        min_value=10,
        max_value=20,
        step=5,
        value=st.session_state.down_pct,
    )

    if st.session_state.down_pct < 20:
        st.warning("A down payment below 20% may require Private Mortgage Insurance (PMI).")

    # Recompute plan with updated values
    budget = st.session_state.budget
    plan = compute_savings_plan(
        budget=budget,
        down_pct=st.session_state.down_pct,
        current_savings=savings,
        timeline_years=st.session_state.timeline_years_plan,
    )

    # -------- Savings Timeline Card --------
    st.subheader("Savings Timeline")
    col_today, _, col_goal = st.columns([1, 2, 1])
    with col_today:
        st.caption("Today")
        st.markdown(f"**${savings:,.0f}**")
    with col_goal:
        st.caption("Goal")
        st.markdown(f"**${plan['down_payment_amount']:,.0f}**")

    st.caption(
        f"Target Date: **{plan['target_date'].strftime('%B %Y')}** "
        f"({st.session_state.timeline_years_plan} years)"
    )

    st.divider()

    # -------- AI Recommendation Card --------
    st.subheader("AI Recommendation")
    st.write(
        f"To reach your goal in **{st.session_state.timeline_years_plan} years**, "
        "we recommend saving:"
    )
    st.markdown(
        f"<div style='font-size:32px;font-weight:800;color:#D96932;text-align:center;'>"
        f"${plan['monthly_savings']:,.0f}/month</div>",
        unsafe_allow_html=True,
    )
    st.caption("Confidence: **High** (demo).")

    st.write(
        "💡 **Why this level?** Your income and debts produce a debt-to-income profile that "
        "supports this plan, plus predictable regional assumptions for this demo."
    )

    # -------- LLM Explanation Layer --------
    st.markdown("#### Explanation (AI Coach)")

    payload = {
        "home_budget": budget,
        "down_payment_percent": st.session_state.down_pct,
        "down_payment_amount": plan["down_payment_amount"],
        "current_savings": savings,
        "remaining_need": plan["remaining_need"],
        "timeline_years": st.session_state.timeline_years_plan,
        "timeline_months": plan["months"],
        "recommended_monthly_savings": round(plan["monthly_savings"]),
        "annual_income": income,
        "monthly_debt": debt,
    }

    if st.button("Explain this plan with AI"):
        with st.spinner("Generating a plain-English explanation..."):
            explanation = explain_savings_plan(payload)
            st.session_state.llm_explanation = explanation

    if st.session_state.llm_explanation:
        st.info(st.session_state.llm_explanation)
    else:
        st.caption("Click the button above to see the AI Coach explain your plan in simple terms.")

    # -------- See the AI's Math (Accordion-like) --------
    with st.expander("See the AI's Math"):
        st.write(
            f"**Target Down Payment ({st.session_state.down_pct}%):** "
            f"${plan['down_payment_amount']:,.0f}"
        )
        st.write(f"**− Current Savings:** ${savings:,.0f}")
        st.write(f"**= Remaining Need:** ${plan['remaining_need']:,.0f}")
        st.write(
            f"**÷ Months in Timeline ({plan['months']}):** "
            f"${plan['monthly_base']:,.0f} base / month"
        )
        st.write(
            f"**+ Estimated Closing Costs (≈ $5,000 / {plan['months']}):** "
            f"+ ${plan['closing_costs']:,.0f}"
        )
        st.write(
            f"**− Projected Interest on Savings (@1.5% APY):** "
            f"− ${plan['interest_earnings']:,.0f}"
        )
        st.write(
            f"**= Recommended Monthly Savings:** "
            f"${plan['monthly_savings']:,.0f} / month"
        )

    st.divider()

    col1, col2 = st.columns(2)
    with col1:
        st.button("⬅ Back", on_click=prev_step, key="back_5")
    with col2:
        st.button("Next Steps ➜", on_click=next_step, key="next_5")


# ---------------------- SCREEN 6 ----------------------
elif step == 6:
    st.header("Tools to Help You Succeed")
    st.write("Based on your plan, here are products and support options to accelerate progress.")

    savings = to_int(st.session_state.current_savings, 12000)

    # Product 1
    st.subheader("PRODUCT – High-Yield “Home Fund” Savings Account (Concept)")
    st.write(
        f"Based on your current savings of **${savings:,.0f}**, you could earn "
        f"**more interest** per year in a high-yield account compared to a basic account."
    )
    st.caption("Why we suggest this: to maximize your savings growth with minimal risk.")

    st.divider()

    # Product 2
    st.subheader("SERVICE – Consider a Home Loan Pre-approval (Concept)")
    st.write(
        "A pre-approval helps clarify your budget. Many buyers who succeed "
        "get pre-approved early in the process."
    )
    st.caption("Why we suggest this: to reduce uncertainty and strengthen your offers.")

    st.divider()

    # Human contact form (replaces modal)
    st.subheader("Prefer to Speak with a Human Agent?")
    with st.form("contact_form"):
        name = st.text_input("Full name")
        email = st.text_input("Email address")
        phone = st.text_input("Phone (optional)")
        msg = st.text_area("How can we assist you?")
        submitted = st.form_submit_button("Send Message")
        if submitted:
            st.success(
                f"Thanks {name or 'there'}! A human agent would contact you at "
                f"{email or phone or 'your provided contact'} in a real deployment."
            )

    st.divider()

    col1, col2 = st.columns(2)
    with col1:
        st.button("⬅ Back to Plan", on_click=prev_step, key="back_6")
    with col2:
        st.button("Save & Estimate Value ➜", on_click=next_step, key="next_6")


# ---------------------- SCREEN 7 ----------------------
elif step == 7:
    st.header("Your Estimated Home Value")

    budget = st.session_state.budget
    down_pct = st.session_state.down_pct
    timeline_years = st.session_state.timeline_years_plan
    location = st.session_state.location

    est_value, confidence = compute_estimated_home_value(
        budget=budget,
        timeline_years=timeline_years,
        location=location,
        down_pct=down_pct,
    )

    st.markdown(
        f"<div style='background:#FFF8EC;border-radius:12px;padding:18px;"
        f"box-shadow:0 6px 18px rgba(217,105,50,0.08);text-align:center;'>"
        f"<h2 style='font-size:34px;color:#D96932;'>${est_value:,.0f}</h2>"
        f"<p>Estimated probable home value given your inputs (illustrative only).</p>"
        f"<p>Confidence: <strong>{confidence}</strong></p>"
        f"</div>",
        unsafe_allow_html=True,
    )

    st.subheader("How we estimate")
    st.write(
        "We combine your desired budget, regional assumptions (if you supplied a location), "
        "and timeline to produce a probable value. In production, MLS or public-record APIs "
        "would provide more accurate estimates."
    )

    colA, colB = st.columns(2)
    with colA:
        st.markdown("**Input Budget**")
        st.markdown(f"${budget:,.0f}")
    with colB:
        st.markdown("**Down Payment (approx)**")
        st.markdown(f"${round(budget * down_pct / 100):,.0f}")

    st.divider()

    col1, col2 = st.columns(2)
    with col1:
        st.button("⬅ Back", on_click=prev_step, key="back_7")
    with col2:
        if st.button("Finish & Save"):
            st.success("Your plan summary has been saved (demo only). Refresh to start over.")