import gradio as gr
import os
from analyst import DataAnalystAgent
from analyst_tools import (
    copy_uploaded_file,
    cleanup_session_sandbox,
    normalize_message_text,
    collect_charts,
)
# ── Helpers ────────────────────────────────────────────────────────────────────
def _init_agent():
    agent = DataAnalystAgent()
    agent.setup()
    return agent
def _show_status(text: str) -> str:
    """Return a small HTML status badge."""
    return (
        f'<div class="status-badge">'
        f'<span class="pulse"></span> {text}</div>'
    )
def _status_done(text: str) -> str:
    return f'<div class="status-badge done">{text}</div>'
def process_message(agent, message, success_criteria, dataset_file, history):
    """Run the agent and return updated UI state with HTML report + notebook."""
    if not message.strip():
        return (
            history, agent, gr.update(),
            gr.update(visible=False),
            gr.update(value="", visible=False),
            gr.update(visible=False),
            gr.update(value=None, visible=False),
            gr.update(value=""),
        )
    dataset_filename = None
    if dataset_file is not None:
        original_name = os.path.basename(dataset_file)
        copy_uploaded_file(dataset_file, original_name, agent.session_dir)
        dataset_filename = original_name
    updated_history, evaluator_feedback, _, nb_path, html_report = agent.run(
        message, success_criteria, dataset_filename, history,
    )
    evaluator_feedback = normalize_message_text(evaluator_feedback)
    # Charts gallery
    charts = collect_charts(agent.session_dir)
    chart_update = gr.update(value=charts, visible=True) if charts else gr.update(visible=False)
    # HTML report rendered in-browser
    report_update = gr.update(value=html_report, visible=True) if html_report else gr.update(visible=False)
    # Notebook download
    nb_update = gr.update(value=nb_path, visible=True) if nb_path else gr.update(value=None, visible=False)
    # Status line
    chart_count = len(charts)
    status_parts = ["Analysis complete"]
    if chart_count:
        status_parts.append(f"{chart_count} chart{'s' if chart_count != 1 else ''}")
    if nb_path:
        status_parts.append("notebook ready")
    status_text = _status_done(" | ".join(status_parts))
    return (
        updated_history,
        agent,
        gr.update(value=""),
        chart_update,
        gr.update(value=evaluator_feedback, visible=bool(evaluator_feedback)),
        report_update,
        nb_update,
        gr.update(value=status_text),
    )
def reset_agent(agent):
    cleanup_session_sandbox(agent.session_dir)
    new_agent = agent.reset()
    return (
        [],
        new_agent,
        gr.update(value=""),
        gr.update(value=""),
        gr.update(value=None),
        gr.update(visible=False),
        gr.update(value="", visible=False),
        gr.update(visible=False),
        gr.update(value=None, visible=False),
        gr.update(value=""),
    )
