import hashlib
import io

import numpy as np
import pandas as pd
import plotly.express as px
import streamlit as st
from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import cosine_similarity

# ----------------------------------------------------------------------------
# Page config
# ----------------------------------------------------------------------------
st.set_page_config(
    page_title="Document Intelligence Explorer",
    layout="wide",
)

st.title(
    "Document Intelligence Explorer",
    help=(
        "Upload a document and this tool will split it into chunks, embed each "
        "chunk with a sentence-transformer model, then show you two views of the "
        "same data side by side: an unsupervised K-Means clustering of the chunks, "
        "and a zero-shot classification of the chunks against labels you define. "
        "Use this to compare how the document organizes itself naturally versus "
        "how it maps onto categories you care about."
    ),
)
st.caption(
    "Upload a document, configure the parameters in the sidebar, then click "
    "Run Analysis to compare zero-shot classification against unsupervised "
    "semantic clustering on the same embedding space."
)

# ----------------------------------------------------------------------------
# Cached resource loaders
# ----------------------------------------------------------------------------
@st.cache_resource(show_spinner="Loading embedding model...")
def load_embedding_model():
    from sentence_transformers import SentenceTransformer
    return SentenceTransformer("all-MiniLM-L6-v2")


# ----------------------------------------------------------------------------
# Text extraction
# ----------------------------------------------------------------------------
def extract_text(uploaded_file) -> str:
    name = uploaded_file.name.lower()
    raw_bytes = uploaded_file.getvalue()

    if name.endswith(".txt"):
        return raw_bytes.decode("utf-8", errors="ignore")

    if name.endswith(".pdf"):
        from pypdf import PdfReader
        reader = PdfReader(io.BytesIO(raw_bytes))
        pages_text = []
        for page in reader.pages:
            try:
                pages_text.append(page.extract_text() or "")
            except Exception:
                pages_text.append("")
        return "\n".join(pages_text)

    raise ValueError("Unsupported file type. Please upload a .txt or .pdf file.")


def chunk_text(text: str, chunk_size: int, overlap: int = 50) -> list[str]:
    """Word-based chunking with light overlap for context continuity."""
    words = text.split()
    if not words:
        return []

    chunks = []
    step = max(chunk_size - overlap, 1)
    for start in range(0, len(words), step):
        chunk_words = words[start:start + chunk_size]
        if chunk_words:
            chunks.append(" ".join(chunk_words))
        if start + chunk_size >= len(words):
            break
    return chunks


class _DummyUpload:
    """Lightweight wrapper so cached function can reuse extract_text()."""
    def __init__(self, data: bytes, name: str):
        self._data = data
        self.name = name

    def getvalue(self):
        return self._data


# ----------------------------------------------------------------------------
# Cached pipeline functions
# ----------------------------------------------------------------------------
@st.cache_data(show_spinner="Generating embeddings...")
def compute_chunks_and_embeddings(file_bytes: bytes, file_name: str, chunk_size: int):
    file_hash = hashlib.md5(file_bytes).hexdigest()  # noqa: S324 - cache key only
    _ = file_hash

    text = extract_text(_DummyUpload(file_bytes, file_name))
    chunks = chunk_text(text, chunk_size=chunk_size)

    if not chunks:
        return [], np.empty((0, 384))

    model = load_embedding_model()
    embeddings = model.encode(chunks, show_progress_bar=False, normalize_embeddings=True)
    return chunks, np.array(embeddings)


@st.cache_data(show_spinner="Computing 2D projection (t-SNE)...")
def compute_2d_projection(embeddings: np.ndarray, seed: int = 42):
    n_samples = embeddings.shape[0]
    if n_samples < 2:
        return np.zeros((n_samples, 2))

    from sklearn.manifold import TSNE
    perplexity = max(2, min(30, n_samples - 1))
    reducer = TSNE(n_components=2, perplexity=perplexity, random_state=seed, init="pca")
    return reducer.fit_transform(embeddings)


@st.cache_data(show_spinner="Running K-Means clustering...")
def compute_clusters(embeddings: np.ndarray, k: int, seed: int = 42):
    k = min(k, max(1, embeddings.shape[0]))
    model = KMeans(n_clusters=k, random_state=seed, n_init="auto")
    labels = model.fit_predict(embeddings)
    return labels


