import os
import io
import time
import base64

import gradio as gr
import pandas as pd
from PIL import Image
from gradio_client import Client

# ------------------------- Constants -------------------------
TABLEAU_URL = "https://public.tableau.com/views/InsuranceDashboard_17677520784850/ExecutiveOverviewDashbiard?:showVizHome=no&:embed=y&:toolbar=no"


BACKEND_SPACE_ID = (os.getenv("BACKEND_SPACE_ID") or "").strip()
HF_TOKEN = (os.getenv("HF_TOKEN") or "").strip()

QUICK_QUESTIONS = [
    "Loss ratio by product_type",
    "Loss ratio by region",
    "Claim frequency by product_type",
    "Average claim severity",
    "Top 10 agents by premium",
]
# ---------------------Initialize once---------------------------------
client = Client(BACKEND_SPACE_ID, token=HF_TOKEN)

# ------------------------- Backend helpers -------------------------
def decode_chart(chart_b64: str):
    """Decode a base64 PNG/JPEG string returned by the backend into a PIL image."""
    if not chart_b64:
        return None

    raw = chart_b64
    if "," in raw and raw.strip().startswith("data:image"):
        raw = raw.split(",", 1)[1]

    try:
        img_bytes = base64.b64decode(raw)
        return Image.open(io.BytesIO(img_bytes))
    except Exception:
        return None


def call_backend(query: str) -> tuple[pd.DataFrame, object, str]:
    """
    Expected backend response format:
    {
      "data": [ { ...row... }, ... ],
      "chart": "<base64-encoded image or data URL>",
      "message": "optional summary"
    }
    """   
    response = client.predict(
    query,
    api_name="/query"
)

    if not isinstance(response, dict):
        raise ValueError("Backend returned an unexpected response format.")

    df_res = pd.DataFrame(response.get("data", []))
    chart_img = decode_chart(response.get("chart"))
    message = response.get("message", "Results generated.")

    return df_res, chart_img, message


# ------------------------- UI + Streaming -------------------------
def _js_autoscroll():
    return """
<script>
(function(){
  const t = document.getElementById("chat_thread");
  if (t) t.scrollTop = t.scrollHeight;
})();
</script>
"""


def _set_interactive(flag: bool):
    return gr.update(interactive=flag)


def _thinking_steps(query: str):
    base = f"Query: {query}\nThinking"
    dots = ["", ".", "..", "..."]
    for i in range(4):
        yield f"{base}{dots[i]}"
        time.sleep(0.10)

    steps = [
        "• Calling backend",
        "• Fetching results",
        "• Rendering results",
    ]

    acc = f"{base}...\n"
    for s in steps:
        acc = acc + s + "\n"
        yield acc.rstrip()
        time.sleep(0.15)


def _to_messages(history_pairs: list[list[str]]):
    msgs = []
    for pair in history_pairs or []:
        if not pair:
            continue
        user = pair[0] if len(pair) > 0 else ""
        assistant = pair[1] if len(pair) > 1 else ""
        if user:
            msgs.append({"role": "user", "content": str(user)})
        if assistant:
            msgs.append({"role": "assistant", "content": str(assistant)})
    return msgs


def run_query(user_text: str, history_pairs: list):
    history_pairs = history_pairs or []
    q = (user_text or "").strip()

    if not q:
        yield (
            gr.update(value=""),
            history_pairs,
            gr.update(value=_to_messages(history_pairs)),
            gr.update(value=pd.DataFrame(), visible=False),
            gr.update(value=None, visible=False),
            gr.update(selected=0),
            _set_interactive(False),
            *[_set_interactive(True) for _ in range(len(QUICK_QUESTIONS))],
            _js_autoscroll(),
        )
        return

    disabled_send = _set_interactive(False)
    disabled_chips = [_set_interactive(False) for _ in range(len(QUICK_QUESTIONS))]

    for frame in _thinking_steps(q):
        yield (
            gr.update(value=frame),
            history_pairs,
            gr.update(value=_to_messages(history_pairs)),
            gr.update(value=pd.DataFrame(), visible=False),
            gr.update(value=None, visible=False),
            gr.update(),
            disabled_send,
            *disabled_chips,
            _js_autoscroll(),
        )

    try:
        df_res, chart_img, assistant_msg = call_backend(q)
    except Exception as e:
        df_res = pd.DataFrame({"error": [str(e)]})
        chart_img = None
        assistant_msg = f"Backend request failed: {e}"

    history_pairs = history_pairs + [[q, assistant_msg]]

    if (df_res is None or df_res.empty) and assistant_msg:
        df_res = pd.DataFrame({"message": [assistant_msg]})
    show_df = df_res is not None and not df_res.empty
    show_fig = chart_img is not None
    tab_selected = 1 if show_fig else 0

    enabled_chips = [_set_interactive(True) for _ in range(len(QUICK_QUESTIONS))]
    send_after = _set_interactive(False)

    yield (
        gr.update(value=""),
        history_pairs,
        gr.update(value=_to_messages(history_pairs)),
        gr.update(value=(df_res if df_res is not None else pd.DataFrame()), visible=show_df),
        gr.update(value=chart_img, visible=show_fig),
        gr.update(selected=tab_selected),
        send_after,
        *enabled_chips,
        _js_autoscroll(),
    )


