import os
import json
import csv
from datetime import datetime

import numpy as np
import pandas as pd
import faiss
import streamlit as st
import altair as alt
from sentence_transformers import SentenceTransformer


# =========================================================
# PAGE CONFIG
# =========================================================
st.set_page_config(
    page_title="IGPA Legislation Explorer",
    layout="wide",
    initial_sidebar_state="expanded"
)


# =========================================================
# PASSWORD GATE
# =========================================================
APP_PASSWORD = "bills"

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

if not st.session_state.authed:
    st.markdown(
        """
        <style>
        .login-title {
            font-size: 2rem;
            font-weight: 700;
            text-align: center;
            margin-top: 3rem;
            margin-bottom: 0.35rem;
        }
        .login-subtitle {
            text-align: center;
            opacity: 0.8;
            margin-bottom: 2rem;
        }
        </style>
        """,
        unsafe_allow_html=True,
    )

    st.markdown('<div class="login-title">IGPA Legislation Explorer</div>', unsafe_allow_html=True)
    st.markdown(
        """
        <div class="login-subtitle">
        🔐 Quick riddle to enter:<br><br>
        <b>“What do lawmakers pass to create laws?”</b><br><br>
        <span style="opacity:0.8;">Answer format: all lowercase, one word</span>
        </div>
        """,
        unsafe_allow_html=True,
    )

    _, center_col, _ = st.columns([1, 1.2, 1])

    with center_col:
        pw = st.text_input("Your answer", type="password")
        if st.button("Enter", use_container_width=True):
            if not APP_PASSWORD:
                st.error("Server misconfigured: APP_PASSWORD not set in Space secrets.")
            elif pw.strip().lower() == APP_PASSWORD:
                st.session_state.authed = True
                st.rerun()
            else:
                st.error("Incorrect password.")

    st.stop()


# =========================================================
# THEME-SAFE STYLES
# =========================================================
st.markdown(
    """
    <style>
    .block-container {
        padding-top: 3rem;
        padding-bottom: 2rem;
        padding-left: 1.5rem;
        padding-right: 1.5rem;
    }
    .main-title {
        font-size: 2rem;
        font-weight: 700;
        margin-bottom: 0.25rem;
    }
    .sub-title {
        font-size: 1rem;
        margin-bottom: 1rem;
        opacity: 0.85;
    }
    .section-title {
        font-size: 1.15rem;
        font-weight: 700;
        margin-top: 1rem;
        margin-bottom: 0.7rem;
    }
    .kpi-card {
        padding: 1rem 1rem 0.8rem 1rem;
        border-radius: 14px;
        border: 1px solid rgba(128,128,128,0.25);
        background: rgba(255,255,255,0.04);
        margin-bottom: 0.75rem;
    }
    .kpi-label {
        font-size: 0.95rem;
        opacity: 0.8;
        margin-bottom: 0.35rem;
    }
    .kpi-value {
        font-size: 2rem;
        font-weight: 700;
        line-height: 1.1;
    }
    .result-meta {
        font-size: 0.9rem;
        opacity: 0.8;
        margin-top: 0.5rem;
        margin-bottom: 0.25rem;
    }
    div[data-testid="stDataFrame"] {
        border: 1px solid rgba(128,128,128,0.25);
        border-radius: 12px;
        overflow: hidden;
    }
    div.stButton > button {
        border-radius: 10px;
        font-weight: 600;
    }
    div[data-baseweb="select"] > div,
    div[data-baseweb="popover"] > div,
    div[data-testid="stDateInput"] > div,
    div[data-baseweb="input"] > div {
        border-radius: 10px;
    }
    div[data-baseweb="tag"] {
        border-radius: 999px !important;
    }
    div[data-testid="metric-container"] svg {
        display: none;
    }
    button[data-baseweb="tab"] {
        color: #000 !important;
        opacity: 1 !important;
        font-weight: 600 !important;
    }
    button[data-baseweb="tab"][aria-selected="true"] {
        border-bottom: 3px solid #2563eb !important;
    }
    </style>
    """,
    unsafe_allow_html=True
)


# =========================================================
# CONFIG
# =========================================================
DB_DIR = "."
FEEDBACK_CSV = os.path.join(DB_DIR, "impact_feedback.csv")
DEFAULT_TOP_K = 10

IMPACT_ORDER = [
    "Not Impactful",
    "Slightly Impactful",
    "Moderately Impactful",
    "Very Impactful"
]

DATE_COL = "status_date"
SUMMARY_COL = "Llama Summary"
LINK_COL = "ftp_url"

CATEGORY_COL = "category_std"
SUBCATEGORY_COL = "subcategory_std"

POLICY_COL = "policy_domain_final"
BENEFICIARY_CATEGORY_COL = "intended_beneficiaries_category"
BENEFICIARY_SUBCATEGORY_COL = "intended_beneficiaries_subcategory"
BENEFICIARY_KEYWORD_COL = "intended_beneficiaries_keyword"
BENEFICIARY_COL = "intended_beneficiaries_final"

STRATEGY_COL = "legislative_strategy_std"
GOAL_COL = "legislative_goal_std"
INTENT_COL = "intent_std"
INCREASING_COL = "increasing_aspects_std"
DECREASING_COL = "decreasing_aspects_std"
MOTIVATION_COL = "motivation_std"

