import streamlit as st
import pandas as pd
import altair as alt
from datetime import datetime
from decimal import Decimal
import io

# PDF export via ReportLab
from reportlab.lib.pagesizes import letter
from reportlab.pdfgen import canvas

# ─── Helpers ───────────────────────────────────────────────────────────────────
def normalize_percentages(raw_dict):
    cats, vals = list(raw_dict.keys()), list(raw_dict.values())
    total = sum(vals)
    normalized = {}
    if total <= 0:
        each = round(100 / len(cats), 2)
        for c in cats:
            normalized[c] = each
        diff = 100 - sum(normalized.values())
        normalized[cats[-1]] += diff
    else:
        cum = 0.0
        for i, c in enumerate(cats):
            if i < len(cats) - 1:
                p = round((raw_dict[c] / total) * 100, 2)
                normalized[c] = p
                cum += p
            else:
                normalized[c] = round(100 - cum, 2)
    return normalized

def max_to_str(x):
    if isinstance(x, (int, float, Decimal)):
        return f"${float(x):,.2f}"
    return x  # e.g. "No Max"

def df_to_pdf(df: pd.DataFrame) -> io.BytesIO:
    buffer = io.BytesIO()
    c = canvas.Canvas(buffer, pagesize=letter)
    width, height = letter
    x_offset, y_offset = 40, height - 40
    line_height = 14

    # Header
    for i, col in enumerate(df.columns):
        c.drawString(x_offset + i*100, y_offset, str(col))
    y_offset -= line_height

    # Rows
    for _, row in df.iterrows():
        for i, cell in enumerate(row):
            c.drawString(x_offset + i*100, y_offset, str(cell))
        y_offset -= line_height
        if y_offset < 40:
            c.showPage()
            y_offset = height - 40

    c.save()
    buffer.seek(0)
    return buffer

# ─── Page setup ────────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Priority Budget Allocator",
    page_icon="💸",
    layout="wide"
)
st.title("💸 Priority Budget Allocator")
st.subheader("We budget for you!")
st.markdown(
    "Enter your balances and categories below, then **Generate Budget** to see your allocation. "
    "When you’re happy, download your results as **CSV**, **Excel** or **PDF**."
)

# ─── STEP 1: Balances ──────────────────────────────────────────────────────────
with st.expander("Step 1: Account Balances", expanded=True):
    num_accounts = st.number_input("How many accounts?", min_value=1, max_value=10, step=1, value=1)
    account_balances = [
        st.number_input(
            f"Account {i+1} balance ($)",
            min_value=0.0, format="%.2f", key=f"acct_{i}"
        )
        for i in range(num_accounts)
    ]
    total_balance = sum(account_balances)
    st.success(f"Total Available Balance: **${total_balance:,.2f}**")

# ─── STEP 2: Categories ─────────────────────────────────────────────────────────
invalid_max_min = False
with st.expander("Step 2: Define Spending Categories", expanded=True):
    num_categories = st.number_input("How many categories?", min_value=1, max_value=15, step=1, value=1)
    categories = []
    for i in range(num_categories):
        st.subheader(f"Category {i+1}")
        name = st.text_input("Name", key=f"name_{i}").strip() or f"Category {i+1}"
        minimum = st.number_input("Minimum ($)", min_value=0.0, format="%.2f", key=f"min_{i}")
        has_max = st.checkbox("Has a maximum?", key=f"has_max_{i}")
        max_amt = None
        if has_max:
            max_amt = st.number_input("Maximum ($)", min_value=0.0, format="%.2f", key=f"max_{i}")
            if max_amt < minimum:
                st.warning("⚠️ Maximum < Minimum—please adjust.")
                invalid_max_min = True
        categories.append({
            "Category": name,
            "Min": minimum,
            "Has_Max": has_max,
            "Max": max_amt
        })

# ─── Optional: Mandatory Savings ────────────────────────────────────────────────
with st.expander("Optional: Mandatory Savings", expanded=False):
    include_savings = st.checkbox("Include a mandatory 'Savings' category?")
    if include_savings:
        savings_pct = st.number_input(
            "Savings (% of total balance)", min_value=0.0, max_value=100.0,
            value=10.0, format="%.2f", key="savings_pct"
        )
        st.info("This will reserve that % before other allocations.")

# ─── STEP 2b: Surplus % for “No Max” ────────────────────────────────────────────
no_max = [c["Category"] for c in categories if not c["Has_Max"]]
raw = {}
if no_max:
    with st.expander("Step 2b: % Distribution for ‘No Max’ Categories"):
        st.markdown("They’ll be normalized to sum to 100%.")
        for c in no_max:
            raw[c] = st.number_input(
                f"% for {c}", min_value=0.0, max_value=100.0,
                value=round(100/len(no_max), 2), key=f"raw_{c}"
            )