@st.cache_data(show_spinner="Classifying chunks against labels...")
def compute_classification(_embeddings_key: str, embeddings: np.ndarray, labels_tuple: tuple):
    model = load_embedding_model()
    label_embeddings = model.encode(list(labels_tuple), normalize_embeddings=True)
    sims = cosine_similarity(embeddings, label_embeddings)
    predicted_idx = sims.argmax(axis=1)
    predicted_labels = [labels_tuple[i] for i in predicted_idx]
    confidence = sims.max(axis=1)
    return predicted_labels, confidence


# ----------------------------------------------------------------------------
# Session state init
# ----------------------------------------------------------------------------
if "result_df" not in st.session_state:
    st.session_state.result_df = None
if "chunks" not in st.session_state:
    st.session_state.chunks = None
if "has_labels" not in st.session_state:
    st.session_state.has_labels = False

# ----------------------------------------------------------------------------
# Sidebar controls
# ----------------------------------------------------------------------------
with st.sidebar:
    st.header("Configuration")

    uploaded_file = st.file_uploader(
        "Upload a document", type=["txt", "pdf"], accept_multiple_files=False
    )

    st.subheader(
        "Chunking",
        help=(
            "The document is split into smaller pieces of text (chunks) before "
            "embedding. Larger chunks give each piece more context but reduce the "
            "number of data points available for clustering and classification."
        ),
    )
    chunk_size = st.slider(
        "Chunk size (words)", min_value=50, max_value=500, value=150, step=10
    )

    st.subheader(
        "Clustering",
        help=(
            "K-Means groups chunks into K clusters based on embedding similarity, "
            "with no labels involved. Increasing K produces more, finer-grained "
            "groups; decreasing it produces fewer, broader groups."
        ),
    )
    n_clusters = st.slider("Number of clusters (K)", min_value=2, max_value=15, value=4)

    st.subheader("Classification labels")
    labels_input = st.text_input(
        "Comma-separated target labels",
        value="",
        placeholder="e.g. Technical, Billing, General, Feedback",
    )

    if st.session_state.result_df is not None:
        st.divider()
        st.subheader("Plot options")
        has_labels_sidebar = st.session_state.has_labels
        view_options = ["Unsupervised Clusters"]
        if has_labels_sidebar:
            view_options.insert(0, "Predicted Classification Labels")
        view_mode = st.radio("Color points by:", view_options)
    else:
        view_mode = "Unsupervised Clusters"

    st.divider()
    run_clicked = st.button("Run Analysis", type="primary", use_container_width=True)

# ----------------------------------------------------------------------------
# Run pipeline only on button click
# ----------------------------------------------------------------------------
if run_clicked:
    if uploaded_file is None:
        st.warning("Please upload a document first.")
    else:
        file_bytes = uploaded_file.getvalue()

        try:
            chunks, embeddings = compute_chunks_and_embeddings(
                file_bytes, uploaded_file.name, chunk_size
            )
        except Exception as e:
            st.error(f"Failed to process file: {e}")
            st.stop()

        if len(chunks) == 0:
            st.warning("No extractable text was found in this document.")
            st.stop()

        if len(chunks) < 3:
            st.warning(
                f"Only {len(chunks)} chunk(s) were generated. "
                "Try a smaller chunk size or a longer document for meaningful "
                "clustering and projection."
            )
            st.stop()

        raw_labels = [lbl.strip() for lbl in labels_input.split(",") if lbl.strip()]
        has_labels = len(raw_labels) >= 2

        coords = compute_2d_projection(embeddings)
        df = pd.DataFrame(coords, columns=["x", "y"])
        df["chunk_index"] = np.arange(len(chunks))
        df["snippet"] = [c[:120] + ("..." if len(c) > 120 else "") for c in chunks]

        cluster_labels = compute_clusters(embeddings, n_clusters)
        df["cluster"] = [f"Cluster {c}" for c in cluster_labels]

        if has_labels:
            cache_key = "|".join(raw_labels)
            predicted_labels, confidence = compute_classification(
                cache_key, embeddings, tuple(raw_labels)
            )
            df["classification"] = predicted_labels
            df["confidence"] = confidence
        else:
            df["classification"] = "N/A"
            df["confidence"] = 0.0

        st.session_state.result_df = df
        st.session_state.chunks = chunks
        st.session_state.has_labels = has_labels
        st.rerun()