IMPACT_COL_PREFERRED = "impact_rating_std"
IMPACT_COL_FALLBACK = "Impact Rating"


# =========================================================
# SIDEBAR TOP ACTION
# =========================================================
with st.sidebar:
    if st.button("Logout", use_container_width=True):
        st.session_state.authed = False
        st.rerun()


# =========================================================
# LOAD VECTOR DB
# =========================================================
@st.cache_resource
def load_vector_db(db_dir: str = DB_DIR):
    with open(os.path.join(db_dir, "config.json"), "r", encoding="utf-8") as f:
        cfg = json.load(f)

    index = faiss.read_index(os.path.join(db_dir, "faiss_index.bin"))
    meta = pd.read_parquet(os.path.join(db_dir, "metadata.parquet"))
    meta = meta.reset_index(drop=True)

    if "vec_id" not in meta.columns:
        meta = meta.reset_index().rename(columns={"index": "vec_id"})

    model = SentenceTransformer(cfg["embedding_model_name"])
    return index, meta, model, cfg


index, meta_df, embed_model, cfg = load_vector_db()

IMPACT_COL = IMPACT_COL_PREFERRED if IMPACT_COL_PREFERRED in meta_df.columns else IMPACT_COL_FALLBACK

if DATE_COL in meta_df.columns:
    meta_df[DATE_COL] = pd.to_datetime(meta_df[DATE_COL], errors="coerce")


# =========================================================
# HELPERS
# =========================================================
def impact_to_score(x):
    if pd.isna(x):
        return np.nan
    x = str(x).strip().lower()
    mapping = {
        "not impactful": 0,
        "slightly impactful": 1,
        "moderately impactful": 2,
        "very impactful": 3
    }
    return mapping.get(x, np.nan)


if "impact_rating_score" not in meta_df.columns and IMPACT_COL in meta_df.columns:
    meta_df["impact_rating_score"] = meta_df[IMPACT_COL].apply(impact_to_score)


DEFAULT_FILTERS = {
    "categories": [],
    "subcategories": [],
    "beneficiary_categories": [],
    "beneficiary_subcategories": [],
    "policy_domains": [],
    "impact_selected": [],
    "bill_statuses": [],
    "date_range": (
        meta_df[DATE_COL].min().date()
        if DATE_COL in meta_df.columns and pd.notna(meta_df[DATE_COL].min())
        else datetime.utcnow().date(),
        meta_df[DATE_COL].max().date()
        if DATE_COL in meta_df.columns and pd.notna(meta_df[DATE_COL].max())
        else datetime.utcnow().date(),
    ),
    "num_search_results": DEFAULT_TOP_K,
}

for key, value in DEFAULT_FILTERS.items():
    if key not in st.session_state:
        st.session_state[key] = value

if "search_results" not in st.session_state:
    st.session_state.search_results = None
if "current_query" not in st.session_state:
    st.session_state.current_query = ""
if "history" not in st.session_state:
    st.session_state.history = []


def embed_query(query: str):
    return embed_model.encode(
        [query],
        normalize_embeddings=True,
        convert_to_numpy=True
    ).astype("float32")


def append_feedback_row(
    bill_id,
    predicted_impact,
    user_response,
    corrected_impact=None,
    comment=None,
    path=FEEDBACK_CSV,
):
    try:
        file_exists = os.path.isfile(path)
        with open(path, "a", newline="", encoding="utf-8") as f:
            writer = csv.writer(f)
            if not file_exists:
                writer.writerow(
                    [
                        "timestamp",
                        "bill_id",
                        "predicted_impact",
                        "user_response",
                        "corrected_impact",
                        "comment",
                    ]
                )
            writer.writerow(
                [
                    datetime.utcnow().isoformat(),
                    bill_id,
                    predicted_impact,
                    user_response,
                    corrected_impact if corrected_impact else "",
                    comment if comment else "",
                ]
            )
    except Exception as e:
        st.error(f"Failed to save feedback: {str(e)}")


def get_multiselect_options(df, col_name):
    if col_name not in df.columns:
        return []
    vals = df[col_name].dropna().astype(str).str.strip()
    return sorted([v for v in vals.unique().tolist() if v])


def get_dependent_subcategories(df, category_col, subcategory_col, selected_categories):
    if subcategory_col not in df.columns:
        return []

    sub_df = df.copy()
    if selected_categories and category_col in df.columns:
        sub_df = sub_df[sub_df[category_col].isin(selected_categories)]

    vals = sub_df[subcategory_col].dropna().astype(str).str.strip()
    return sorted([v for v in vals.unique().tolist() if v])


def get_beneficiary_categories(df):
    if BENEFICIARY_CATEGORY_COL not in df.columns:
        return []
    vals = df[BENEFICIARY_CATEGORY_COL].dropna().astype(str).str.strip()
    return sorted([v for v in vals.unique().tolist() if v])


def get_beneficiary_subcategories(df, selected_categories):
    return get_dependent_subcategories(
        df,
        BENEFICIARY_CATEGORY_COL,
        BENEFICIARY_SUBCATEGORY_COL,
        selected_categories
    )