# ── CSS ────────────────────────────────────────────────────────────────────────
CSS = """
/* ── Page background ─────────────────────────────────────────────────────── */
body {
    background:
        radial-gradient(ellipse at 10% 0%, rgba(59,130,246,.18), transparent 52%),
        radial-gradient(ellipse at 90% 0%, rgba(14,165,233,.14), transparent 42%),
        linear-gradient(180deg, #f5faff 0%, #eaf4ff 100%);
    font-family: "Manrope", system-ui, sans-serif;
}
.gradio-container { max-width: 1200px !important; }
/* ── Header ──────────────────────────────────────────────────────────────── */
#app-header {
    text-align: center;
    padding: 1.2rem 0 .6rem;
}
#app-header h1 {
    font-size: 1.8rem;
    font-family: "Space Grotesk", "Manrope", sans-serif;
    background: linear-gradient(135deg, #0f3fb8 0%, #2563eb 48%, #0ea5e9 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    background-clip: text;
    margin: 0;
    letter-spacing: -0.02em;
    font-weight: 700;
}
#app-header p {
    font-family: "Manrope", system-ui, sans-serif;
    color: #64748b;
    font-size: .92rem;
    margin: .3rem 0 0;
}
/* ── Cards ───────────────────────────────────────────────────────────────── */
.card {
    background: rgba(255,255,255,.92);
    backdrop-filter: blur(12px);
    border: 1px solid rgba(59,130,246,.12);
    border-radius: 16px;
    box-shadow: 0 8px 28px rgba(37,99,235,.08);
    padding: 1rem;
}
/* ── Chatbot ─────────────────────────────────────────────────────────────── */
#chatbot {
    border-radius: 16px !important;
    border: 1px solid rgba(59,130,246,.14) !important;
    box-shadow: 0 10px 34px rgba(37,99,235,.10) !important;
    background: linear-gradient(180deg, #ffffff 0%, #f2f8ff 100%) !important;
}
/* ── Input area ──────────────────────────────────────────────────────────── */
#input-group {
    background: rgba(255,255,255,.88);
    backdrop-filter: blur(10px);
    border: 1px solid rgba(59,130,246,.12);
    border-radius: 16px;
    box-shadow: 0 8px 24px rgba(37,99,235,.07);
    padding: .8rem 1rem;
}
/* ── Buttons ─────────────────────────────────────────────────────────────── */
#go-btn {
    min-width: 140px;
    background: linear-gradient(135deg, #1d4ed8 0%, #2563eb 58%, #0ea5e9 100%) !important;
    color: #fff !important;
    border: none !important;
    border-radius: 12px !important;
    font-weight: 600 !important;
    font-size: .95rem !important;
    box-shadow: 0 8px 20px rgba(37,99,235,.28) !important;
    transition: transform .15s, box-shadow .15s;
}
#go-btn:hover {
    transform: translateY(-1px);
    box-shadow: 0 10px 26px rgba(37,99,235,.34) !important;
}
#reset-btn {
    border: 1px solid rgba(59,130,246,.20) !important;
    background: rgba(239,246,255,.92) !important;
    color: #1d4ed8 !important;
    border-radius: 12px !important;
    font-weight: 500 !important;
    transition: background .15s;
}
#reset-btn:hover { background: rgba(219,234,254,.96) !important; }
/* ── Tabs ────────────────────────────────────────────────────────────────── */
.gr-tab-nav button {
    font-family: "Space Grotesk", "Manrope", sans-serif !important;
    font-weight: 700 !important;
    font-size: .98rem !important;
    letter-spacing: .01em !important;
    color: #335c8a !important;
    background: rgba(255,255,255,.9) !important;
    border: 1px solid rgba(59,130,246,.10) !important;
    border-radius: 10px 10px 0 0 !important;
    padding-top: .55rem !important;
    padding-bottom: .55rem !important;
    transition: color .15s, transform .15s;
}
.gr-tab-nav button:hover {
    color: #1d4ed8 !important;
    transform: translateY(-1px);
}
.gr-tab-nav button.selected {
    color: #1d4ed8 !important;
    background: #ffffff !important;
    border-bottom: 2px solid #1d4ed8 !important;
    box-shadow: 0 8px 20px rgba(37,99,235,.10) !important;
}
.gr-tab-nav button *, .gr-tab-nav button span, .gr-tab-nav button div {
    color: inherit !important;
}
/* ── Sidebar labels ──────────────────────────────────────────────────────── */
.sidebar-label {
    font-size: .82rem;
    font-weight: 700;
    text-transform: uppercase;
    letter-spacing: .06em;
    color: #2563eb;
    margin: .6rem 0 .3rem;
}
/* ── Status badge ────────────────────────────────────────────────────────── */
.status-badge {
    display: inline-flex;
    align-items: center;
    gap: 6px;
    font-size: .82rem;
    font-weight: 600;
    color: #1d4ed8;
    padding: 4px 12px;
    border-radius: 999px;
    background: rgba(59,130,246,.10);
}
.status-badge.done { color: #0369a1; background: rgba(14,165,233,.12); }
.pulse {
    width: 8px; height: 8px;
    border-radius: 50%;
    background: #2563eb;
    animation: pulse-anim 1.4s infinite;
}
@keyframes pulse-anim {
    0%   { opacity: 1; transform: scale(1); }
    50%  { opacity: .4; transform: scale(1.4); }
    100% { opacity: 1; transform: scale(1); }
}
/* ── Gallery ─────────────────────────────────────────────────────────────── */
.gr-gallery {
    border-radius: 12px !important;
    border: 1px solid rgba(59,130,246,.10) !important;
}
/* ── File upload / download ──────────────────────────────────────────────── */
.gr-file {
    border-radius: 12px !important;
    border: 1px dashed rgba(59,130,246,.24) !important;
    background: rgba(240,249,255,.95) !important;
}
.gr-file, .gr-file * {
    color: #1e293b !important;
}
.gr-button, .gr-button * {
    color: inherit !important;
}
.gr-download-button, .gr-download-button * {
    color: #1e293b !important;
}
.gr-markdown, .gr-markdown * {
    color: #1e293b;
}
/* ── Report embed ────────────────────────────────────────────────────────── */
#tab-report > div {
    width: 100% !important;
    max-width: 100% !important;
}
#tab-report .prose {
    max-width: 100% !important;
}
/* Force dark text inside the HTML report for readability */
#tab-report .rpt,
#tab-report .rpt * {
    color: #1e293b;
}
#tab-report .rpt h1, #tab-report .rpt h2 {
    color: #1d4ed8;
}
#tab-report .rpt .findings {
    background: #f3f8ff;
    border: 1px solid #cfe3ff;
}
/* ── Tab panel backgrounds ───────────────────────────────────────────────── */
.gr-tabitem {
    background: rgba(255,255,255,.72) !important;
    border-radius: 0 0 12px 12px;
}
/* ── Misc ────────────────────────────────────────────────────────────────── */
footer { display: none !important; }
.gr-examples .gr-samples-table { border-radius: 12px !important; }
"""
# ── Build UI ───────────────────────────────────────────────────────────────────
THEME = gr.themes.Soft(
    primary_hue=gr.themes.colors.indigo,
    secondary_hue=gr.themes.colors.amber,
    neutral_hue=gr.themes.colors.slate,
    font=[
        gr.themes.GoogleFont("Manrope"),
        gr.themes.GoogleFont("Space Grotesk"),
        "system-ui",
        "sans-serif",
    ],
)