def chip_run(chip_text: str, history_pairs: list):
    yield from run_query(chip_text, history_pairs)


def clear_all():
    empty_pairs = []
    return (
        gr.update(value=""),
        empty_pairs,
        gr.update(value=_to_messages(empty_pairs)),
        gr.update(value=pd.DataFrame(), visible=False),
        gr.update(value=None, visible=False),
        gr.update(selected=0),
        gr.update(interactive=False),
        *[gr.update(interactive=True) for _ in range(len(QUICK_QUESTIONS))],
        _js_autoscroll(),
    )


def send_enabled(text: str):
    ok = bool((text or "").strip())
    return gr.update(interactive=ok)


# ------------------------- CSS -------------------------
CSS = """
:root{
  --bg: #f6f7fb;
  --card: #ffffff;
  --muted: #64748b;
  --border: rgba(15, 23, 42, 0.10);
  --shadow: 0 10px 30px rgba(15, 23, 42, 0.08);
  --radius: 16px;
}

body{ background: var(--bg) !important; }

#app_wrap{
  height: 100vh;
  max-width: 1600px;
  margin: 0 auto;
  padding: 14px;
  box-sizing: border-box;
}
.split_row{ height: calc(100vh - 28px); gap: 14px; }

#left_tableau{
  height: calc(100vh - 28px);
  border: 1px solid var(--border);
  border-radius: var(--radius);
  overflow: hidden;
  background: var(--card);
  box-shadow: var(--shadow);
}
#left_tableau iframe{ width: 100%; height: calc(100vh - 28px); border: 0; }

#right_chat{
  height: calc(100vh - 28px);
  border: 1px solid var(--border);
  border-radius: var(--radius);
  background: var(--card);
  box-shadow: var(--shadow);
  display: flex;
  flex-direction: column;
  overflow: visible;
}

.copilot_header{
  padding: 12px 14px;
  border-bottom: 1px solid rgba(15, 23, 42, 0.08);
  background: linear-gradient(180deg, #ffffff, #fbfbfe);
  font-weight: 700;
}

/* ONLY scroll container on right */
#chat_thread{
  flex: 1 1 auto;
  min-height: 0;
  overflow-y: auto;
  overflow-x: hidden;
  padding: 14px 14px 120px 14px;
  background: var(--bg);
  box-sizing: border-box;
}

/* REMOVE chat message bubbles + toolbar */
#chat_thread .message,
#chat_thread .message-wrap,
#chat_thread .message-row,
#chat_thread .bubble,
#chat_thread .toolbar,
#chat_thread .icon-row {
  display: none !important;
}

/* Collapse leftover Chatbot spacing */
#chat_thread .chatbot,
#chat_thread .chatbot > div {
  margin: 0 !important;
  padding: 0 !important;
  height: 0 !important;
}

/* Dock */
#dock{
  position: sticky;
  bottom: 0;
  z-index: 50;
  padding: 12px 14px 12px 14px;
  background: linear-gradient(180deg, rgba(246,247,251,0.95), rgba(246,247,251,1));
  border-top: 1px solid rgba(15, 23, 42, 0.10);
}

#composer_card{
  display: flex;
  gap: 10px;
  align-items: flex-end;
  padding: 10px 10px 10px 12px;
  border-radius: 18px;
  background: #ffffff;
  border: 1px solid rgba(15, 23, 42, 0.12);
  box-shadow: 0 10px 24px rgba(15, 23, 42, 0.08);
}

#composer{ flex: 1 1 auto; }
#composer .wrap{ padding: 0 !important; }
#composer textarea{
  width: 100% !important;
  border: none !important;
  outline: none !important;
  background: transparent !important;
  font-size: 14px !important;
  line-height: 1.35 !important;
  min-height: 56px !important;
  max-height: 160px !important;
  resize: none !important;
  padding: 6px 6px 6px 2px !important;
}

#composer_actions{
  display: flex;
  gap: 8px;
  align-items: center;
  justify-content: flex-end;
  flex: 0 0 auto;
}
#send_btn button{
  height: 38px !important;
  padding: 0 14px !important;
  border-radius: 999px !important;
  border: 1px solid #111827 !important;
  background: #111827 !important;
  color: #fff !important;
  font-weight: 800 !important;
}
#clear_btn button{
  height: 38px !important;
  padding: 0 12px !important;
  border-radius: 999px !important;
  border: 1px solid rgba(15, 23, 42, 0.18) !important;
  background: #ffffff !important;
  color: #111827 !important;
  font-weight: 800 !important;
}
#clear_btn button:hover{ background: #f8fafc !important; }

#quick_chips .chips-row > div,
#quick_chips .chips-row > .gradio-column,
#quick_chips .chips-row .gradio-column{
  flex: 0 0 auto !important;
  width: auto !important;
  min-width: 0 !important;
}

#chips{ margin-top: 10px; }
.chips-wrap{ display: grid; gap: 8px; }
.chips-title{
  font-size: 12px;
  font-weight: 700;
  color: var(--muted);
  letter-spacing: 0.02em;
}
.chips-row{
  display: flex;
  flex-wrap: wrap;
  gap: 8px;
  align-items: center;
}
.chip{
  appearance: none;
  border: 1px solid rgba(15, 23, 42, 0.16);
  background: #ffffff;
  border-radius: 999px;
  padding: 7px 11px;
  font-size: 12px;
  font-weight: 700;
  cursor: pointer;
  box-shadow: 0 2px 10px rgba(15, 23, 42, 0.06);
}
.chip:hover{ background: #f8fafc; }

.card{
  background: var(--card);
  border: 1px solid var(--border);
  border-radius: var(--radius);
  box-shadow: 0 6px 18px rgba(15, 23, 42, 0.05);
  padding: 12px;
}

#result_chart{
  min-height: 360px;
  width: 100%;
}

#output_tabs{ margin-top: 4px !important; }
"""

