from __future__ import annotations

import streamlit as st

from quering.core import (
    build_output_zip_bytes,
    infer_query_prefix,
    load_markdown_zip_bytes,
    load_pooled_zip_bytes,
    query_pooled_records,
    supported_models,
)

st.set_page_config(
    page_title="quering",
    page_icon="🔎",
    layout="wide",
)

st.title("quering")
st.caption(
    "Query pooled embeddings against your original markdown chunks so you can pick the best context for an LLM."
)

pooled_uploaded = st.file_uploader("Upload pooled embedding zip", type=["zip"], key="pooled")
markdown_uploaded = st.file_uploader("Upload markdown chunk zip", type=["zip"], key="markdown")

if pooled_uploaded is None or markdown_uploaded is None:
    st.info("Upload both the pooled embedding zip and the markdown chunk zip to continue.")
    st.stop()

pooled_bytes = pooled_uploaded.read()
markdown_bytes = markdown_uploaded.read()

try:
    pooled_records, pooled_summary = load_pooled_zip_bytes(pooled_bytes)
    markdown_chunks = load_markdown_zip_bytes(markdown_bytes)
except Exception as exc:
    st.error(str(exc))
    st.stop()

source_model_name = pooled_summary.get("source_model_name") or pooled_summary.get("model_name") or "unknown"

with st.sidebar:
    st.header("Settings")
    st.write("Source model from pooled zip:")
    st.code(str(source_model_name))

    model_options = supported_models()
    labels = [f"{name} — {desc}" for name, desc in model_options]
    names = [name for name, _ in model_options]
    default_index = names.index(source_model_name) if source_model_name in names else 0

    use_source_model = st.checkbox("Use source model from pooled zip", value=source_model_name in names)
    if use_source_model and source_model_name not in names:
        st.warning("Source model is not in the built-in list. Enable custom model below if you want to use it.")

    selected_label = st.selectbox("Embedding model", labels, index=default_index)
    selected_model = model_options[labels.index(selected_label)][0]

    use_custom = st.checkbox("Use a custom Hugging Face / SentenceTransformers model name", value=use_source_model and source_model_name not in names)
    custom_default = str(source_model_name) if (use_source_model and source_model_name not in names) else ""
    custom_model = st.text_input("Custom model name", value=custom_default) if use_custom else ""

    model_name = custom_model.strip() or (source_model_name if use_source_model and source_model_name in names else selected_model)
    if use_source_model and custom_model.strip():
        model_name = custom_model.strip()
    elif use_source_model and source_model_name in names:
        model_name = source_model_name

    prefix_mode = st.selectbox("Query prefix", ["auto", "none", "custom"], index=0)
    if prefix_mode == "auto":
        query_prefix = infer_query_prefix(model_name)
    elif prefix_mode == "custom":
        query_prefix = st.text_input("Custom query prefix", value=infer_query_prefix(model_name))
    else:
        query_prefix = ""

    top_k = st.slider("Top matches", min_value=1, max_value=min(20, len(pooled_records)), value=min(5, len(pooled_records)))
    normalize_query = st.checkbox("Normalize query embedding", value=True)
    include_matched_chunks = st.checkbox("Include matched markdown chunks in export zip", value=True)
    formats = st.multiselect(
        "Extra output formats inside export zip",
        options=["jsonl", "npz"],
        default=["jsonl", "npz"],
        help="CSV and prompt text are always included.",
    )

query_text = st.text_area("Query sentence", height=120, placeholder="Ask one sentence, for example: What is RagRetriever?")

col_a, col_b, col_c = st.columns(3)
col_a.metric("Pooled candidates", len(pooled_records))
col_b.metric("Markdown chunks", len(markdown_chunks))
col_c.metric("Source model", source_model_name)

if source_model_name and model_name and source_model_name != model_name:
    st.warning(
        "The selected query embedding model differs from the source model used to create the pooled embeddings. "
        "Similarity scores may not be meaningful unless both sides come from the same vector space."
    )

if st.button("Run query", type="primary"):
    with st.spinner("Embedding query and ranking pooled chunks..."):
        try:
            manifest, results, summary = query_pooled_records(
                pooled_records,
                markdown_chunks,
                query_text=query_text,
                model_name=model_name,
                top_k=int(top_k),
                query_prefix=query_prefix,
                normalize_query=normalize_query,
            )
            export_bytes = build_output_zip_bytes(
                manifest,
                results,
                summary,
                pooled_zip_name=pooled_uploaded.name,
                markdown_zip_name=markdown_uploaded.name,
                formats=formats,
                include_matched_chunks=include_matched_chunks,
            )
        except Exception as exc:
            st.error(str(exc))
            st.stop()

    left, right = st.columns([1, 2])
    with left:
        st.metric("Top matches returned", summary["top_k"])
        st.metric("Best score", f"{summary['best_score']:.4f}" if summary["best_score"] is not None else "n/a")
        st.metric("Query vector dim", summary["query_vector_dim"])
        st.download_button(
            label="Download quering .zip",
            data=export_bytes,
            file_name=f"{pooled_uploaded.name.rsplit('.', 1)[0]}_querying.zip",
            mime="application/zip",
            use_container_width=True,
        )
    with right:
        st.subheader("Top results")
        st.dataframe(manifest, use_container_width=True)

    with st.expander("LLM prompt preview", expanded=True):
        from quering.core import build_llm_prompt
        st.code(build_llm_prompt(query_text, results), language="text")

    with st.expander("Matched chunk previews", expanded=False):
        for r in results:
            st.markdown(f"### Rank {r.rank} — {r.source_markdown_file or r.group_key} — score {r.score:.4f}")
            if r.text:
                st.code(r.text[:1200], language="markdown")
            else:
                st.info("No matching markdown chunk text was found for this pooled group key.")
