src/streamlit_app.py

import asyncio
import sys
import uuid
from datetime import datetime
from pathlib import Path

import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import streamlit as st

# ── Path setup so `src` is importable when running from src/ or project root ──
_here = Path(__file__).resolve().parent  # src/
_root = _here.parent  # project root
for _p in [str(_here), str(_root)]:
    if _p not in sys.path:
        sys.path.insert(0, _p)

from controllers._data_extractor import DataExtractorController, UserQuery
from schemas import Message

INCIDENT_COLOR_MAP = {
    "1": ("#ef4444", "Fire"),
    "2": ("#eab308", "Rupt/Exp"),
    "3": ("#3b82f6", "EMS"),
    "4": ("#22c55e", "Hazardous"),
    "5": ("#a855f7", "Public Assist"),
    "6": ("#06b6d4", "Good Intent"),
    "7": ("#9ca3af", "False Alarms"),
    "8": ("#84cc16", "Weather"),
    "9": ("#c0c0c0", "Special Type"),
}
INCIDENT_NULL_COLOR = ("#111111", "Not Entered")
INCIDENT_BLANK_COLOR = ("#f97316", "Blank")

INICIDENT_CATEGORY_NAMES = [
    "Fire",
    "Rupt/Exp",
    "EMS",
    "Hazardous",
    "Public Assist",
    "Good Intent",
    "False Alarms",
    "Weather",
    "Special Type",
    "Not Entered",
    "Blank",
]

INCIDENT_NAME_COLOR_MAP = {
    "fire": ("#ef4444", "Fire"),
    "rupt/exp": ("#eab308", "Rupt/Exp"),
    "ems": ("#3b82f6", "EMS"),
    "hazardous": ("#22c55e", "Hazardous"),
    "public assist": ("#a855f7", "Public Assist"),
    "good intent": ("#06b6d4", "Good Intent"),
    "false alarms": ("#9ca3af", "False Alarms"),
    "weather": ("#84cc16", "Weather"),
    "special type": ("#c0c0c0", "Special Type"),
}

INCIDENT_COL_NAMES = {
    "incidenttype",
    "incident_type",
    "incidentclassification",
    "incident_classification",
    "incident_category",
    "incidentcategory",
}


def detect_date_columns(df: pd.DataFrame):
    for col in df.columns:
        if pd.api.types.is_numeric_dtype(df[col]):
            continue
        
        try:
            numeric_like_ratio = pd.to_numeric(df[col], errors="coerce").notna().mean()
            if numeric_like_ratio > 0.8:
                continue
        except Exception:
            pass

        if pd.api.types.is_datetime64_any_dtype(df[col]):
            return col
        else:
            try:
                converted = pd.to_datetime(df[col], errors="coerce")
                valid_ratio = converted.notna().mean()
                if valid_ratio > 0.8:
                    return col
            except Exception:
                pass

    return None

def get_date_range(df: pd.DataFrame):
    date_col = detect_date_columns(df)

    if not date_col:
        return None

    series = pd.to_datetime(df[date_col], errors="coerce").dropna()
    if series.empty:
        return None

    col_min = series.min()
    col_max = series.max()

    if col_min and col_max:
        return f"{col_min.strftime('%b %d, %Y')} → {col_max.strftime('%b %d, %Y')}"

    return None

def _detect_incident_col(df: pd.DataFrame) -> str | None:
    incident_cols = [col for col in df.columns if col.strip().lower() in INCIDENT_COL_NAMES]
    return incident_cols if incident_cols else None


