frontend/app.py

from __future__ import annotations

import sys
from pathlib import Path
import pandas as pd
import gradio as gr

import matplotlib.pyplot as plt
# Make project importable
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from services.service import BIService

service = BIService()


# -----------------------------
# Helpers
# -----------------------------
def ensure_list(value):
    if isinstance(value, list):
        return [str(v) for v in value if str(v).strip()]
    if isinstance(value, str) and value.strip():
        return [value.strip()]
    return []


def get_indexed_documents():
    try:
        return service.list_documents()
    except Exception:
        return []


# -----------------------------
# File Upload + RAG Indexing
# -----------------------------
def index_documents(files):
    if not files:
        return "⚠️ No files uploaded"

    try:
        file_data = [(file.name, file.read()) for file in files]
        result = service.ingest_documents(file_data)

        messages = []
        for item in result:
            messages.append(f"{item.get('status')}: {item.get('message')}")

        return "\n".join(messages)

    except Exception as e:
        return f"❌ Error: {str(e)}"


# -----------------------------
# Main Query Function
# -----------------------------
def run_query(question, use_rag):
    if not question.strip():
        return "⚠️ Please enter a question", None, "", "", "", ""

    try:
        try:
            result = service.ask(question.strip(), use_rag=use_rag)
        except TypeError:
            result = service.ask(question.strip())

        sql_payload = result.get("sql", {})
        analysis = result.get("analysis", {})
        visualization = result.get("visualization", {})
        reflection = result.get("reflection", {})
        result_payload = sql_payload.get("result", {})
        dataframe = pd.DataFrame(result_payload.get("rows", []))
       
        fig = render_chart_gradio(dataframe, visualization)
        df = pd.DataFrame(result_payload.get("rows", []))

        summary = analysis.get("summary", "No summary")

        sql_query = sql_payload.get("sql", "No SQL generated")

        insights = "\n".join(ensure_list(analysis.get("insights", [])))
        followups = "\n".join(ensure_list(analysis.get("follow_ups", [])))

        risks = "\n".join(reflection.get("risks", [])) if reflection.get("risks") else ""

        return (
        analysis.get("summary", "No summary"),
        dataframe,
        sql_query,
        insights,
        fig,
        followups,
        risks,
        sql_payload.get("sql", ""),
        "\n".join(analysis.get("insights", []))
    )

    except Exception as e:
        return f"❌ Error: {str(e)}", None, "", "", "", ""

def validate_visualization(dataframe: pd.DataFrame, viz: dict) -> dict:
    if dataframe.empty:
        return {}

    columns = list(dataframe.columns)

    chart_type = viz.get("chart_type")
    x_axis = viz.get("x_axis")
    y_axis = viz.get("y_axis")

    # Fix missing fields
    if x_axis not in columns:
        x_axis = columns[0]

    if y_axis not in columns:
        y_axis = columns[1] if len(columns) > 1 else None

    # Auto-detect chart type if invalid
    if chart_type not in ["line", "bar", "scatter"]:
        if "date" in x_axis.lower() or "time" in x_axis.lower():
            chart_type = "line"
        elif dataframe[x_axis].dtype == "object":
            chart_type = "bar"
        else:
            chart_type = "scatter"

    return {
        "chart_type": chart_type,
        "x_axis": x_axis,
        "y_axis": y_axis,
    }
def render_chart_gradio(dataframe: pd.DataFrame, viz: dict):
    if dataframe.empty:
        return None

    viz = validate_visualization(dataframe, viz)

    chart_type = viz.get("chart_type")
    x = viz.get("x_axis")
    y = viz.get("y_axis")

    if not x or not y:
        return None

    fig, ax = plt.subplots()

    if chart_type == "line":
        ax.plot(dataframe[x], dataframe[y])

    elif chart_type == "bar":
        ax.bar(dataframe[x], dataframe[y])

    elif chart_type == "scatter":
        ax.scatter(dataframe[x], dataframe[y])

    ax.set_xlabel(x)
    ax.set_ylabel(y)
    ax.set_title(f"{chart_type.upper()} Chart")

    return fig
# -----------------------------
# UI Layout
# -----------------------------
with gr.Blocks(title="BI Agent") as demo:

    gr.Markdown("# 📊 Business Intelligence Agent")
    gr.Markdown("Ask database questions → SQL → Results → Insights 🚀")

    # Runtime Info
    runtime = service.runtime_status()
    gr.Markdown(
        f"""
        **Runtime Status**
        - Groq: `{runtime.get('has_api_key', False)}`
        - ClickHouse: `{runtime.get('has_database', False)}`
        - Chroma: `{runtime.get('has_chroma', False)}`
        - Model: `{runtime.get('groq_model', 'unknown')}`
        """
    )

    # -----------------------------
    # File Upload Section (Optional RAG)
    # -----------------------------
    with gr.Accordion("📄 Upload Documents (Optional RAG)", open=False):

        file_input = gr.File(
            file_count="multiple",
            file_types=[".pdf", ".docx", ".txt", ".md"]
        )

        upload_btn = gr.Button("Index Documents")
        upload_output = gr.Textbox(label="Indexing Status")

        upload_btn.click(
            index_documents,
            inputs=file_input,
            outputs=upload_output
        )

    # -----------------------------
    # Question Input
    # -----------------------------
    question_input = gr.Textbox(
        label="Ask your question",
        placeholder="Which popular but poorly rated products should we investigate?",
        lines=3
    )

    use_rag_checkbox = gr.Checkbox(
        label="Use uploaded documents (RAG)",
        value=False
    )

    run_btn = gr.Button("Run Analysis 🚀")

    # -----------------------------
    # Outputs
    # -----------------------------
    summary_output = gr.Textbox(label="Summary")
    table_output = gr.Dataframe(label="Result Table")
    sql_output = gr.Code(label="Generated SQL", language="sql")
    insights_output = gr.Textbox(label="Insights")
    visualization=gr.Plot(label="Visualization")
    followups_output = gr.Textbox(label="Follow-ups")
    risks_output = gr.Textbox(label="Risks")

    run_btn.click(
        run_query,
        inputs=[question_input, use_rag_checkbox],
        outputs=[
            summary_output,
            table_output,
            sql_output,
            insights_output,
            visualization,
            followups_output,
            risks_output
        ]
    )

# Launch
if __name__ == "__main__":
    demo.launch(share=True)