# app.py  —  BizIntel AI Ultra  
# Supports: CSV/Excel/DB ingestion, date+metric plotting, ARIMA forecasting,
# safe Plotly writes, Gemini 1.5 Pro strategy, KPI cards, optional EDA.

import os
import tempfile
from typing import Literal

import pandas as pd
import streamlit as st
import google.generativeai as genai
import plotly.graph_objects as go

# ──────────────────────────────────────────────────────────────
# 0) Redirect all Plotly write_image() into a writable temp dir
# ──────────────────────────────────────────────────────────────
TMP = tempfile.gettempdir()
_original_write = go.Figure.write_image

def _safe_write(self, path, *args, **kwargs):
    filename = os.path.basename(path)
    safe_path = os.path.join(TMP, filename)
    return _original_write(self, safe_path, *args, **kwargs)

go.Figure.write_image = _safe_write

# ──────────────────────────────────────────────────────────────
# 1) Imports for tools & DB connector
# ──────────────────────────────────────────────────────────────
from tools.csv_parser      import parse_csv_tool
from tools.plot_generator  import plot_metric_tool      # (csv_path, date_col, metric_col) → Figure or error
from tools.forecaster      import forecast_metric_tool  # (csv_path, date_col, metric_col) → text + /tmp/forecast_plot.png
from tools.visuals         import histogram_tool, scatter_matrix_tool, corr_heatmap_tool
from db_connector          import fetch_data_from_db, list_tables, SUPPORTED_ENGINES

# ──────────────────────────────────────────────────────────────
# 2) Gemini 1.5 Pro initialization
# ──────────────────────────────────────────────────────────────
genai.configure(api_key=os.getenv("GEMINI_APIKEY"))
gemini = genai.GenerativeModel(
    "gemini-1.5-pro-latest",
    generation_config={
        "temperature": 0.7,
        "top_p": 0.9,
        "response_mime_type": "text/plain",
    },
)

# ──────────────────────────────────────────────────────────────
# 3) Streamlit page setup
# ──────────────────────────────────────────────────────────────
st.set_page_config(page_title="BizIntel AI Ultra", layout="wide")
st.title("📊 BizIntel AI Ultra – Advanced Analytics + Gemini 1.5 Pro")
TEMP_DIR = tempfile.gettempdir()

# ──────────────────────────────────────────────────────────────
# 4) Data source selection: CSV/Excel or SQL Database
# ──────────────────────────────────────────────────────────────
source = st.radio("Select data source", ["Upload CSV / Excel", "Connect to SQL Database"])
csv_path = None

if source == "Upload CSV / Excel":
    upload = st.file_uploader("Upload CSV or Excel (≤ 500 MB)", type=["csv","xlsx","xls"])
    if upload:
        tmp_file = os.path.join(TEMP_DIR, upload.name)
        with open(tmp_file, "wb") as f:
            f.write(upload.read())

        if upload.name.lower().endswith(".csv"):
            csv_path = tmp_file
        else:
            try:
                df_xl = pd.read_excel(tmp_file, sheet_name=0)
                csv_path = os.path.splitext(tmp_file)[0] + ".csv"
                df_xl.to_csv(csv_path, index=False)
            except Exception as e:
                st.error(f"Excel parsing failed: {e}")
                st.stop()
        st.success(f"{upload.name} saved ✅")

else:
    engine = st.selectbox("DB engine", SUPPORTED_ENGINES)
    conn   = st.text_input("SQLAlchemy connection string")
    if conn:
        try:
            tables = list_tables(conn)
            tbl    = st.selectbox("Table", tables)
            if st.button("Fetch table"):
                csv_path = fetch_data_from_db(conn, tbl)
                st.success(f"Fetched **{tbl}** as CSV ✅")
        except Exception as e:
            st.error(f"Connection failed: {e}")
            st.stop()

if not csv_path:
    st.stop()

# Download the working CSV
with open(csv_path, "rb") as f:
    st.download_button("⬇️ Download working CSV", f, file_name=os.path.basename(csv_path))