def _incident_label_and_color(value) -> tuple[str, str]:
    """Return (display_label, hex_color) for a raw incidenttype value."""
    if value is None or value == "nan" or (isinstance(value, float) and pd.isna(value)):
        return INCIDENT_NULL_COLOR[1], INCIDENT_NULL_COLOR[0]
    s = str(value).strip()
    if s == "":
        return INCIDENT_BLANK_COLOR[1], INCIDENT_BLANK_COLOR[0]
    prefix = s[0].upper() if s[0].isalpha() else s[0]
    key = s[0].upper() if s[0].upper() == "S" else s[0]
    if key in INCIDENT_COLOR_MAP:
        color, label = INCIDENT_COLOR_MAP[key]
        return label, color
    elif any(
        category_name.lower() in s.lower() for category_name in INICIDENT_CATEGORY_NAMES
    ):
        category_name_found = next(
            category_name
            for category_name in INICIDENT_CATEGORY_NAMES
            if category_name.lower() in s.lower()
        )
        name_key = category_name_found.lower()
        if name_key in INCIDENT_NAME_COLOR_MAP:
            color, label = INCIDENT_NAME_COLOR_MAP[name_key]
            return label, color
    return s, "#6b7280"


def _add_incident_category(df: pd.DataFrame, col: str) -> pd.DataFrame:
    df = df.copy()
    if col not in df.columns or df.empty:
        return df
    df = df.reset_index(drop=True)
    plain = df[col].astype(str).where(df[col].notna(), other=None)
    mapped = plain.map(_incident_label_and_color)
    df["_incident_label"] = mapped.apply(lambda x: x[0])
    df["_incident_color"] = mapped.apply(lambda x: x[1])
    return df


# ── Page config ──────────────────────────────────────────────────────────────
st.set_page_config(
    page_title="Data Extractor AI",
    page_icon="🔍",
    layout="wide",
    initial_sidebar_state="expanded",
)


# ── Plotly template (adapts to Streamlit theme) ──────────────────────────────
def _get_plotly_template():
    """Return a Plotly template that works with Streamlit's current theme."""
    return "plotly_white"


# ── Chart rendering ───────────────────────────────────────────────────────────

CHART_ALIASES = {
    "pie": "pie",
    "bar": "bar",
    "line": "line",
    "bar_chart": "bar",
    "vertical_bar": "bar",
    "column": "bar",
    "grouped_bar": "bar",
    "stacked_bar": "bar",
    "line_chart": "line",
    "time_series": "line",
    "trend": "line",
    "donut": "pie",
    "doughnut": "pie",
}


def _normalise_chart_type(raw: str | None) -> str | None:
    if not raw:
        return None
    return CHART_ALIASES.get(raw.lower().strip(), raw.lower().strip())


def _guess_chart_type(df: pd.DataFrame) -> str:
    cols = list(df.columns)
    n_cols = len(cols)
    n_rows = len(df)

    numeric = df.select_dtypes(include="number").columns.tolist()
    if n_cols == 1 and numeric:
        return "histogram"

    cat = df.select_dtypes(exclude="number").columns.tolist()
    if n_cols == 2 and len(cat) == 1 and len(numeric) == 1:
        return "bar" if n_rows <= 50 else "line"

    if len(numeric) >= 2:
        return "line"

    return "bar"

def normalize(x):
    try:
        num = float(x)
        if num.is_integer():
            return str(int(num))   
        return str(num)           
    except:
        return x
    
def sort_key(x):
    try:
        return (0, float(x)) 
    except:
        return (1, x)          


def get_categorical_columns(df: pd.DataFrame, column_name: str) -> list[str]:
    df[column_name] = df[column_name].astype(str)
    df[column_name] = df[column_name].map(normalize)
    cats = sorted(df[column_name].dropna().unique(), key=sort_key)
    df[column_name] = pd.Categorical(df[column_name], categories=cats, ordered=True)
    s = df.sort_values(by=column_name).reset_index(drop=True)[column_name]
    return s.tolist()

