pages/6_advanced_explorer.py

"""
Advanced Data Explorer — LLM-Powered Analytical Queries
Decomposes natural language into structured execution plans,
runs whitelisted operations, and presents results with exports.
"""
import json
import copy
import streamlit as st
import pandas as pd
from datetime import datetime

from core.db_connector import get_db_connector
from analysis.query_planner import QueryPlanner
from analysis.execution_engine import ExecutionEngine
from analysis.result_interpreter import ResultInterpreter
from analysis.operations import OperationOutput
from core.export_utils import package_multi_window_csv, estimate_csv_size
from ui.components import page_header, analysis_card, follow_up_section
from analysis.llm_analyzer import LLMCycleAnalyzer
from ui.plotly_charts import create_timeseries_chart


# ── Helper Functions (defined first for Streamlit top-to-bottom flow) ──


def _plan_to_serializable(plan: dict) -> dict:
    """Convert plan with datetimes to JSON-serializable dict."""
    p = copy.deepcopy(plan)
    dr = p.get("data_requirements", {})
    for key in ("start_time", "end_time"):
        if isinstance(dr.get(key), datetime):
            dr[key] = dr[key].isoformat()
    return p


def _render_dict_result(data: dict, op_result: OperationOutput):
    """Render a dict result (e.g., pump strokes, statistics)."""
    # Special rendering for pump stroke results
    if "total_revolutions" in data:
        cols = st.columns(4)
        cols[0].metric("Total Revolutions", f"{data.get('total_revolutions', 0):,}")
        cols[1].metric("Avg RPM", f"{data.get('avg_rpm', 0):.1f}")
        cols[2].metric("Peak RPM", f"{data.get('peak_rpm', 0):.1f}")
        cols[3].metric("Active Time", f"{data.get('active_minutes', 0):.1f} min")

        if data.get("note"):
            st.info(data["note"])
        return

    # Blocks result from find_peak_pressure_blocks
    if "blocks" in data:
        blocks = data["blocks"]
        if blocks:
            st.markdown(f"Found **{len(blocks)}** time blocks")
            block_df = pd.DataFrame(blocks)
            st.dataframe(block_df, use_container_width=True, hide_index=True)
        else:
            st.info("No blocks found matching the criteria.")
        return

    # Generic dict
    for key, val in data.items():
        st.markdown(f"**{key}:** {val}")


def _render_list_result(data: list, op_result: OperationOutput):
    """Render a list-of-dicts result (fills, periods, ramps)."""
    st.markdown(f"Found **{len(data)}** results")

    result_df = pd.DataFrame(data)

    # Format datetime columns for display
    for col in result_df.columns:
        if pd.api.types.is_datetime64_any_dtype(result_df[col]):
            result_df[col] = result_df[col].dt.strftime("%Y-%m-%d %H:%M:%S")

    st.dataframe(result_df, use_container_width=True, hide_index=True)


def _render_multi_df_result(data: list, op_result: OperationOutput):
    """Render multiple DataFrames (extract_windows)."""
    labels = op_result.metadata.get("labels", [f"Window {i+1}" for i in range(len(data))])
    rows = op_result.metadata.get("rows_per_window", [len(d) for d in data])

    st.markdown(f"Extracted **{len(data)}** windows")

    for i, (window_df, label) in enumerate(zip(data, labels)):
        with st.expander(f"{label} ({rows[i] if i < len(rows) else len(window_df):,} rows)"):
            tags_in_data = [c for c in window_df.columns if c != "timestamp" and pd.api.types.is_numeric_dtype(window_df[c])]
            if "timestamp" in window_df.columns and tags_in_data:
                fig = create_timeseries_chart(window_df, tags_in_data[:5], title=label)
                st.plotly_chart(fig, use_container_width=True)
            st.dataframe(window_df, use_container_width=True, height=200)

    # Multi-window zip download
    if data:
        zip_bytes = package_multi_window_csv(data, labels)
        st.download_button(
            label=f"Download All Windows (ZIP, {len(data)} files)",
            data=zip_bytes,
            file_name="csh2_windows_export.zip",
            mime="application/zip",
            key="v2_download_zip",
        )