# ----------------------------------------------------------------------------
# Main area
# ----------------------------------------------------------------------------
if st.session_state.result_df is None:
    st.info("Please upload a .txt or .pdf file and click Run Analysis in the sidebar to begin.")
else:
    df = st.session_state.result_df
    chunks = st.session_state.chunks
    has_labels = st.session_state.has_labels

    color_col = "classification" if view_mode == "Predicted Classification Labels" else "cluster"

    if not has_labels and view_mode == "Predicted Classification Labels":
        st.caption(
            "Enter at least two comma-separated labels in the sidebar and rerun "
            "to enable zero-shot classification coloring."
        )

    # ------------------------------------------------------------------------
    # Scatter plot
    # ------------------------------------------------------------------------
    fig = px.scatter(
        df,
        x="x",
        y="y",
        color=color_col,
        hover_data={"snippet": True, "chunk_index": True, "x": False, "y": False},
        custom_data=["chunk_index"],
        title=f"Semantic Map (t-SNE) - colored by {view_mode}",
        template="plotly_dark",
        height=560,
    )
    fig.update_traces(
        marker=dict(size=10, opacity=0.8, line=dict(width=1, color="DarkSlateGrey"))
    )
    fig.update_layout(
        legend_title_text=view_mode,
        xaxis_title="Dimension 1",
        yaxis_title="Dimension 2",
        margin=dict(l=10, r=10, t=50, b=10),
    )

    event = st.plotly_chart(
        fig,
        use_container_width=True,
        on_select="rerun",
        key="semantic_scatter_chart",
        selection_mode="points",
    )

    clicked_index = None
    if event and "selection" in event and event["selection"].get("points"):
        point = event["selection"]["points"][0]
        try:
            clicked_index = int(point["customdata"][0])
        except (KeyError, IndexError, TypeError, ValueError):
            clicked_index = None

    # ------------------------------------------------------------------------
    # Side-by-side: distribution + snippet viewer
    # ------------------------------------------------------------------------
    st.divider()
    col_left, col_right = st.columns([1, 1.4])

    with col_left:
        st.subheader("Group Distribution")
        group_counts = df[color_col].value_counts(normalize=True).sort_index()
        for group_name, pct in group_counts.items():
            st.write(f"**{group_name}** - {pct * 100:.1f}%")
            st.progress(min(max(pct, 0.0), 1.0))

        if color_col == "classification" and has_labels:
            avg_conf = df["confidence"].mean()
            st.metric("Avg. classification confidence", f"{avg_conf:.2f}")

    with col_right:
        st.subheader("Text Snippets")

        if clicked_index is None:
            st.info("Click a data point on the scatter plot to inspect its text chunk.")
        elif clicked_index >= len(chunks):
            st.warning("Could not resolve the selected point. Please click another point.")
        else:
            clicked_group = df.loc[df["chunk_index"] == clicked_index, color_col].values[0]

            with st.expander(
                f"Selected Chunk {clicked_index} (Group: {clicked_group})", expanded=True
            ):
                st.write(chunks[clicked_index])
                if color_col == "classification" and has_labels:
                    conf_val = df.loc[
                        df["chunk_index"] == clicked_index, "confidence"
                    ].values[0]
                    st.caption(f"Classification confidence: {conf_val:.3f}")

            st.markdown(f"**Other chunks in '{clicked_group}':**")
            same_group_df = df[
                (df[color_col] == clicked_group) & (df["chunk_index"] != clicked_index)
            ]

            if same_group_df.empty:
                st.caption("No other chunks share this group.")
            else:
                for _, row in same_group_df.iterrows():
                    idx = int(row["chunk_index"])
                    with st.expander(f"Chunk {idx}"):
                        st.write(chunks[idx])