def render_chart(
    df: pd.DataFrame,
    incident_cols: list[str] | None = None,
    chart_type_raw: str | None = None,
    key_prefix: str = "chart",
):
    if df.empty:
        st.info("No data to chart.")
        return

    df = df.copy()
    for c in df.columns:
        if hasattr(df[c], "cat"):
            df[c] = df[c].astype(str).replace("nan", None)

    chart_type = _normalise_chart_type(chart_type_raw) or _guess_chart_type(df)

    incident_cols = incident_cols or []
    active_incident_col: str | None = None
    df_plot = df.copy()

    cols = list(df.columns)
    numeric = df.select_dtypes(include="number").columns.tolist()
    cat = df.select_dtypes(exclude="number").columns.tolist()

    default_x = (
        incident_cols[0]
        if incident_cols and incident_cols[0] in cols
        else (cat[0] if cat else cols[0])
    )
    default_y = numeric[0] if numeric else (cols[1] if len(cols) > 1 else cols[0])

    # ── Column selectors ──
    st.caption("⚙️ Configure columns")

    if chart_type == "pie":
        c1, c2 = st.columns(2)
        with c1:
            x_col = st.selectbox(
                "Labels",
                options=cols,
                index=cols.index(default_x) if default_x in cols else 0,
                key=f"{key_prefix}_pie_x",
            )
        with c2:
            val_opts = numeric if numeric else cols
            for col in cols:
                try:
                    if pd.to_numeric(df[col], errors="coerce").notna().all() and col not in val_opts:
                        val_opts.append(col)
                except Exception:
                    pass
            y_col = st.selectbox(
                "Values",
                options=val_opts,
                index=val_opts.index(default_y) if default_y in val_opts else 0,
                key=f"{key_prefix}_pie_y",
            )
        color_col = None
    else:
        c1, c2, c3 = st.columns(3)
        with c1:
            x_col = st.selectbox(
                "X axis",
                options=cols,
                index=cols.index(default_x) if default_x in cols else 0,
                key=f"{key_prefix}_x",
            )
            view_all_labels = st.checkbox(
                    "View All Labels",
                    key=f"{key_prefix}_x_all_labels",
                )
        with c2:
            y_opts = numeric if numeric else cols
            y_col = st.selectbox(
                "Y axis",
                options=y_opts,
                index=y_opts.index(default_y) if default_y in y_opts else 0,
                key=f"{key_prefix}_y",
            )
        with c3:
            color_options = ["None"] + [c for c in cols if c not in (x_col, y_col)]
            color_sel = st.selectbox(
                "Color / Group",
                options=color_options,
                index=0,
                key=f"{key_prefix}_color",
            )
            view_horizontal_stacked = st.checkbox(
                "Horizontal Stacked",
                key=f"{key_prefix}_stacked",
            )
            color_col = None if color_sel == "None" else color_sel

    # ── Incident color mapping ─────────────────────────
    incident_color_map = None
    if active_incident_col and "_incident_label" in df_plot.columns:
        incident_color_map = dict(
            zip(df_plot["_incident_label"], df_plot["_incident_color"])
        )

    # ── Build chart ────────────────────────────────────
    fig = None
    tmpl = _get_plotly_template()

    date_range = get_date_range(df_plot)

    title = f"{y_col} by {x_col}"

    if date_range:
        title += f" ({date_range})"

    try:
        # ───── BAR ─────
        if chart_type == "bar":
            if view_all_labels:
                df_plot[x_col] = df_plot[x_col].astype(str)
            active_incident_col = x_col if x_col in incident_cols else None
            df_plot = (
                _add_incident_category(df, active_incident_col)
                if active_incident_col
                else df.copy()
            )
            incident_color_map = (
                dict(
                    zip(
                        df_plot["_incident_label"].tolist(),
                        df_plot["_incident_color"].tolist(),
                    )
                )
                if active_incident_col and "_incident_label" in df_plot.columns
                else None
            )

            if color_col:
                df_plot_copy = df_plot.copy()
                color_incident_color_map = None

                if color_col in incident_cols:
                    df_plot_copy = _add_incident_category(df_plot_copy, color_col)
                    color_incident_color_map = dict(
                        zip(
                            df_plot_copy["_incident_label"].tolist(),
                            df_plot_copy["_incident_color"].tolist(),
                        )
                    )
                
                if not color_incident_color_map:
                    color_incident_color_map = incident_color_map

                if active_incident_col is not None:
                    x_col = "_incident_label"
                if color_col in incident_cols:
                    color_col = "_incident_label"

                if view_horizontal_stacked:
                    df_plot_copy[color_col] = df_plot_copy[color_col].astype(str)

                if df_plot_copy.duplicated(subset=[x_col, color_col]).any():
                    df_plot_copy[y_col] = (
                        df_plot_copy.groupby([x_col, color_col])[y_col]
                        .transform("sum")
                    )
                    df_plot_copy = df_plot_copy.drop_duplicates(subset=[x_col, color_col])

                bar_kwargs = dict(
                    x=x_col, y=y_col, color=color_col, barmode="group", template=tmpl, text=y_col
                )
                category_orders = {}
                if active_incident_col is not None:
                    bar_kwargs["x"] = "_incident_label"
                    color_incident_color_map = incident_color_map
                    bar_kwargs["color_discrete_map"] = color_incident_color_map
                if (
                    color_incident_color_map 
                    and color_col in incident_cols
                ) or (
                    color_incident_color_map
                    and color_col in ["_incident_label"]
                ):
                    bar_kwargs["x"] = (
                        "_incident_label" if active_incident_col is not None else x_col
                    )
                    bar_kwargs["color"] = "_incident_label"
                    bar_kwargs["color_discrete_map"] = color_incident_color_map
                    category_orders["_incident_label"] = INICIDENT_CATEGORY_NAMES
                x_order = get_categorical_columns(df_plot_copy, bar_kwargs["x"])
                category_orders[bar_kwargs["x"]] = x_order
                if category_orders:
                    bar_kwargs["category_orders"] = category_orders
                fig = px.bar(df_plot_copy, **bar_kwargs)
                group_x = bar_kwargs["x"]
                group_y = bar_kwargs["y"]
                group_color = bar_kwargs.get("color")
                x_values = category_orders.get(group_x, df_plot_copy[group_x].drop_duplicates().tolist())

                if group_color:
                    total_base = (
                        df_plot_copy[[group_x, group_color, group_y]]
                        .drop_duplicates()
                    )
                else:
                    total_base = (
                        df_plot_copy[[group_x, group_y]]
                        .drop_duplicates()
                    )

                group_totals = (
                    total_base.groupby(group_x)[group_y]
                    .sum()
                    .reset_index()
                )
                totals_map = dict(
                    zip(group_totals[group_x], group_totals[group_y])
                )

                for x_val in x_values:
                    if x_val in totals_map:
                        fig.add_annotation(
                            x=x_val,
                            y=totals_map[x_val],
                            text=f"{totals_map[x_val]}",
                            showarrow=False,
                            yshift=15,
                            font=dict(size=14),
                            xanchor="center"
                        )
            elif incident_color_map and active_incident_col is not None:
                df_plot_copy = df_plot.copy()
                if df_plot[x_col].duplicated().any(): 
                    df_plot_copy[y_col] = df_plot_copy.groupby("_incident_label")[y_col].transform("sum")
                    df_plot_copy.drop_duplicates(subset=["_incident_label"], inplace=True)
                fig = px.bar(
                    df_plot_copy,
                    x="_incident_label",
                    y=y_col,
                    color="_incident_label",
                    template=tmpl,
                    color_discrete_map=incident_color_map,
                    text=y_col,
                    category_orders={"_incident_label": INICIDENT_CATEGORY_NAMES}
                )
            else:
                df_plot_copy = df_plot.copy()
                if df_plot[x_col].duplicated().any(): 
                    df_plot_copy[y_col] = df_plot_copy.groupby(x_col)[y_col].transform("sum")
                    df_plot_copy.drop_duplicates(subset=[x_col], inplace=True)
                fig = px.bar(df_plot_copy, x=x_col, y=y_col, template=tmpl, text=y_col)
            fig.update_traces(textposition="outside")
            if view_all_labels:
                fig.update_xaxes(
                    type="category",
                    tickmode="array",
                    tickvals=get_categorical_columns(df_plot, x_col)
                )
            if active_incident_col is not None:
                fig.update_xaxes(
                    categoryorder="array",
                    categoryarray=INICIDENT_CATEGORY_NAMES
                )

        # ───── LINE ─────
        elif chart_type == "line":
            active_incident_col = x_col if x_col in incident_cols else None
            df_plot = (
                _add_incident_category(df, active_incident_col)
                if active_incident_col
                else df.copy()
            )
            incident_color_map = (
                dict(
                    zip(
                        df_plot["_incident_label"].tolist(),
                        df_plot["_incident_color"].tolist(),
                    )
                )
                if active_incident_col and "_incident_label" in df_plot.columns
                else None
            )

            if color_col:
                line_kwargs = dict(
                    x=x_col, y=y_col, color=color_col, markers=True, template=tmpl
                )
                if incident_color_map and color_col in incident_cols:
                    line_kwargs["x"] = (
                        "_incident_label" if active_incident_col is not None else x_col
                    )
                    line_kwargs["color"] = "_incident_label"
                    line_kwargs["color_discrete_map"] = incident_color_map
                fig = px.line(df_plot, **line_kwargs)
            elif incident_color_map and active_incident_col is not None:
                fig = px.line(
                    df_plot,
                    x="_incident_label",
                    y=y_col,
                    color="_incident_label",
                    markers=True,
                    template=tmpl,
                    color_discrete_map=incident_color_map,
                )
            else:
                fig = px.line(df_plot, x=x_col, y=y_col, markers=True, template=tmpl)

        # ───── PIE ─────
        elif chart_type == "pie":
            active_incident_col = x_col if x_col in incident_cols else None
            df_plot = (
                _add_incident_category(df, active_incident_col)
                if active_incident_col
                else df.copy()
            )
            incident_color_map = (
                dict(
                    zip(
                        df_plot["_incident_label"].tolist(),
                        df_plot["_incident_color"].tolist(),
                    )
                )
                if active_incident_col and "_incident_label" in df_plot.columns
                else None
            )

            if y_col:
                df_plot[y_col] = pd.to_numeric(df_plot[y_col], errors="coerce").fillna(0)

            if incident_color_map and active_incident_col is not None:
                fig = px.pie(
                    df_plot,
                    names="_incident_label",
                    values=y_col,
                    hole=0.35,
                    template=tmpl,
                    color="_incident_label",
                    color_discrete_map=incident_color_map,
                )
            else:
                fig = px.pie(
                    df_plot, names=x_col, values=y_col, hole=0.35, template=tmpl
                )
            fig.update_traces(textinfo="percent+label")

        # ───── FALLBACK ─────
        else:
            fig = px.bar(
                df_plot,
                x=x_col,
                y=y_col,
                template=tmpl,
                title=f"Chart type '{chart_type_raw}' not recognized",
            )
        fig.update_layout(
            title={
                "text": title.replace("_", " ").capitalize(),
                "x": 0.5,
                "xanchor": "center",
                "font": {
                    "size": 24
                }
            }
        )

    except Exception as e:
        st.warning(f"Could not render `{chart_type}` chart: {e}")
        return

    if fig:
        st.plotly_chart(
            fig,
            use_container_width=True,
            config={"displayModeBar": False},
            key=f"{uuid.uuid4()}_plot",
        )