def get_categories(df):
    if CATEGORY_COL not in df.columns:
        return []
    vals = df[CATEGORY_COL].dropna().astype(str).str.strip()
    return sorted([v for v in vals.unique().tolist() if v])


def get_subcategories(df, selected_categories):
    return get_dependent_subcategories(
        df,
        CATEGORY_COL,
        SUBCATEGORY_COL,
        selected_categories
    )


def build_filter_mask(df):
    mask = pd.Series(True, index=df.index)

    selected_categories = st.session_state.get("categories", [])
    if selected_categories and CATEGORY_COL in df.columns:
        mask &= df[CATEGORY_COL].isin(selected_categories)

    selected_subcategories = st.session_state.get("subcategories", [])
    if selected_subcategories and SUBCATEGORY_COL in df.columns:
        mask &= df[SUBCATEGORY_COL].isin(selected_subcategories)

    selected_beneficiary_categories = st.session_state.get("beneficiary_categories", [])
    if selected_beneficiary_categories and BENEFICIARY_CATEGORY_COL in df.columns:
        mask &= df[BENEFICIARY_CATEGORY_COL].isin(selected_beneficiary_categories)

    selected_beneficiary_subcategories = st.session_state.get("beneficiary_subcategories", [])
    if selected_beneficiary_subcategories and BENEFICIARY_SUBCATEGORY_COL in df.columns:
        mask &= df[BENEFICIARY_SUBCATEGORY_COL].isin(selected_beneficiary_subcategories)

    selected_policy_domains = st.session_state.get("policy_domains", [])
    if selected_policy_domains and POLICY_COL in df.columns:
        mask &= df[POLICY_COL].isin(selected_policy_domains)

    selected_impact = st.session_state.get("impact_selected", [])
    if selected_impact and IMPACT_COL in df.columns:
        mask &= df[IMPACT_COL].isin(selected_impact)

    selected_statuses = st.session_state.get("bill_statuses", [])
    if selected_statuses and "status_desc" in df.columns:
        mask &= df["status_desc"].isin(selected_statuses)

    if "date_range" in st.session_state and st.session_state.date_range and DATE_COL in df.columns:
        dr = st.session_state.date_range
        if isinstance(dr, (tuple, list)) and len(dr) == 2:
            start, end = dr
        else:
            start = end = dr

        start = pd.to_datetime(start)
        end = pd.to_datetime(end)
        mask &= df[DATE_COL].between(start, end)

    return mask


def reset_filters():
    for key, value in DEFAULT_FILTERS.items():
        st.session_state[key] = value
    st.session_state.search_results = None
    st.session_state.current_query = ""
    st.rerun()


def get_first_available(row, cols, default=""):
    for c in cols:
        if c in row.index:
            val = row.get(c)
            if pd.notna(val) and str(val).strip():
                return val
    return default


# =========================================================
# SIDEBAR FILTERS
# =========================================================
with st.sidebar:
    st.header("Filters")

    if st.button("Reset Filters", use_container_width=True):
        reset_filters()

    st.caption("Tip: Apply filters here before reviewing bills and charts.")

    st.markdown("### Category & Subcategory")

    has_category_cols = CATEGORY_COL in meta_df.columns and SUBCATEGORY_COL in meta_df.columns

    if has_category_cols:
        category_options = get_categories(meta_df)
        st.multiselect(
            "Category",
            options=category_options,
            key="categories"
        )

        subcategory_options = get_subcategories(
            meta_df,
            st.session_state.categories
        )

        current_subcats = st.session_state.get("subcategories", [])
        valid_current_subcats = [x for x in current_subcats if x in subcategory_options]
        if current_subcats != valid_current_subcats:
            st.session_state["subcategories"] = valid_current_subcats

        st.multiselect(
            "Subcategory",
            options=subcategory_options,
            key="subcategories"
        )
    else:
        st.info("Category and Subcategory filters are unavailable because 'category_std' and 'subcategory_std' are not present in the dataset.")

    st.markdown("### Intended Beneficiaries")

    beneficiary_category_options = get_beneficiary_categories(meta_df)
    st.multiselect(
        "Intended Beneficiary Category",
        options=beneficiary_category_options,
        key="beneficiary_categories"
    )

    beneficiary_subcategory_options = get_beneficiary_subcategories(
        meta_df,
        st.session_state.beneficiary_categories
    )

    current_beneficiary_subcats = st.session_state.get("beneficiary_subcategories", [])
    valid_beneficiary_subcats = [x for x in current_beneficiary_subcats if x in beneficiary_subcategory_options]
    if current_beneficiary_subcats != valid_beneficiary_subcats:
        st.session_state["beneficiary_subcategories"] = valid_beneficiary_subcats

    st.multiselect(
        "Intended Beneficiary Subcategory",
        options=beneficiary_subcategory_options,
        key="beneficiary_subcategories"
    )

    st.markdown("### Policy, Impact, Status")

    st.multiselect(
        "Policy Area",
        options=get_multiselect_options(meta_df, POLICY_COL),
        key="policy_domains"
    )

    st.multiselect(
        "Impact Rating (STD)",
        options=IMPACT_ORDER,
        key="impact_selected"
    )

    st.multiselect(
        "Bill Status",
        options=get_multiselect_options(meta_df, "status_desc"),
        key="bill_statuses"
    )

    st.markdown("### Time Filter")

    min_date = (
        meta_df[DATE_COL].min().date()
        if DATE_COL in meta_df.columns and pd.notna(meta_df[DATE_COL].min())
        else datetime.utcnow().date()
    )
    max_date = (
        meta_df[DATE_COL].max().date()
        if DATE_COL in meta_df.columns and pd.notna(meta_df[DATE_COL].max())
        else datetime.utcnow().date()
    )

    st.date_input(
        "Status Date Range",
        value=st.session_state.get("date_range", (min_date, max_date)),
        min_value=min_date,
        max_value=max_date,
        key="date_range"
    )

    if os.path.exists(FEEDBACK_CSV):
        try:
            df_feedback = pd.read_csv(FEEDBACK_CSV)
            st.info(f"Feedback records: {len(df_feedback)}")
            with open(FEEDBACK_CSV, "rb") as f:
                st.download_button(
                    label="Download impact_feedback.csv",
                    data=f.read(),
                    file_name="impact_feedback.csv",
                    mime="text/csv"
                )
        except Exception:
            st.info("Feedback CSV ready")

    with st.expander("Search History"):
        for i, item in enumerate(reversed(st.session_state.history[-5:]), 1):
            st.write(f"{i}. {item.get('query', '')}")