def _render_df_result(data: pd.DataFrame, op_result: OperationOutput):
    """Render a single DataFrame result."""
    st.markdown(f"**{len(data):,}** rows")

    tags_in_data = [c for c in data.columns if c != "timestamp" and pd.api.types.is_numeric_dtype(data[c])]

    if "timestamp" in data.columns and tags_in_data:
        fig = create_timeseries_chart(data, tags_in_data[:5], title=op_result.label)
        st.plotly_chart(fig, use_container_width=True)

    with st.expander("Data Preview"):
        st.dataframe(data.head(500), use_container_width=True, height=300)


def _render_operation_result(op_result: OperationOutput, index: int):
    """Render a single operation's results."""
    st.subheader(f"Step {index + 1}: {op_result.label}")

    if op_result.metadata.get("error"):
        st.warning(f"Issue: {op_result.metadata['error']}")
        return

    data = op_result.data

    if isinstance(data, dict):
        _render_dict_result(data, op_result)
    elif isinstance(data, list) and data and isinstance(data[0], dict):
        _render_list_result(data, op_result)
    elif isinstance(data, list) and data and isinstance(data[0], pd.DataFrame):
        _render_multi_df_result(data, op_result)
    elif isinstance(data, pd.DataFrame):
        _render_df_result(data, op_result)

    # Metadata display
    meta = {k: v for k, v in op_result.metadata.items() if k != "error"}
    if meta:
        with st.expander("Details"):
            for k, v in meta.items():
                st.markdown(f"**{k}:** {v}")

    # Export button for this operation
    if op_result.exportable_df is not None and not op_result.exportable_df.empty:
        csv = op_result.exportable_df.to_csv(index=False)
        size = estimate_csv_size(op_result.exportable_df)
        safe_name = op_result.op_name.replace(" ", "_")
        st.download_button(
            label=f"Download {safe_name}.csv ({size})",
            data=csv,
            file_name=f"csh2_{safe_name}.csv",
            mime="text/csv",
            key=f"v2_download_{index}",
        )


# ── Page Header ──────────────────────────────────────────────────

page_header(
    "Advanced Data Explorer",
    "LLM-powered analytical queries with post-processing and export"
)

db = get_db_connector()
planner = QueryPlanner()


# ── Example Queries ──────────────────────────────────────────────

EXAMPLE_QUERIES = [
    "Find time blocks where PT130 exceeded 900 bar",
    "Find where both AOV140 and MC130_VFD_Speed were constant for 5+ minutes",
    "Export PT130, TT110, FT140 at 2Hz as CSV for Sep 25 10am-11am",
    "Identify fills where kWh per kg dispensed exceeded 0.6",
    "In Fill 3, what was the total number of pump strokes?",
    "Between Sep 24 15:00 and Sep 24 18:00, total pump strokes?",
]


# ── Session State Initialization ─────────────────────────────────

for key in ("v2_plan", "v2_result", "v2_query", "v2_run_plan", "v2_interpretation", "v2_interp_prompt"):
    if key not in st.session_state:
        st.session_state[key] = None
# v2_run_plan is a bool flag, default False
if st.session_state.v2_run_plan is None:
    st.session_state.v2_run_plan = False


# ── Query Input Section ──────────────────────────────────────────


def _set_example_query(text: str):
    """Callback: set the widget's own key so text appears on next render."""
    st.session_state.v2_query_input = text
    st.session_state.v2_query = text


st.subheader("Ask a Complex Question")

query_col, btn_col = st.columns([5, 1])
with query_col:
    user_query = st.text_input(
        "Describe your analysis",
        value=st.session_state.v2_query or "",
        placeholder='e.g. "Find time blocks where PT130 exceeded 900 bar"',
        key="v2_query_input",
        label_visibility="collapsed",
    )
    # Sync widget value back to logical query state
    st.session_state.v2_query = user_query
with btn_col:
    analyze_btn = st.button("Analyze", type="primary", use_container_width=True, key="v2_analyze_btn")