with gr.Blocks(title="Data Analyst Agent") as ui:
    # ── Header ─────────────────────────────────────────────────────────────────
    gr.HTML(
        '<div id="app-header">'
        "<h1>Data Analyst Agent</h1>"
        "<p>Upload a CSV, ask a question — the agent writes &amp; runs Python, "
        "then delivers insights you can save as a notebook.</p>"
        "</div>"
    )
    agent_state = gr.State()
    # ── Top bar: dataset upload + status ───────────────────────────────────────
    with gr.Row(equal_height=True):
        with gr.Column(scale=3):
            dataset_upload = gr.File(
                label="Dataset (CSV)",
                file_types=[".csv"],
                type="filepath",
                elem_classes=["card"],
            )
        with gr.Column(scale=2):
            status_html = gr.HTML(value="", elem_id="status-bar")
    # ── Main content: left = chat, right = results tabs ────────────────────────
    with gr.Row(equal_height=False):
        # ── Left: conversation ─────────────────────────────────────────────────
        with gr.Column(scale=5, min_width=440):
            chatbot = gr.Chatbot(
                label="Conversation",
                elem_id="chatbot",
                height=500,
                avatar_images=(
                    None,
                    "https://api.dicebear.com/9.x/bottts/svg?seed=analyst",
                ),
            )
            # ── Input area ─────────────────────────────────────────────────────
            with gr.Group(elem_id="input-group"):
                message = gr.Textbox(
                    show_label=False,
                    placeholder="Ask a question about your data ...",
                    lines=1,
                    max_lines=4,
                    scale=4,
                )
                success_criteria = gr.Textbox(
                    show_label=False,
                    placeholder="Success criteria (optional) — e.g. Include correlation matrix and at least one chart",
                    lines=1,
                    max_lines=2,
                    scale=4,
                )
                with gr.Row():
                    reset_btn = gr.Button(
                        "Reset",
                        variant="stop",
                        elem_id="reset-btn",
                        size="sm",
                    )
                    go_btn = gr.Button(
                        "Analyse",
                        variant="primary",
                        elem_id="go-btn",
                        size="lg",
                    )
        # ── Right: results tabs ────────────────────────────────────────────────
        with gr.Column(scale=4, min_width=360):
            with gr.Tabs():
                # Tab 1 — Report
                with gr.Tab("Report", id="tab-report"):
                    report_html = gr.HTML(
                        label="Analysis Report",
                        visible=False,
                    )
                    gr.Markdown(
                        "<p style='text-align:center;color:#94a3b8;font-size:.85rem'>"
                        "Run an analysis to see the report here.</p>",
                        elem_id="report-placeholder",
                    )
                # Tab 2 — Charts
                with gr.Tab("Charts", id="tab-charts"):
                    chart_gallery = gr.Gallery(
                        label="Generated Charts",
                        columns=2,
                        rows=2,
                        visible=False,
                        height=400,
                        object_fit="contain",
                    )
                    gr.Markdown(
                        "<p style='text-align:center;color:#94a3b8;font-size:.85rem'>"
                        "Charts will appear here after analysis.</p>",
                        elem_id="charts-placeholder",
                    )
                # Tab 3 — Save / Download
                with gr.Tab("Save", id="tab-save"):
                    gr.HTML(
                        '<p class="sidebar-label">Download Notebook</p>'
                        '<p style="color:#64748b;font-size:.85rem">'
                        "The notebook (.ipynb) bundles your analysis code, "
                        "charts, and findings in one file.</p>"
                    )
                    notebook_download = gr.File(
                        label="Notebook (.ipynb)",
                        visible=False,
                        interactive=False,
                    )
                # Tab 4 — Evaluator
                with gr.Tab("Evaluator", id="tab-eval"):
                    gr.Markdown(
                        "<p style='color:#64748b;font-size:.85rem'>"
                        "Internal quality evaluator feedback.</p>"
                    )
                    evaluator_feedback = gr.Markdown(visible=False)
    # ── Example prompts ────────────────────────────────────────────────────────
    gr.Examples(
        examples=[
            [
                "What are the key trends in this dataset? Detect any outliers.",
                "Include at least one chart",
            ],
            [
                "Run a correlation analysis and identify the top 3 most correlated pairs.",
                "",
            ],
            [
                "Forecast the next 3 periods using a simple linear trend.",
                "",
            ],
        ],
        inputs=[message, success_criteria],
        label="Quick-start prompts",
    )
    # ── Events ─────────────────────────────────────────────────────────────────
    ui.load(_init_agent, [], [agent_state])
    shared_inputs = [agent_state, message, success_criteria, dataset_upload, chatbot]
    shared_outputs = [
        chatbot, agent_state, message, chart_gallery, evaluator_feedback,
        report_html, notebook_download, status_html,
    ]
    go_btn.click(process_message, shared_inputs, shared_outputs)
    message.submit(process_message, shared_inputs, shared_outputs)
    reset_btn.click(
        reset_agent,
        [agent_state],
        [
            chatbot, agent_state, message, success_criteria, dataset_upload,
            chart_gallery, evaluator_feedback, report_html, notebook_download,
            status_html,
        ],
    )
if __name__ == "__main__":
    ui.launch(theme=THEME, css=CSS)