# =========================================================
# FILTERED DATA
# =========================================================
filtered_df = meta_df[build_filter_mask(meta_df)].copy()


# =========================================================
# TABS
# =========================================================
tab_search, tab_trends = st.tabs(["Search & Results", "Trends & Insights"])


# =========================================================
# SEARCH TAB
# =========================================================
with tab_search:
    st.markdown('<div class="main-title">IGPA Legislation Explorer</div>', unsafe_allow_html=True)
    st.markdown(
        '<div class="sub-title">Search, filter, and analyze legislative bills by impact, policy area, category, subcategory, and intended beneficiaries.</div>',
        unsafe_allow_html=True
    )

    col1, col2, col3, col4 = st.columns(4)

    total_bills = len(filtered_df)
    policy_count = filtered_df[POLICY_COL].nunique() if POLICY_COL in filtered_df.columns else 0
    beneficiary_count = filtered_df[BENEFICIARY_COL].nunique() if BENEFICIARY_COL in filtered_df.columns else 0
    very_impactful = (filtered_df[IMPACT_COL] == "Very Impactful").sum() if IMPACT_COL in filtered_df.columns else 0

    with col1:
        st.markdown(f"""
            <div class="kpi-card">
                <div class="kpi-label">Total Bills</div>
                <div class="kpi-value">{total_bills}</div>
            </div>
        """, unsafe_allow_html=True)

    with col2:
        st.markdown(f"""
            <div class="kpi-card">
                <div class="kpi-label">Policy Domains</div>
                <div class="kpi-value">{policy_count}</div>
            </div>
        """, unsafe_allow_html=True)

    with col3:
        st.markdown(f"""
            <div class="kpi-card">
                <div class="kpi-label">Beneficiary Groups</div>
                <div class="kpi-value">{beneficiary_count}</div>
            </div>
        """, unsafe_allow_html=True)

    with col4:
        st.markdown(f"""
            <div class="kpi-card">
                <div class="kpi-label">Very Impactful Bills</div>
                <div class="kpi-value">{very_impactful}</div>
            </div>
        """, unsafe_allow_html=True)

    st.markdown('<div class="section-title">Most Impacted Beneficiary Keywords</div>', unsafe_allow_html=True)

    if BENEFICIARY_KEYWORD_COL in filtered_df.columns and "impact_rating_score" in filtered_df.columns:
        impact_df = (
            filtered_df.dropna(subset=[BENEFICIARY_KEYWORD_COL, "impact_rating_score"])
            .groupby(BENEFICIARY_KEYWORD_COL)
            .agg(
                avg_impact=("impact_rating_score", "mean"),
                bills=("bill_id", "count"),
                top_bills=("title", lambda x: "; ".join(x.head(5)))
            )
            .reset_index()
            .sort_values(["avg_impact", "bills"], ascending=[False, False])
            .head(15)
        )

        if not impact_df.empty:
            chart = (
                alt.Chart(impact_df)
                .mark_bar()
                .encode(
                    x=alt.X(f"{BENEFICIARY_KEYWORD_COL}:N", sort="-y", title="Beneficiary Keyword"),
                    y=alt.Y("avg_impact:Q", title="Average Impact Score"),
                    color=alt.Color(
                        "avg_impact:Q",
                        scale=alt.Scale(domain=[0, 3], range=["#fde68a", "#dc2626"]),
                        legend=alt.Legend(title="Average Impact")
                    ),
                    tooltip=[
                        alt.Tooltip(f"{BENEFICIARY_KEYWORD_COL}:N", title="Beneficiary Keyword"),
                        alt.Tooltip("avg_impact:Q", format=".2f", title="Average Impact"),
                        alt.Tooltip("bills:Q", title="Number of Bills"),
                        alt.Tooltip("top_bills:N", title="Top Bills")
                    ]
                )
                .properties(height=350)
            )
            st.altair_chart(chart, use_container_width=True)
        else:
            st.info("No beneficiary keyword impact data available for the current filters.")
    else:
        st.info("Beneficiary keyword or impact score columns are not available.")

    table_header_col1 = st.columns([1])[0]

    with table_header_col1:
        st.markdown('<div class="section-title">Bills Matching Selected Filters</div>', unsafe_allow_html=True)

    st.caption(
        "These rows reflect the active toolbar filters. ILGA columns are source legislative fields. "
        "Llama columns are AI-generated summaries or standardized outputs."
    )

    display_cols = {
        "bill_number": "Bill Number (ILGA)",
        "title": "Title (ILGA)",
        "description": "Description (ILGA)",
        CATEGORY_COL: "Category",
        SUBCATEGORY_COL: "Subcategory",
        BENEFICIARY_CATEGORY_COL: "Intended Beneficiary Category",
        BENEFICIARY_SUBCATEGORY_COL: "Intended Beneficiary Subcategory",
        SUMMARY_COL: "Llama Summary",
        "Potential Impact": "Potential Impact (Llama)",
        IMPACT_COL: "Impact Rating (STD)",
        "status_desc": "Status (ILGA)",
        LINK_COL: "Bill Link (ILGA)"
    }

    available_cols = [c for c in display_cols if c in filtered_df.columns]

    filter_bill_df = (
        filtered_df[available_cols]
        .rename(columns=display_cols)
        .copy()
    )

    st.dataframe(
        filter_bill_df,
        use_container_width=True,
        column_config={
            "Bill Link (ILGA)": st.column_config.LinkColumn(
                label="Bill Link (ILGA)",
                display_text="Open Bill"
            )
        } if "Bill Link (ILGA)" in filter_bill_df.columns else None
    )

    csv = filter_bill_df.to_csv(index=False).encode("utf-8")

    st.download_button(
        label="Download as CSV",
        data=csv,
        file_name="filtered_bills.csv",
        mime="text/csv"
    )

    search_header_col1, search_header_col2 = st.columns([6, 2])

    with search_header_col1:
        st.markdown('<div class="section-title">Search Bills</div>', unsafe_allow_html=True)

    with search_header_col2:
        st.slider(
            "Number of Results – Search Query Results",
            5, 50, st.session_state.get("num_search_results", DEFAULT_TOP_K), 5,
            key="num_search_results"
        )

    search_col, button_col = st.columns([6, 1])

    with search_col:
        query = st.text_input(
            "Ask a question about legislation",
            value=st.session_state.current_query,
            placeholder="Example: bills related to higher education funding, healthcare workers, or workforce policy"
        )

    with button_col:
        st.markdown("<div style='height: 28px;'></div>", unsafe_allow_html=True)
        search_clicked = st.button("Search", use_container_width=True)

    if search_clicked and query.strip():
        st.session_state.current_query = query
        st.session_state.history.append({"query": query})

        q_vec = embed_query(query)
        top_k = st.session_state.num_search_results
        n_search = min(len(meta_df), top_k * 5)

        scores, ids = index.search(q_vec, n_search)
        ids, scores = ids[0], scores[0]

        allowed = set(filtered_df.index)
        kept = [(i, s) for i, s in zip(ids, scores) if i in allowed][:top_k]

        if not kept:
            st.warning("No results found for this query under the current filters.")
            st.session_state.search_results = None
        else:
            results = meta_df.loc[[i for i, _ in kept]].copy()
            results["similarity"] = [s for _, s in kept]
            st.session_state.search_results = results

    if st.session_state.search_results is not None:
        results = st.session_state.search_results

        st.markdown('<div class="section-title">Filtered Results Table</div>', unsafe_allow_html=True)
        st.caption("Search results are ranked by semantic similarity after applying the active toolbar filters.")
        st.caption("ILGA fields come from source legislation data. Llama fields are AI-generated summaries or standardized outputs.")

        review_cols = [
            "bill_number",
            "title",
            "description",
            CATEGORY_COL,
            SUBCATEGORY_COL,
            BENEFICIARY_CATEGORY_COL,
            BENEFICIARY_SUBCATEGORY_COL,
            SUMMARY_COL,
            GOAL_COL,
            "Potential Impact",
            INCREASING_COL,
            DECREASING_COL,
            IMPACT_COL,
            "similarity",
            LINK_COL
        ]

        review_df = results[[c for c in review_cols if c in results.columns]].copy()

        review_df.rename(
            columns={
                "bill_number": "Bill Number",
                "title": "Title",
                "description": "Description from Legiscan",
                CATEGORY_COL: "Category",
                SUBCATEGORY_COL: "Subcategory",
                BENEFICIARY_CATEGORY_COL: "Main Category (LLaMA Generated Response)",
                BENEFICIARY_SUBCATEGORY_COL: "Beneficiary Subcategory",
                SUMMARY_COL: "LLaMA Summary",
                GOAL_COL: "Legislative Goal",
                "Potential Impact": "Potential Impact (LLaMA)",
                INCREASING_COL: "Increasing Aspects (STD)",
                DECREASING_COL: "Decreasing Aspects (STD)",
                IMPACT_COL: "Impact Rating (STD)",
                "similarity": "Similarity Score",
                LINK_COL: "Bill URL (ILGA)"
            },
            inplace=True
        )

        st.dataframe(
            review_df,
            use_container_width=True,
            column_config={
                "Bill URL (ILGA)": st.column_config.LinkColumn(
                    "Bill URL (ILGA)",
                    display_text="Open bill"
                )
            } if "Bill URL (ILGA)" in review_df.columns else None
        )

        csv_data = review_df.to_csv(index=False).encode("utf-8")
        st.download_button(
            label="Download CSV",
            data=csv_data,
            file_name="search_results.csv",
            mime="text/csv"
        )

        st.markdown("---")
        st.markdown('<div class="section-title">Top Matching Bills</div>', unsafe_allow_html=True)

        for idx, row in results.iterrows():
            bill_number_val = get_first_available(row, ["bill_number"])
            title_val = get_first_available(row, ["title"])
            description_val = get_first_available(row, ["description"])
            category_val = get_first_available(row, [CATEGORY_COL])
            subcategory_val = get_first_available(row, [SUBCATEGORY_COL])
            summary_val = get_first_available(row, [SUMMARY_COL])

            st.markdown(f"**Bill Number**  \n{bill_number_val if bill_number_val else 'N/A'}")
            st.markdown(f"**Title**  \n{title_val if title_val else 'N/A'}")

            if description_val:
                st.markdown(f"**Description from Legiscan**  \n{description_val}")

            if category_val:
                st.markdown(f"**Category**  \n{category_val}")

            if subcategory_val:
                st.markdown(f"**Subcategory**  \n{subcategory_val}")

            if pd.notna(row.get(LINK_COL)):
                st.markdown(f"[Open Full Bill]({row.get(LINK_COL)})")

            if summary_val:
                with st.expander("Summary from LLaMA", expanded=True):
                    st.write(summary_val)

            detail_map = {
                "Status": "status_desc",
                "Legislative Goal": "Legislative Goal",
                "Key Provisions": "Key Provisions",
                "Increasing Aspects": "Increasing Aspects",
                "Decreasing Aspects": "Decreasing Aspects",
                "Category & Subcategory": "Category & Subcategory",
                "Ideological Alignment": "Ideological Alignment",
                "Potential Impact": "Potential Impact",
                "Original Law": "Original Law",
                "committee": "committee",
                "last_action_date": "last_action_date",
                "last_action": "last_action",
                "ILGA State Link": "state_link",
                "Legiscan URL": "url",
                "ftp_url": "ftp_url",
            }

            with st.expander("More Details"):
                for label, col_name in detail_map.items():
                    if col_name in results.columns:
                        val = row.get(col_name)
                        if pd.notna(val) and str(val).strip():
                            if label in ["ILGA State Link", "Legiscan URL", "ftp_url"]:
                                st.markdown(f"**{label}:** [Open Link]({val})")
                            else:
                                st.write(f"**{label}:** {val}")

            with st.expander("Impact Rating Accuracy", expanded=False):
                st.markdown("**Is this impact rating accurate?**")
                predicted_impact = row.get(IMPACT_COL, "")
                bill_id_safe = str(row.get("bill_id", idx))

                feedback_submitted = st.session_state.get(f"feedback_done_{bill_id_safe}", False)

                if feedback_submitted:
                    st.success("Thank you for your feedback.")
                    st.caption(f"Bill: {row.get('bill_number', 'N/A')} | Saved to impact_feedback.csv")
                else:
                    col_yes, col_no = st.columns(2)

                    with col_yes:
                        if st.button("Yes - Accurate", key=f"yes_{bill_id_safe}", use_container_width=True):
                            append_feedback_row(
                                bill_id=bill_id_safe,
                                predicted_impact=predicted_impact,
                                user_response="Yes",
                                corrected_impact=None,
                                comment=None,
                            )
                            st.session_state[f"feedback_done_{bill_id_safe}"] = True
                            st.rerun()

                    with col_no:
                        if st.button("No - Incorrect", key=f"no_{bill_id_safe}", use_container_width=True):
                            st.session_state[f"show_corrected_{bill_id_safe}"] = True
                            st.rerun()

                    if st.session_state.get(f"show_corrected_{bill_id_safe}", False):
                        corrected_value = st.selectbox(
                            "Correct impact rating",
                            IMPACT_ORDER,
                            key=f"corrected_{bill_id_safe}",
                        )

                        comment = st.text_area(
                            "Optional correction note",
                            max_chars=250,
                            key=f"comment_{bill_id_safe}",
                            placeholder="Add a short note explaining why the impact rating should change"
                        )

                        col_submit, col_cancel = st.columns([3, 1])

                        with col_submit:
                            if st.button("Submit Feedback", key=f"submit_{bill_id_safe}", type="primary"):
                                append_feedback_row(
                                    bill_id=bill_id_safe,
                                    predicted_impact=predicted_impact,
                                    user_response="No",
                                    corrected_impact=corrected_value,
                                    comment=comment,
                                )
                                st.session_state[f"feedback_done_{bill_id_safe}"] = True
                                st.session_state[f"show_corrected_{bill_id_safe}"] = False
                                st.rerun()

                        with col_cancel:
                            if st.button("Cancel", key=f"cancel_{bill_id_safe}"):
                                st.session_state[f"show_corrected_{bill_id_safe}"] = False
                                st.rerun()

            st.markdown("---")