# Example query chips
st.caption("Try an example:")
example_cols = st.columns(3)
for i, example in enumerate(EXAMPLE_QUERIES):
    col = example_cols[i % 3]
    with col:
        st.button(
            example,
            key=f"v2_example_{i}",
            use_container_width=True,
            on_click=_set_example_query,
            args=(example,),
        )

st.divider()


# ── Planning Phase ───────────────────────────────────────────────

should_plan = (analyze_btn and user_query) or st.session_state.v2_run_plan
if should_plan:
    query_to_plan = user_query or st.session_state.v2_query or ""
    st.session_state.v2_run_plan = False  # Reset flag immediately
    st.session_state.v2_query = query_to_plan
    st.session_state.v2_plan = None
    st.session_state.v2_result = None
    st.session_state.v2_interpretation = None
    st.session_state.v2_interp_prompt = None
    st.session_state.pop("v2_followups", None)

    if not query_to_plan:
        st.warning("Please enter a query or click an example.")
    elif not planner.api_available:
        st.error("Claude API key not configured. Add ANTHROPIC_API_KEY to your .env or Streamlit secrets.")
    else:
        with st.spinner("Planning execution..."):
            all_sensors = db.get_all_tag_names()
            plan = planner.plan(query_to_plan, all_sensors)

        if plan and "error" in plan:
            st.error(f"Planning failed: {plan['error']}")
        elif plan:
            st.session_state.v2_plan = plan
        else:
            st.error("Could not create an execution plan. Try rephrasing your query.")


# ── Plan Review Card ─────────────────────────────────────────────

plan = st.session_state.v2_plan

if plan and "error" not in plan and st.session_state.v2_result is None:
    qt = plan.get("query_type", "FETCH")
    explanation = plan.get("explanation", "")
    dr = plan.get("data_requirements", {})
    ops = plan.get("operations", [])
    export_info = plan.get("export", {})

    # Query type badge
    badge_colors = {"FETCH": "#5BA3B5", "ANALYZE": "#D4A04A", "EXPORT": "#5B9A6E"}
    badge_color = badge_colors.get(qt, "#6B6358")

    st.markdown(f"""
    <div style="background: #1A1714; border: 1px solid #2A2520; border-radius: 8px; padding: 16px; margin-bottom: 16px;">
        <div style="display: flex; align-items: center; gap: 12px; margin-bottom: 12px;">
            <span style="background: {badge_color}; color: #0F0D0B; font-weight: 700; padding: 3px 10px;
                         border-radius: 4px; font-size: 0.75rem; letter-spacing: 1px;">{qt}</span>
            <span style="color: #B0A898; font-size: 0.9rem;">{explanation}</span>
        </div>
    </div>
    """, unsafe_allow_html=True)

    # Data requirements
    sensors = dr.get("sensors", [])
    resolution = dr.get("resolution", "auto")
    start_time = dr.get("start_time", "")
    end_time = dr.get("end_time", "")

    if isinstance(start_time, datetime):
        start_str = start_time.strftime("%b %d, %Y %H:%M")
    else:
        start_str = str(start_time)
    if isinstance(end_time, datetime):
        end_str = end_time.strftime("%b %d, %Y %H:%M")
    else:
        end_str = str(end_time)

    info_cols = st.columns(4)
    info_cols[0].markdown(f"**Sensors:** {', '.join(sensors[:6])}{'...' if len(sensors) > 6 else ''}")
    info_cols[1].markdown(f"**From:** {start_str}")
    info_cols[2].markdown(f"**To:** {end_str}")
    info_cols[3].markdown(f"**Resolution:** `{resolution}`")

    if ops:
        st.markdown("**Execution Steps:**")
        for i, op in enumerate(ops, 1):
            op_label = op.get("label", op.get("op", ""))
            st.markdown(f"&nbsp;&nbsp;`{i}.` {op_label}")

    if export_info.get("enabled"):
        st.markdown(f"**Export:** `{export_info.get('filename_hint', 'export')}.csv`")

    # Action buttons
    exec_col, cancel_col = st.columns([2, 1])
    with exec_col:
        execute_btn = st.button("Execute Plan", type="primary", use_container_width=True, key="v2_execute_btn")
    with cancel_col:
        cancel_btn = st.button("Cancel", use_container_width=True, key="v2_cancel_btn")

    with st.expander("Debug JSON"):
        display_plan = _plan_to_serializable(plan)
        st.code(json.dumps(display_plan, indent=2), language="json")

    if cancel_btn:
        st.session_state.v2_plan = None
        st.session_state.v2_result = None
        st.rerun()

    if execute_btn:
        engine = ExecutionEngine()
        with st.spinner("Executing plan..."):
            exec_result = engine.execute(plan)
        st.session_state.v2_result = exec_result
        st.rerun()