# ──────────────────────────────────────────────────────────────
# 5) Show first rows & let user pick date + numeric metric
# ──────────────────────────────────────────────────────────────
df_head = pd.read_csv(csv_path, nrows=5)
st.dataframe(df_head)

date_col     = st.selectbox("Select date/time column", df_head.columns)
numeric_cols = df_head.select_dtypes("number").columns.tolist()
metric_col   = st.selectbox("Select numeric metric column", numeric_cols)

# ──────────────────────────────────────────────────────────────
# 6) Local analysis: summary, trend chart, forecast
# ──────────────────────────────────────────────────────────────
with st.spinner("Parsing dataset…"):
    summary_text = parse_csv_tool(csv_path)

with st.spinner("Generating trend chart…"):
    trend_fig = plot_metric_tool(csv_path, date_col, metric_col)
if isinstance(trend_fig, go.Figure):
    st.subheader("📈 Trend")
    st.plotly_chart(trend_fig, use_container_width=True)
else:
    st.warning(trend_fig)

with st.spinner("Running forecast…"):
    forecast_text = forecast_metric_tool(csv_path, date_col, metric_col)

st.subheader(f"🔮 {metric_col} Forecast")
forecast_png = os.path.join(TEMP_DIR, "forecast_plot.png")
if os.path.exists(forecast_png):
    st.image(forecast_png, use_container_width=True)
else:
    st.warning("Forecast image not found.")

# ──────────────────────────────────────────────────────────────
# 7) Gemini-driven strategy recommendations
# ──────────────────────────────────────────────────────────────
prompt = (
    f"You are **BizIntel Strategist AI**.\n\n"
    f"### Dataset Summary\n```\n{summary_text}\n```\n\n"
    f"### {metric_col} Forecast\n```\n{forecast_text}\n```\n\n"
    "Return **Markdown** with:\n"
    "1. Five key insights\n"
    "2. Three actionable strategies\n"
    "3. Risk factors or anomalies\n"
    "4. Suggested additional visuals\n"
)

st.subheader("🚀 Strategy Recommendations (Gemini 1.5 Pro)")
with st.spinner("Generating insights…"):
    strategy_md = gemini.generate_content(prompt).text
st.markdown(strategy_md)
st.download_button("⬇️ Download Strategy (.md)", strategy_md, file_name="strategy.md")

# ──────────────────────────────────────────────────────────────
# 8) KPI cards + detailed Stats
# ──────────────────────────────────────────────────────────────
full_df    = pd.read_csv(csv_path, low_memory=False)
total_rows = len(full_df)
num_cols   = full_df.shape[1]
missing_pct = full_df.isna().mean().mean() * 100

st.markdown("---")
st.subheader("📑 Dataset Overview")
c1, c2, c3 = st.columns(3)
c1.metric("Rows",     f"{total_rows:,}")
c2.metric("Columns",  str(num_cols))
c3.metric("Missing %", f"{missing_pct:.1f}%")

with st.expander("🔎 Detailed descriptive statistics"):
    stats_df = full_df.describe().T.reset_index().rename(columns={"index":"Feature"})
    st.dataframe(stats_df.style.format(precision=2).background_gradient(cmap="Blues"),
                 use_container_width=True)

# ──────────────────────────────────────────────────────────────
# 9) Optional Exploratory Visuals
# ──────────────────────────────────────────────────────────────
st.markdown("---")
st.subheader("🔍 Optional Exploratory Visuals")

if st.checkbox("Histogram"):
    hcol = st.selectbox("Variable", numeric_cols, key="hist")
    st.plotly_chart(histogram_tool(csv_path, hcol), use_container_width=True)

if st.checkbox("Scatter-matrix"):
    sel = st.multiselect("Choose columns", numeric_cols, default=numeric_cols[:3])
    if sel:
        st.plotly_chart(scatter_matrix_tool(csv_path, sel), use_container_width=True)

if st.checkbox("Correlation heat-map"):
    st.plotly_chart(corr_heatmap_tool(csv_path), use_container_width=True)