import streamlit as st
from utils.api import get_available_datasets, load_datasets

_LABEL_MAP = {
    "pubmedqa":             "PubMedQA",
    "mental_health":        "Mental Health Counseling",
    "mediqa":               "Medical MediQA",
    "medqa_usmle":          "MedQA-USMLE",
}


def render_dataset_panel() -> None:
    st.markdown("### Knowledge Base")

    try:
        datasets = get_available_datasets()
    except Exception:
        st.caption("API offline — start the server to manage datasets.")
        return

    if not datasets:
        st.caption("No datasets available.")
        return

    search = st.text_input(
        "search_kb",
        placeholder="Search datasets…",
        label_visibility="collapsed",
        key="kb_search",
    )

    filtered = [
        ds for ds in datasets
        if not search or search.lower() in ds["name"].lower() or search.lower() in ds["description"].lower()
    ]

    if not filtered:
        st.caption("No datasets match your search.")
        return

    st.caption("Select datasets to load into the knowledge base:")

    selected = [
        ds["name"]
        for ds in filtered
        if st.checkbox(
            _LABEL_MAP.get(ds["name"], ds["name"]),
            key=f"ds_{ds['name']}",
            help=ds["description"],
        )
    ]

    max_samples = st.slider(
        "Max samples per dataset",
        min_value=100,
        max_value=2000,
        value=500,
        step=100,
        help="Higher values improve coverage but take longer to embed.",
    )

    if st.button("Load Selected", disabled=not selected, use_container_width=True):
        with st.spinner(f"Embedding {len(selected)} dataset(s)… this may take a few minutes"):
            try:
                response = load_datasets(selected, max_samples)
                for name, result in response.get("results", {}).items():
                    label = _LABEL_MAP.get(name, name)
                    if result["status"] == "success":
                        st.success(f"{label}: {result['records_upserted']} records indexed")
                    elif result["status"] == "skipped":
                        st.warning(f"{label}: skipped — {result.get('detail', '')}")
                    else:
                        st.error(f"{label}: {result.get('detail', 'failed')}")
            except Exception as exc:
                st.error(f"Failed to load datasets: {exc}")