def render_crosstab(df: pd.DataFrame):
    if df.empty:
        st.info("No data to summarise.")
        return

    numeric = df.select_dtypes(include="number").columns.tolist()
    cat = df.select_dtypes(exclude="number").columns.tolist()

    try:
        if len(cat) >= 2 and len(numeric) >= 1:
            pivot = df.pivot_table(
                index=cat[0],
                columns=cat[1],
                values=numeric[0],
                aggfunc="sum",
                fill_value=0,
            )
            st.caption(f"Crosstab — {cat[0]} x {cat[1]} (sum of {numeric[0]})")
            st.dataframe(pivot, use_container_width=True)

        elif len(cat) == 1 and len(numeric) >= 1:
            summary = df.groupby(cat[0])[numeric].agg(["sum", "mean", "count"])
            summary.columns = [f"{v}_{f}" for v, f in summary.columns]
            summary = summary.reset_index()
            st.caption(f"Summary — grouped by {cat[0]}")
            st.dataframe(summary, use_container_width=True, hide_index=True)

        elif len(numeric) >= 2:
            corr = df[numeric].corr().round(3)
            st.caption("Correlation Matrix")
            st.dataframe(
                corr.style.background_gradient(cmap="Blues", axis=None),
                use_container_width=True,
            )

        else:
            desc = df.describe(include="all").T.reset_index()
            desc.rename(columns={"index": "column"}, inplace=True)
            st.caption("Statistical Summary")
            st.dataframe(desc, use_container_width=True, hide_index=True)

    except Exception as e:
        st.warning(f"Could not build crosstab: {e}")
        st.dataframe(df.describe(include="all").T, use_container_width=True)


