import gradio as gr
from app.main import app as fastapi_app
from app.retriever import RETRIEVAL_MODES, store

# Ensure indexes are loaded
store.load()


def search_fn(question: str, mode: str, top_k: int):
    """Gradio handler for full hybrid retrieval."""
    if not question.strip():
        return {"error": "Please enter a valid search query."}

    result = store.answer(question, top_k=int(top_k), mode=mode)
    return result


# Build Gradio UI
with gr.Blocks(title="Calibrated Hybrid Retrieval API") as demo:
    gr.Markdown(
        """
    # 🔬 Calibrated Entropy-Weighted Hybrid Retrieval API
    A research-oriented hybrid retrieval system combining BM25 sparse search, FAISS dense vector search, 
    corpus-level CDF score calibration, entropy-weighted adaptive fusion, and cross-encoder reranking.
    """
    )

    with gr.Row():
        with gr.Column(scale=2):
            query_input = gr.Textbox(
                lines=2,
                placeholder="Enter your question (e.g., '0-dimensional biomaterials show inductive properties.')...",
                label="Search Question / Claim",
            )
            mode_dropdown = gr.Dropdown(
                choices=list(RETRIEVAL_MODES.keys()),
                value="hybrid_calibrated_rerank",
                label="Retrieval Mode",
            )
            top_k_slider = gr.Slider(
                minimum=1, maximum=20, value=3, step=1, label="Top K Results"
            )
            search_btn = gr.Button("⚡ Run Retrieval Pipeline", variant="primary")

        with gr.Column(scale=3):
            output_json = gr.JSON(label="Retrieval Results & Telemetry Output")

    search_btn.click(
        fn=search_fn,
        inputs=[query_input, mode_dropdown, top_k_slider],
        outputs=output_json,
    )

# Mount FastAPI app into Gradio so REST API endpoints (/docs, /query, /health) are also live!
app = gr.mount_gradio_app(fastapi_app, demo, path="/ui")

if __name__ == "__main__":
    demo.launch()