# ─── STEP 3: Generate & Display ─────────────────────────────────────────────────
if st.button("📊 Generate Budget"):
    if invalid_max_min:
        st.error("❌ Please fix category errors (Max ≥ Min) before generating budget.")
    else:
        # Insert mandatory savings category if requested
        if include_savings:
            savings_amt = round(total_balance * savings_pct / 100, 2)
            categories.insert(0, {
                "Category": "Savings",
                "Min": savings_amt,
                "Has_Max": False,
                "Max": None
            })

        df = pd.DataFrame(categories)
        df["Allocation"] = df["Min"].copy()
        sum_min = df["Min"].sum()

        if total_balance < sum_min:
            st.warning("⚠️ Balance < sum of minimums—allocating proportionally to Min.")
            df["Allocation"] = (df["Min"] / sum_min * total_balance).round(2)
        else:
            remaining = total_balance - sum_min
            # fill to Max
            for idx, row in df.iterrows():
                if row["Has_Max"]:
                    cap = row["Max"] - row["Min"]
                    add = min(cap, remaining)
                    df.at[idx, "Allocation"] += round(add, 2)
                    remaining -= add
            # distribute leftover
            if remaining > 0 and no_max:
                norm = normalize_percentages(raw)
                st.subheader("🔢 Normalized % Distribution")
                dist_df = (
                    pd.DataFrame.from_dict(norm, orient="index", columns=["Pct"])
                    .rename_axis("Category").reset_index()
                )
                st.dataframe(dist_df, use_container_width=True)
                for idx, row in df.iterrows():
                    if not row["Has_Max"]:
                        df.at[idx, "Allocation"] += round(remaining * norm[row["Category"]] / 100, 2)
                remaining = 0

        df["Surplus/Deficit"] = (df["Allocation"] - df["Min"]).round(2)

        # ── Metrics ──────────────────────────────────────────────────────────────
        alloc_sum = df["Allocation"].sum()
        unalloc = total_balance - alloc_sum
        st.header("Key Metrics")
        c1, c2, c3 = st.columns(3)
        c1.metric("Total Balance", f"${total_balance:,.2f}")
        c2.metric("Total Allocated", f"${alloc_sum:,.2f}")
        c3.metric(
            "Unallocated",
            f"${unalloc:,.2f}" if unalloc >= 0 else f"-${abs(unalloc):,.2f}"
        )

        # ── Explain Terms ─────────────────────────────────────────────────────────
        with st.expander("❓ What do Allocation & Surplus/Deficit mean?", expanded=False):
            st.markdown(
                "- **Allocation**: The dollar amount assigned to each category based on your inputs.\n"
                "- **Surplus/Deficit**: Allocation minus your Minimum. A positive number means you have extra above your minimum; negative means you fell short."
            )

        # ── Display Table ────────────────────────────────────────────────────────
        disp = df.copy()
        disp["Max"] = disp["Max"].fillna("No Max").apply(max_to_str)
        disp["Min"] = disp["Min"].apply(lambda x: f"${x:,.2f}")
        disp["Allocation"] = disp["Allocation"].apply(lambda x: f"${x:,.2f}")
        disp["Surplus/Deficit"] = disp["Surplus/Deficit"].apply(lambda x: f"${x:,.2f}")
        st.header("📋 Allocation Breakdown")
        st.dataframe(disp, use_container_width=True)

        # ── Charts ───────────────────────────────────────────────────────────────
        st.header("📊 Allocation by Category")
        chart1 = (
            alt.Chart(df)
            .mark_bar()
            .encode(
                x=alt.X("Category:N", sort=None),
                y=alt.Y("Allocation:Q", title="Allocated ($)"),
                tooltip=[
                    alt.Tooltip("Category:N"),
                    alt.Tooltip("Allocation:Q", format="$,.2f"),
                    alt.Tooltip("Min:Q", format="$,.2f"),
                    alt.Tooltip("Surplus/Deficit:Q", format="$,.2f"),
                ]
            )
            .properties(height=300)
        )
        st.altair_chart(chart1, use_container_width=True)

        st.header("📈 Surplus / Deficit by Category")
        sd = df[["Category", "Surplus/Deficit"]]
        chart2 = (
            alt.Chart(sd)
            .mark_bar()
            .encode(
                x="Category:N",
                y=alt.Y("Surplus/Deficit:Q", title="Surplus / Deficit ($)"),
                color=alt.condition(
                    alt.datum["Surplus/Deficit"] >= 0,
                    alt.value("#4caf50"),
                    alt.value("#e15759"),
                ),
            )
            .properties(height=300)
        )
        st.altair_chart(chart2, use_container_width=True)

        st.header("🍰 Allocation Distribution")
        pie_df = df[df["Allocation"] > 0]
        chart3 = (
            alt.Chart(pie_df)
            .mark_arc()
            .encode(
                theta="Allocation:Q",
                color=alt.Color("Category:N", legend=alt.Legend(title="Category")),
                tooltip=["Category", "Allocation"],
            )
            .properties(height=300)
        )
        st.altair_chart(chart3, use_container_width=True)

        # ── Download buttons ─────────────────────────────────────────────────────
        csv_data = df.to_csv(index=False).encode("utf-8")
        st.download_button(
            label="📥 Download as CSV",
            data=csv_data,
            file_name=f"budget_{datetime.now():%Y%m%d_%H%M%S}.csv",
            mime="text/csv"
        )

        to_excel = io.BytesIO()
        with pd.ExcelWriter(to_excel, engine="xlsxwriter") as writer:
            df.to_excel(writer, index=False, sheet_name="Budget")
        to_excel.seek(0)
        st.download_button(
            label="📥 Download as Excel",
            data=to_excel.getvalue(),
            file_name=f"budget_{datetime.now():%Y%m%d_%H%M%S}.xlsx",
            mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"
        )

        pdf_buffer = df_to_pdf(disp)
        st.download_button(
            label="📥 Download as PDF",
            data=pdf_buffer,
            file_name=f"budget_{datetime.now():%Y%m%d_%H%M%S}.pdf",
            mime="application/pdf"
        )

        # ── Future Feature Placeholder ──────────────────────────────────────────
        st.info("📈 Investment recommendations coming soon!")