# =========================================================
# TRENDS TAB
# =========================================================
with tab_trends:
    top_policy = (
        filtered_df[POLICY_COL].value_counts().head(1)
        if POLICY_COL in filtered_df.columns else pd.Series(dtype=int)
    )

    top_beneficiaries = (
        filtered_df[BENEFICIARY_COL].value_counts().head(1)
        if BENEFICIARY_COL in filtered_df.columns else pd.Series(dtype=int)
    )

    strategy_impact = (
        filtered_df[filtered_df[IMPACT_COL].notna()]
        .groupby(STRATEGY_COL)[IMPACT_COL]
        .apply(lambda x: (x == "Very Impactful").sum())
        if STRATEGY_COL in filtered_df.columns and IMPACT_COL in filtered_df.columns
        else pd.Series(dtype=int)
    )

    avg_impact_ben = (
        filtered_df.dropna(subset=["impact_rating_score"])
        .groupby(BENEFICIARY_COL)["impact_rating_score"]
        .mean()
        .sort_values(ascending=False)
        if BENEFICIARY_COL in filtered_df.columns and "impact_rating_score" in filtered_df.columns
        else pd.Series(dtype=float)
    )

    total_bills = len(filtered_df)
    total_high_impact = (
        (filtered_df[IMPACT_COL] == "Very Impactful").sum()
        if IMPACT_COL in filtered_df.columns else 0
    )

    st.markdown("### Key Insights")
    st.write(f"**Total Bills Considered:** {total_bills}")
    st.write(f"**Total Very Impactful Bills:** {total_high_impact}")
    st.write(
        f"**Most Active Policy Domain:** {top_policy.index[0]} ({top_policy.iloc[0]} bills)"
        if not top_policy.empty else "No data"
    )
    st.write(
        f"**Most Benefited Group:** {top_beneficiaries.index[0]} ({top_beneficiaries.iloc[0]} bills)"
        if not top_beneficiaries.empty else "No data"
    )
    st.write(
        f"**Strategy Producing Most Very Impactful Bills:** {strategy_impact.idxmax()}"
        if not strategy_impact.empty else "N/A"
    )
    st.write(
        f"**Highest Average Impact (Beneficiary):** {avg_impact_ben.index[0]} ({avg_impact_ben.iloc[0]:.2f})"
        if not avg_impact_ben.empty else "N/A"
    )
    st.markdown("---")

    col1, col2 = st.columns(2)

    with col1:
        st.markdown("### Policy Domain Activity")
        if POLICY_COL in filtered_df.columns:
            policy_agg = (
                filtered_df.groupby(POLICY_COL)
                .agg(
                    Count=("bill_id", "count"),
                    avg_impact=("impact_rating_score", "mean")
                    if "impact_rating_score" in filtered_df.columns else ("bill_id", "count"),
                    top_bills=("title", lambda x: "; ".join(x.head(5))),
                    top_beneficiaries=(
                        BENEFICIARY_COL,
                        lambda x: ", ".join(x.value_counts().head(3).index)
                    ) if BENEFICIARY_COL in filtered_df.columns else ("title", lambda x: ""),
                    recent_date=(
                        DATE_COL,
                        lambda x: x.max().strftime("%Y-%m-%d") if pd.notna(x.max()) else ""
                    ) if DATE_COL in filtered_df.columns else ("title", lambda x: ""),
                    bill_numbers=("bill_number", lambda x: ", ".join(map(str, x.head(5))))
                )
                .reset_index()
                .rename(columns={POLICY_COL: "Policy Domain"})
            )

            policy_chart = (
                alt.Chart(policy_agg)
                .mark_bar()
                .encode(
                    x=alt.X("Policy Domain:N", sort="-y", title="Policy Domain"),
                    y=alt.Y("Count:Q", title="Number of Bills"),
                    color=alt.Color(
                        "avg_impact:Q",
                        title="Average Impact",
                        scale=alt.Scale(scheme="orangered")
                    ),
                    tooltip=[
                        alt.Tooltip("Policy Domain:N"),
                        alt.Tooltip("Count:Q", title="Number of Bills"),
                        alt.Tooltip("avg_impact:Q", format=".2f", title="Average Impact"),
                        alt.Tooltip("top_bills:N", title="Top Bills"),
                        alt.Tooltip("top_beneficiaries:N", title="Top Beneficiaries"),
                        alt.Tooltip("recent_date:N", title="Most Recent Bill"),
                        alt.Tooltip("bill_numbers:N", title="Bill Numbers")
                    ]
                )
                .properties(height=400)
            )
            st.altair_chart(policy_chart, use_container_width=True)
        else:
            st.write("No policy domain data available.")

    with col2:
        st.markdown("### Impact Distribution")
        if IMPACT_COL in filtered_df.columns:
            impact_hover = (
                filtered_df[filtered_df[IMPACT_COL].notna()]
                .groupby(IMPACT_COL)
                .agg(
                    Count=("bill_id", "count"),
                    top_beneficiaries=(
                        BENEFICIARY_KEYWORD_COL,
                        lambda x: ", ".join(x.dropna().astype(str).value_counts().head(3).index)
                    ) if BENEFICIARY_KEYWORD_COL in filtered_df.columns else ("bill_id", lambda x: ""),
                    top_motivation=(
                        MOTIVATION_COL,
                        lambda x: ", ".join(x.dropna().astype(str).value_counts().head(3).index)
                    ) if MOTIVATION_COL in filtered_df.columns else ("bill_id", lambda x: "")
                )
                .reindex(IMPACT_ORDER)
                .reset_index()
            )

            impact_chart = (
                alt.Chart(impact_hover)
                .mark_bar()
                .encode(
                    x=alt.X(f"{IMPACT_COL}:N", sort=IMPACT_ORDER, title="Impact Level"),
                    y=alt.Y("Count:Q"),
                    color=alt.Color("Count:Q", scale=alt.Scale(scheme="reds"), legend=None),
                    tooltip=[
                        alt.Tooltip(f"{IMPACT_COL}:N", title="Impact Level"),
                        alt.Tooltip("Count:Q", title="Count"),
                        alt.Tooltip("top_beneficiaries:N", title="Top Beneficiaries"),
                        alt.Tooltip("top_motivation:N", title="Top Motivation")
                    ]
                )
                .properties(height=300)
            )
            st.altair_chart(impact_chart, use_container_width=True)
        else:
            st.write("No impact rating data available.")

    st.markdown("### Legislative Strategy: Very Impactful Bills")
    if STRATEGY_COL in filtered_df.columns and IMPACT_COL in filtered_df.columns:
        strategy_high_impact = (
            filtered_df[filtered_df[IMPACT_COL].notna()]
            .groupby(STRATEGY_COL)
            .agg(
                Very_Impactful_Bills=(IMPACT_COL, lambda x: (x == "Very Impactful").sum()),
                top_bills=("title", lambda x: "; ".join(x.head(5))),
                top_beneficiaries=(
                    BENEFICIARY_COL,
                    lambda x: ", ".join(x.value_counts().head(3).index)
                ) if BENEFICIARY_COL in filtered_df.columns else ("title", lambda x: ""),
                recent_date=(
                    DATE_COL,
                    lambda x: x.max().strftime("%Y-%m-%d") if pd.notna(x.max()) else ""
                ) if DATE_COL in filtered_df.columns else ("title", lambda x: "")
            )
            .reset_index()
            .rename(columns={STRATEGY_COL: "Strategy"})
        )

        strategy_chart = (
            alt.Chart(strategy_high_impact)
            .mark_bar()
            .encode(
                x=alt.X("Strategy:N", sort="-y", title="Strategy"),
                y=alt.Y("Very_Impactful_Bills:Q", title="Very Impactful Bills"),
                color=alt.Color("Very_Impactful_Bills:Q", scale=alt.Scale(scheme="orangered")),
                tooltip=[
                    alt.Tooltip("Strategy:N"),
                    alt.Tooltip("Very_Impactful_Bills:Q"),
                    alt.Tooltip("top_bills:N", title="Top Bills"),
                    alt.Tooltip("top_beneficiaries:N", title="Top Beneficiaries"),
                    alt.Tooltip("recent_date:N", title="Most Recent Bill")
                ]
            )
            .properties(height=400)
        )
        st.altair_chart(strategy_chart, use_container_width=True)
    else:
        st.write("No legislative strategy data available for selected filters.")

    st.markdown("### Beneficiary Coverage & Average Impact")
    if BENEFICIARY_COL in filtered_df.columns and "impact_rating_score" in filtered_df.columns:
        ben_df = (
            filtered_df.dropna(subset=[BENEFICIARY_COL, "impact_rating_score"])
            .groupby(BENEFICIARY_COL)
            .agg(
                total_bills=("bill_id", "count"),
                avg_impact=("impact_rating_score", "mean"),
                top_bills=("title", lambda x: "; ".join(x.head(5))),
                recent_date=(
                    DATE_COL,
                    lambda x: x.max().strftime("%Y-%m-%d") if pd.notna(x.max()) else ""
                ) if DATE_COL in filtered_df.columns else ("title", lambda x: ""),
                bill_numbers=("bill_number", lambda x: ", ".join(map(str, x.head(5))))
            )
            .reset_index()
        )

        if not ben_df.empty:
            ben_chart = (
                alt.Chart(ben_df)
                .mark_rect()
                .encode(
                    x=alt.X("total_bills:Q", title="Number of Bills"),
                    y=alt.Y(f"{BENEFICIARY_COL}:N", sort="-x", title="Beneficiary Group"),
                    color=alt.Color(
                        "avg_impact:Q",
                        scale=alt.Scale(domain=[0, 3], range=["#fde68a", "#dc2626"]),
                        legend=alt.Legend(title="Average Impact Score")
                    ),
                    tooltip=[
                        alt.Tooltip(f"{BENEFICIARY_COL}:N", title="Beneficiary"),
                        alt.Tooltip("total_bills:Q", title="Number of Bills"),
                        alt.Tooltip("avg_impact:Q", format=".2f", title="Average Impact"),
                        alt.Tooltip("top_bills:N", title="Top Bills"),
                        alt.Tooltip("recent_date:N", title="Most Recent Bill"),
                        alt.Tooltip("bill_numbers:N", title="Bill Numbers")
                    ]
                )
                .properties(height=400)
            )
            st.altair_chart(ben_chart, use_container_width=True)
        else:
            st.write("No beneficiary impact data available for selected filters.")
    else:
        st.write("No beneficiary data available for selected filters.")