# ── Controller singleton ─────────────────────────────────────────────────────
@st.cache_resource
def get_controller():
    return DataExtractorController()


controller = get_controller()

# ── Session state ────────────────────────────────────────────────────────────
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "total_queries" not in st.session_state:
    st.session_state.total_queries = 0
if "successful_queries" not in st.session_state:
    st.session_state.successful_queries = 0


# ── Helpers ──────────────────────────────────────────────────────────────────
def build_message_history() -> list[Message]:
    return [
        Message(role=msg["role"], content=msg["content"])
        for msg in st.session_state.chat_history
    ]


def call_controller(user_query: str):
    uq = UserQuery(user_query=user_query)
    history = build_message_history()
    response = asyncio.run(controller.extrcat(user_query=uq, message_history=history))
    return response


def render_message(msg):
    is_user = msg["role"] == "user"
    avatar = "👤" if is_user else "🤖"
    role = "user" if is_user else "assistant"

    with st.chat_message(role, avatar=avatar):
        # Status badge
        if "status" in msg and not is_user:
            if msg["status"] == "success":
                st.success("Query executed successfully", icon="✅")
            else:
                st.error("Query failed", icon="❌")

        st.markdown(msg["content"])

        # SQL block
        if msg.get("sql"):
            with st.expander("Generated SQL", expanded=False):
                st.code(msg["sql"], language="sql")

        # Timestamp
        if msg.get("ts"):
            st.caption(msg["ts"])

    # ── Multi-view data panel ──────────────────────────────────────────────
    data = msg.get("data", [])
    chart_hint = msg.get("best_suitable_chart")

    if data and len(data) > 0 and msg.get("status") == "success":
        df = pd.DataFrame(data)

        tab_table, tab_crosstab, tab_chart = st.tabs(
            ["📋 Table", "📐 Crosstab / Summary", "📊 Chart"]
        )

        with tab_table:
            st.dataframe(df, use_container_width=True, hide_index=True)

        with tab_crosstab:
            render_crosstab(df)

        with tab_chart:
            if chart_hint:
                icon_map = {"BAR": "📊", "PIE": "🥧", "LINE": "📈"}
                icon = icon_map.get(str(chart_hint).upper(), "📊")
                st.caption(f"{icon} {chart_hint} (AI suggested)")
            else:
                st.caption("📊 Auto-detected chart type")
            charts = [
                "BAR",
                "LINE",
                "PIE"
            ]
            default_option = chart_hint if chart_hint in charts else "BAR"
            default_index = charts.index(default_option)
            charts[default_index] = f"{charts[default_index]} (AI suggested)"

            chart_hint = st.selectbox(
                "Select chart type",
                options=charts,
                index=default_index,
                key=f"chart_type_{msg.get('ts', 'x')}",
            )
            chart_hint = chart_hint.replace(" (AI suggested)", "") if chart_hint else None

            chart_key = f"chart_{msg.get('ts', 'x').replace(':', '_')}"
            incident_cols = []
            if msg.get("incident_col"):
                incident_cols.append(msg["incident_col"])
            detected = _detect_incident_col(df)
            if detected and detected not in incident_cols:
                incident_cols.extend(detected)

            for col in incident_cols:
                if col in df.columns:
                    df[col] = df[col].astype("category")
            render_chart(df, incident_cols, chart_hint, key_prefix=chart_key)

    elif msg.get("status") == "success" and not data:
        st.info("Query returned 0 rows.")


