from main import rag_pipeline
import gradio as gr
import html


municipalities = [
    "Città di Agrigento",
    "Comun General de Fascia",
    "Comune di Capo D'Orlando",
    "Comune di Casatenovo",
    "Comune di Fonte Nuova",
    "Comune di Galliate",
    "Comune di Grottaferrata",
    "Comune di Gubbio",
    "Comune di Ispica",
    "Comune di Torre Santa Susanna",
    "Comune di Santa Maria Capua Vetere"
]

def check_muni(muni: str) -> str:
    return "" if muni in (None, "", "Tutti") else muni

def _run_rag(raw, muni, feedback):
    """Internal helper: runs your pipeline and returns everything needed."""
    out = rag_pipeline(raw_query=raw, municipality=check_muni(muni), feedback=feedback)
    return {
        "answer": out["final_answer"],
        "cot": out["chain_of_thought"],
        "chunks_html": build_chunks_html(out["retrieved_chunks"]),
        "rewritten": out["rewritten_query"],
        "intent": out["intent"]
    }

def build_chunks_html(chunks):
    html_blocks = []
    for i, c in enumerate(chunks, 1):
        text = html.escape(c["chunk_text"])
        meta = {
            "Document ID": c.get("document_id","N/A"),
            "Municipality": c.get("municipality","N/A"),
            "Section": c.get("section","N/A"),
            "Page": c.get("page","N/A"),
            "Score": f"{c.get('final_score',0):.4f}"
        }
        meta_lines = "".join(f"<li><b>{k}</b>: {v}</li>" for k,v in meta.items())
        html_blocks.append(f"""
          <div style="margin-bottom:1em;">
            <h4>Chunk {i}</h4>
            <ul>{meta_lines}</ul>
            <p style="white-space: pre-wrap;">{text}</p>
          </div><hr/>
        """)
    return "\n".join(html_blocks) or "<i>No chunks retrieved.</i>"


# Stage 1: initial answer
def stage1(query, municipality):
    result = _run_rag(query, municipality, feedback="")
    return (
        result["answer"],
        result["cot"],
        result["chunks_html"],
        result["rewritten"],
        result["intent"],
        gr.update(visible=True),   # show satisfaction
        gr.update(visible=False),  # hide feedback
        gr.update(visible=False)   # hide refine
    )

# When satisfaction changes: show feedback only if “No”
def on_satisfied_change(sat):
    if sat == "👎 No":
        # show feedback box, show refine button (but keep it disabled)
        return gr.update(visible=True), gr.update(visible=True, interactive=False)
    else:
        # hide both if they say "👍 Sì"
        return gr.update(visible=False), gr.update(visible=False, interactive=False)


# Stage 2: enable refine button only once there's feedback text
def on_feedback_change(feedback_text):
    return gr.update(interactive=bool(feedback_text.strip()))

# Stage 3: refinement pass
def stage3(query, municipality, satisfaction, feedback):
    result = _run_rag(query, municipality, feedback)
    return (
        result["answer"],
        result["cot"],
        result["chunks_html"],
        result["rewritten"],
        result["intent"],
        gr.update(visible=True), # satisfaction
        gr.update(visible=True), # feedback box
        gr.update(visible=False, interactive=False)  # disable refine button
    )

# == Gradio Blocks ==
with gr.Blocks() as demo:
    gr.Markdown("## DSPy RAG Demo")

    with gr.Row():
        query_input = gr.Textbox(label="Domanda", placeholder="Scrivi la tua domanda qui...")
        muni_input = gr.Dropdown(
            choices=["Tutti"] + municipalities,
            value="Tutti",
            label="Comune (facoltativo)"
        )

    run_btn = gr.Button("Cerca")

    ans_out = gr.Textbox(label="Risposta Finale", lines=3)  # answer container

    # CoT container inside its accordion
    with gr.Accordion("Catena di Ragionamento", open=False):
        cot_txt = gr.Textbox(label="", interactive=False, lines=6)

    # chunks HTML container inside its accordion
    with gr.Accordion("Chunks Recuperati con Metadata", open=False):
        chunks_html = gr.HTML("<i>No data yet.</i>")

    # Show what the model actually used
    rewritten_out = gr.Textbox(label="Domanda Rielaborata", interactive=False)
    intent_out = gr.Textbox(label="Intento Identificato", interactive=False)

    # Human feedback step
    satisfaction = gr.Radio(
        ["👍 Sì","👎 No"], label="Sei soddisfatto?", visible=False
    )
    feedback_txt = gr.Textbox(
        label="Cosa non va? (facoltativo)", lines=3, visible=False, interactive=True
    )
    refine_btn = gr.Button("Affina la Domanda", visible=False, interactive=False)

  # Wire Stage 1
    run_btn.click(
        fn=stage1,
        inputs=[query_input, muni_input],
        outputs=[
            ans_out, cot_txt, chunks_html,
            rewritten_out, intent_out,
            satisfaction, feedback_txt, refine_btn
        ]
    )

    # When satisfaction toggles, show/hide feedback
    satisfaction.change(
        fn=on_satisfied_change,
        inputs=[satisfaction],
        outputs=[feedback_txt, refine_btn]
    )

    # Only enable refine once feedback is non-empty
    feedback_txt.change(
        fn=on_feedback_change,
        inputs=[feedback_txt],
        outputs=[refine_btn]
    )

    # Wire Stage 3 (refinement)
    refine_btn.click(
        fn=stage3,
        inputs=[query_input, muni_input, satisfaction, feedback_txt],
        outputs=[
            ans_out,
            cot_txt,
            chunks_html,
            rewritten_out,
            intent_out,
            satisfaction,   # now reset and visible
            feedback_txt,   # still visible
            refine_btn      # hidden/disabled
        ]
    )

demo.launch(ssr_mode=False)