# ------------------------- App -------------------------
with gr.Blocks(title="Converse AI", css=CSS) as demo:
    history_state = gr.State([])

    with gr.Column(elem_id="app_wrap"):
        with gr.Row(elem_classes=["split_row"]):
            with gr.Column(scale=2, min_width=650):
                gr.HTML(
                    f"""
                    <div id="left_tableau">
                      <iframe src="{TABLEAU_URL}" allowfullscreen></iframe>
                    </div>
                    """
                )

            with gr.Column(scale=1, min_width=420):
                with gr.Group(elem_id="right_chat"):
                    with gr.Column(elem_id="chat_thread"):
                        chat_ui = gr.Chatbot(
                            value=[],
                            show_label=False,
                            label=None,
                            height=None,
                        )

                        with gr.Tabs(elem_id="output_tabs") as output_tabs:
                            with gr.Tab("Results", id=0):
                                result_df = gr.Dataframe(
                                    value=pd.DataFrame(),
                                    show_label=False,
                                    interactive=False,
                                    visible=False,
                                    wrap=True,
                                    elem_classes=["card"],
                                )
                            with gr.Tab("Chart", id=1):
                                result_chart = gr.Image(
                                    value=None,
                                    show_label=False,
                                    visible=False,
                                    elem_id="result_chart",
                                    elem_classes=["card"],
                                )

                        js = gr.HTML(value=_js_autoscroll())

                    with gr.Column(elem_id="dock"):
                        with gr.Row(elem_id="composer_card", equal_height=False):
                            composer = gr.Textbox(
                                placeholder="Ask a question…",
                                show_label=False,
                                lines=2,
                                max_lines=8,
                                elem_id="composer",
                                container=False,
                            )
                            with gr.Row(elem_id="composer_actions", equal_height=False):
                                send_btn = gr.Button("Send", elem_id="send_btn", interactive=False)
                                clear_btn = gr.Button("Clear", elem_id="clear_btn", variant="secondary")

                        with gr.Column(elem_id="quick_chips"):
                            gr.HTML(
                                """
                                <div class="chips-wrap">
                                  <div class="chips-title">Quick questions</div>
                                </div>
                                """
                            )
                            with gr.Row(elem_classes=["chips-row"]):
                                chip_buttons = []
                                for q in QUICK_QUESTIONS:
                                    chip_buttons.append(gr.Button(q, variant="secondary", elem_classes=["chip"]))

    composer.change(send_enabled, inputs=[composer], outputs=[send_btn])

    send_btn.click(
        run_query,
        inputs=[composer, history_state],
        outputs=[
            composer,
            history_state,
            chat_ui,
            result_df,
            result_chart,
            output_tabs,
            send_btn,
            *chip_buttons,
            js,
        ],
    )
    composer.submit(
        run_query,
        inputs=[composer, history_state],
        outputs=[
            composer,
            history_state,
            chat_ui,
            result_df,
            result_chart,
            output_tabs,
            send_btn,
            *chip_buttons,
            js,
        ],
    )

    for btn, q in zip(chip_buttons, QUICK_QUESTIONS):
        btn.click(
            chip_run,
            inputs=[gr.State(q), history_state],
            outputs=[
                composer,
                history_state,
                chat_ui,
                result_df,
                result_chart,
                output_tabs,
                send_btn,
                *chip_buttons,
                js,
            ],
        )

    clear_btn.click(
        clear_all,
        outputs=[
            composer,
            history_state,
            chat_ui,
            result_df,
            result_chart,
            output_tabs,
            send_btn,
            *chip_buttons,
            js,
        ],
    )

if __name__ == "__main__":
    demo.queue().launch(ssr_mode=False)