# ── Sidebar ───────────────────────────────────────────────────────────────────
with st.sidebar:
    st.header("📊 Session Stats")
    col1, col2 = st.columns(2)
    with col1:
        st.metric("Queries", st.session_state.total_queries)
    with col2:
        st.metric("Success", st.session_state.successful_queries)

    st.divider()
    st.header("💡 Example Prompts")
    examples = [
        "List all fire incidents in last 10 years",
        "Show top 10 incidents by type",
        "Count incidents per year",
        "Find incidents with alarm time after 6pm",
        "List unique incident types",
    ]
    for ex in examples:
        if st.button(ex, use_container_width=True, key=f"ex_{ex[:20]}"):
            st.session_state["prefill"] = ex

    st.divider()
    if st.button("🗑 Clear History", use_container_width=True):
        st.session_state.chat_history = []
        st.session_state.total_queries = 0
        st.session_state.successful_queries = 0
        st.rerun()

    if st.session_state.chat_history:
        with st.expander("🔍 Raw Message History"):
            st.json(build_message_history())


# ── Main layout ───────────────────────────────────────────────────────────────
st.title("Firerms Data Extractor Chatbot")
st.caption("Natural language → SQL → Results")

# ── Chat area ─────────────────────────────────────────────────────────────────
chat_container = st.container()
with chat_container:
    if not st.session_state.chat_history:
        st.markdown("---")
        col1, col2, col3 = st.columns([1, 2, 1])
        with col2:
            st.markdown(
                "<div style='text-align:center;padding:40px 0;'>"
                "<p style='font-size:3rem;'>🔍</p>"
                "<h3>Ask anything about your data</h3>"
                "<p>Type a natural language question and the AI will generate SQL and return results.</p>"
                "</div>",
                unsafe_allow_html=True,
            )
    else:
        for msg in st.session_state.chat_history:
            render_message(msg)