# ── Results Display ──────────────────────────────────────────────

result = st.session_state.v2_result

if result is not None:
    st.divider()

    if result.success:
        st.success(f"Execution completed in {result.execution_time_seconds:.1f}s")
    else:
        st.error(f"Execution failed ({result.execution_time_seconds:.1f}s)")
        for err in result.errors:
            st.error(err)

    # Warnings
    for warn in (result.warnings or []):
        st.warning(warn)

    # Summary metrics
    metric_cols = st.columns(4)
    metric_cols[0].metric("Rows Fetched", f"{result.row_count:,}")
    metric_cols[1].metric("Table", result.table_used or "N/A")
    metric_cols[2].metric("Operations", len(result.operation_results))
    metric_cols[3].metric("Time", f"{result.execution_time_seconds:.1f}s")

    # Explanation
    if result.explanation:
        analysis_card("Plan Summary", result.explanation)

    # Render each operation result
    for i, op_result in enumerate(result.operation_results):
        st.divider()
        _render_operation_result(op_result, i)

    # ── Murphy's Interpretation ──
    if result.success and result.operation_results:
        st.divider()
        if st.session_state.v2_interpretation:
            analysis_card("Murphy's Take", st.session_state.v2_interpretation)

            # Follow-up section for Murphy's Take
            interp_prompt = st.session_state.get("v2_interp_prompt", "")
            _llm_for_followup = LLMCycleAnalyzer()
            follow_up_section(
                session_key="v2_followups",
                llm_analyzer=_llm_for_followup,
                original_prompt=interp_prompt,
                original_analysis=st.session_state.v2_interpretation,
            )
        else:
            if st.button("Get Murphy's Take", type="primary", key="v2_interpret_btn"):
                interpreter = ResultInterpreter()
                if interpreter.api_available:
                    # Store the prompt for follow-up context
                    interp_prompt = interpreter._build_interpretation_prompt(
                        st.session_state.v2_query or "", result
                    )
                    st.session_state.v2_interp_prompt = interp_prompt

                    with st.spinner("Murphy is analyzing the results..."):
                        interpretation = interpreter.interpret(
                            st.session_state.v2_query or "", result
                        )
                    st.session_state.v2_interpretation = interpretation
                    st.rerun()
                else:
                    st.error("Claude API key not configured for interpretation.")

    # Source data overview
    if result.source_data is not None and not result.source_data.empty:
        st.divider()
        with st.expander("Source Data Overview"):
            src = result.source_data
            st.markdown(f"**{len(src):,}** rows, **{len(src.columns)}** columns")
            tags_in_src = [c for c in src.columns if c != "timestamp" and pd.api.types.is_numeric_dtype(src[c])]
            if "timestamp" in src.columns and tags_in_src:
                fig = create_timeseries_chart(src, tags_in_src[:5], title="Source Data")
                st.plotly_chart(fig, use_container_width=True)
            st.dataframe(src.head(200), use_container_width=True, height=200)


# ── Reset Button ─────────────────────────────────────────────────

if st.session_state.v2_plan or st.session_state.v2_result:
    st.divider()
    if st.button("New Query", use_container_width=True, key="v2_reset"):
        st.session_state.v2_plan = None
        st.session_state.v2_result = None
        st.session_state.v2_query = None
        st.session_state.v2_interpretation = None
        st.session_state.v2_interp_prompt = None
        st.session_state.pop("v2_followups", None)
        st.rerun()