# ── Chat input ────────────────────────────────────────────────────────────────
prefill = st.session_state.pop("prefill", "")
prompt = st.chat_input("Ask a question about your data…", key="chat_input")
if not prompt and prefill:
    prompt = prefill

if prompt:
    ts_now = datetime.now().strftime("%H:%M:%S")

    st.session_state.chat_history.append(
        {"role": "user", "content": prompt, "ts": ts_now}
    )
    st.session_state.total_queries += 1

    with st.spinner("Generating SQL and fetching results…"):
        try:
            result = call_controller(prompt)
            status = result.status
            sql = result.sql_query
            data = result.data or []

            try:
                best_chart = result.output.best_suitable_chart.value
            except Exception:
                best_chart = None

            incident_col = None
            if result.output.is_incident_category_required:
                incident_col = result.output.column_name_mapped_with_incident_category

            row_count = len(data)
            content = (
                f"Query executed successfully. Returned **{row_count}** row(s)."
                if status == "success"
                else f"Query returned status: `{status}`."
            )

            st.session_state.chat_history.append(
                {
                    "role": "assistant",
                    "content": content,
                    "sql": sql,
                    "data": data,
                    "status": status,
                    "best_suitable_chart": best_chart,
                    "incident_col": incident_col,
                    "ts": datetime.now().strftime("%H:%M:%S"),
                }
            )

            if status == "success":
                st.session_state.successful_queries += 1

        except Exception as e:
            st.session_state.chat_history.append(
                {
                    "role": "assistant",
                    "content": f"Error: {str(e)}",
                    "status": "error",
                    "ts": datetime.now().strftime("%H:%M:%S"),
                }
            )

    st.rerun()