app.py

#!/usr/bin/env python3
"""Gradio app for structural manifold sidecar: compression, reconstruction, and verification."""

from __future__ import annotations

import io
import textwrap
from pathlib import Path
from typing import Dict, Optional, Tuple

import os
import tempfile
import zipfile
import gradio as gr
import gradio.networking as gr_networking
import gradio_client.utils as grc_utils
import matplotlib.pyplot as plt
import numpy as np

import sys

# Ensure Gradio detects the Spaces runtime so it skips localhost/share checks.
if (os.getenv("SPACE_ID") or os.getenv("HF_SPACE_ID")) and not os.getenv("SYSTEM"):
    os.environ["SYSTEM"] = "spaces"

# HF Spaces sometimes blocks loopback checks; force Gradio to treat local URL as reachable.
gr_networking.url_ok = lambda url: True  # type: ignore[assignment]

# Guard against boolean JSON schemas tripping gradio_client utils.
_orig_json_schema_to_python_type = grc_utils._json_schema_to_python_type
_orig_get_type = grc_utils.get_type


def _safe_json_schema_to_python_type(schema, defs=None):
    if isinstance(schema, bool):
        return "boolean"
    return _orig_json_schema_to_python_type(schema, defs)


def _safe_get_type(schema):
    if isinstance(schema, bool):
        return "boolean"
    return _orig_get_type(schema)


grc_utils._json_schema_to_python_type = _safe_json_schema_to_python_type
grc_utils.get_type = _safe_get_type

REPO_ROOT = Path(__file__).resolve().parent
SRC_PATH = REPO_ROOT / "src"
if str(SRC_PATH) not in sys.path:
    sys.path.insert(0, str(SRC_PATH))
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

WINDOW_BYTES = 128
STRIDE_BYTES = 96
EMBEDDING_MODEL_NAME = "all-MiniLM-L6-v2"
ENABLE_RETRIEVE = os.getenv("ENABLE_RETRIEVE", "0").lower() in {"1", "true", "yes"}

from manifold.sidecar import (
    EncodeResult,
    ManifoldIndex,
    build_index,
    encode_text,
    reconstruct_from_windows,
    verify_snippet,
)

try:
    import pdfplumber  # type: ignore
except Exception:  # pragma: no cover - optional dependency handled by requirements.txt
    pdfplumber = None

try:
    from sentence_transformers import SentenceTransformer  # type: ignore
except Exception:  # pragma: no cover - lazy load handled later
    SentenceTransformer = None


docs_store: Dict[str, str] = {}
encodings_store: Dict[str, EncodeResult] = {}
doc_counter = 0
_embedding_model = None


def _next_doc_id() -> str:
    global doc_counter
    doc_counter += 1
    return f"doc-{doc_counter}"


def _extract_text_from_file(file_obj) -> Tuple[Optional[str], Optional[str]]:
    if file_obj is None:
        return None, None
    path: Optional[Path] = None
    raw_bytes: bytes

    if hasattr(file_obj, "read"):
        # File-like object (older Gradio versions)
        name = getattr(file_obj, "name", None)
        path = Path(name) if name else None
        raw_bytes = file_obj.read()
        try:
            file_obj.seek(0)
        except Exception:
            pass
    else:
        # Gradio v4 returns NamedString (path-like) or plain path string
        if isinstance(file_obj, (str, Path)):
            path = Path(file_obj)
        elif hasattr(file_obj, "name"):
            path = Path(file_obj.name)
        else:
            raise ValueError(f"Unsupported file object type: {type(file_obj)}")
        raw_bytes = path.read_bytes()

    if path is None:
        return None, None

    suffix = path.suffix.lower()
    if suffix in {".txt", ".md"}:
        text = raw_bytes.decode("utf-8", errors="ignore")
        return path.name, text
    if suffix == ".pdf":
        if pdfplumber is None:
            raise RuntimeError("pdfplumber is required for PDF ingestion. Install with `pip install pdfplumber`.")
        with pdfplumber.open(io.BytesIO(raw_bytes)) as pdf:
            pages = [page.extract_text() or "" for page in pdf.pages]
        text = "\n\n".join(pages).strip()
        return path.name, text
    raise ValueError(f"Unsupported file type: {suffix}")


def _make_hazard_plot(hazards):
    fig, ax = plt.subplots(figsize=(5, 3))
    if hazards:
        ax.hist(hazards, bins=20, color="#2f6fff", alpha=0.8)
    ax.set_title("Window hazards")
    ax.set_xlabel("Hazard λ")
    ax.set_ylabel("Window count")
    fig.tight_layout()
    return fig


def _preview(text: str, limit: int = 2000) -> str:
    if len(text) <= limit:
        return text
    return text[:limit] + f"\n\n… [truncated {len(text) - limit} chars]"


def _chunk_text(text: str, chunk_size: int = 512, overlap: int = 128) -> list[tuple[str, str]]:
    chunks = []
    start = 0
    text_len = len(text)
    idx = 0
    while start < text_len:
        end = min(text_len, start + chunk_size)
        chunk = text[start:end]
        chunks.append((f"chunk-{idx}", chunk))
        if end == text_len:
            break
        start = end - overlap
        idx += 1
    return chunks


def _get_embedding_model():
    global _embedding_model
    if _embedding_model is None:
        if SentenceTransformer is None:
            raise RuntimeError(
                "sentence-transformers is required for retrieval demo. Install with `pip install sentence-transformers`."
            )
        _embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME)
    return _embedding_model


def _embed_texts(texts: list[str]) -> np.ndarray:
    model = _get_embedding_model()
    embeddings = model.encode(texts, convert_to_numpy=True, normalize_embeddings=True)
    return embeddings.astype(np.float32)


def handle_reconstruct_download():
    """Reconstruct all ingested docs and bundle as a ZIP for download."""
    if not encodings_store:
        return None
    tmp_dir = tempfile.mkdtemp(prefix="manifold_recon_")
    zip_path = os.path.join(tmp_dir, "reconstructed_corpus.zip")
    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zipf:
        for doc_id, encoded in encodings_store.items():
            text = reconstruct_from_windows(encoded.windows, encoded.prototypes)
            header = f"--- RECONSTRUCTED BY SEP MANIFOLD (ID: {doc_id}) ---\n\n"
            zipf.writestr(f"{doc_id}.txt", header + text)
    return zip_path


def handle_compress(files, raw_text):
    try:
        ingested = []
        combined_hazards: list[float] = []
        total_before = 0
        total_after = 0
        first_original = ""
        first_recon = ""
        last_doc_id = None

        incoming_files = files or []
        if not isinstance(incoming_files, (list, tuple)):
            incoming_files = [incoming_files] if incoming_files else []

        if raw_text and raw_text.strip():
            incoming_files = list(incoming_files) + [("pasted.txt", raw_text)]

        if not incoming_files:
            return "No text provided.", "", "", "", None, gr.update(choices=list(docs_store.keys()), value=None)

        for file_obj in incoming_files:
            # Support pasted text passed as (name, text)
            if isinstance(file_obj, tuple) and len(file_obj) == 2:
                name, text_content = file_obj
            else:
                name, text_content = _extract_text_from_file(file_obj)
            if not text_content or not str(text_content).strip():
                continue

            doc_id = _next_doc_id()
            docs_store[doc_id] = text_content
            encoded = encode_text(
                text_content,
                window_bytes=WINDOW_BYTES,
                stride_bytes=STRIDE_BYTES,
            )
            encodings_store[doc_id] = encoded
            reconstruction = reconstruct_from_windows(encoded.windows, encoded.prototypes)

            if not first_original:
                first_original = text_content
                first_recon = reconstruction
            last_doc_id = doc_id

            unique_sigs = len(encoded.prototypes)
            bytes_before = encoded.original_bytes
            bytes_after = unique_sigs * 9  # approx storage per signature
            total_before += bytes_before
            total_after += bytes_after
            combined_hazards.extend(encoded.hazards)

            compression_ratio = (bytes_before / bytes_after) if bytes_after else 0.0
            ingested.append(
                f"- **{doc_id}** ({name or 'upload'}): windows={len(encoded.windows)}, "
                f"unique sigs={unique_sigs}, hazard gate ≤ {encoded.hazard_threshold:.4f}, "
                f"ratio≈{compression_ratio:.2f}×"
            )

        if not ingested:
            return "No valid text provided.", "", "", "", None, gr.update(choices=list(docs_store.keys()), value=None)

        total_ratio = (total_before / total_after) if total_after else 0.0
        stats = textwrap.dedent(
            f"""
            **Ingested {len(ingested)} document(s)**  
            Original bytes: {total_before} (~{total_before/1e6:.2f} MB)  
            Manifold payload bytes (~signatures): {total_after} (~{total_after/1e6:.2f} MB)  
            Approx compression: {total_ratio:.2f}×

            Details:
            {chr(10).join(ingested)}
            """
        ).strip()

        fig = _make_hazard_plot(combined_hazards)
        if combined_hazards:
            threshold = np.quantile(np.array(combined_hazards), 0.8)
            ax = fig.axes[0]
            ax.axvline(threshold, color="red", linestyle="--", label="hazard gate (80th pct)")
            ax.legend()

        dropdown_update = gr.update(choices=list(docs_store.keys()), value=last_doc_id)
        return (
            f"Stored {len(ingested)} doc(s)",
            _preview(first_original),
            _preview(first_recon),
            stats,
            fig,
            dropdown_update,
        )
    except Exception as exc:  # pragma: no cover - UI surface
        return f"Error: {exc}", "", "", "", None, gr.update(choices=list(docs_store.keys()), value=None)


def _ensure_index() -> Optional[ManifoldIndex]:
    if not docs_store:
        return None
    return build_index(
        docs_store,
        window_bytes=WINDOW_BYTES,
        stride_bytes=STRIDE_BYTES,
    )


def handle_verify(selected_doc, snippet, coverage_threshold):
    if not snippet or not snippet.strip():
        return "Provide a snippet to verify.", ""
    index = _ensure_index()
    if index is None:
        return "No documents ingested yet.", ""

    meta = getattr(index, "meta", {}) if hasattr(index, "meta") else {}
    default_hazard_threshold = float(meta.get("hazard_threshold", 0.8))
    hazard_threshold = handle_verify.hazard_threshold  # type: ignore[attr-defined]
    if hazard_threshold is None:
        hazard_threshold = default_hazard_threshold

    # Dense scan to recover alignment for imperfect copy/paste.
    VERIFY_WINDOW = 128
    VERIFY_STRIDE = 4

    snippet_bytes = len(snippet.encode("utf-8"))
    if snippet_bytes < VERIFY_WINDOW:
        return (
            "<span style='color:orange; font-weight:700;'>⚠️ Snippet too short</span>",
            f"Snippet is {snippet_bytes} bytes; minimum window is {VERIFY_WINDOW} bytes to form a signature.",
        )

    result = verify_snippet(
        snippet,
        index,
        coverage_threshold=coverage_threshold,
        hazard_threshold=hazard_threshold,
        window_bytes=VERIFY_WINDOW,
        stride_bytes=VERIFY_STRIDE,
        include_reconstruction=False,
        min_text_similarity=0.0,
    )
    total = max(result.total_windows, 1)
    valid_matches = [m for m in result.matches if m.get("matched") and m.get("occurrences")]
    raw_hits = len(valid_matches)
    hazard_hits = sum(1 for m in valid_matches if m.get("hazard_ok"))
    raw_coverage = raw_hits / total
    safe_coverage = hazard_hits / total
    verified = safe_coverage >= coverage_threshold
    status = "✅ Verified" if verified else "❌ Not verified"
    status_color = "green" if verified else "red"
    status_line = (
        f"<span style='color:{status_color}; font-weight:700;'>{status}</span> "
        f"(Match density={raw_coverage*100:.1f}%, hazard-safe={safe_coverage*100:.1f}%, "
        f"hazard_gate ≤ {hazard_threshold:.3f})"
    )

    lines = []
    seen_sigs = set()
    for match in valid_matches[:50]:
        sig = str(match.get("signature", ""))
        if sig in seen_sigs:
            continue
        seen_sigs.add(sig)
        hz = float(match.get("hazard", 0.0))
        occ = match.get("occurrences", []) or []
        first_doc = occ[0].get("doc_id") if occ else "unknown"
        lines.append(f"- Signature `{sig[:12]}...` | hazard={hz:.3f} | found in: {first_doc}")
    matches_md = "\n".join(lines) if lines else "_No valid signatures found in index._"
    return status_line, matches_md


def handle_retrieve(question, top_k, coverage_threshold, hazard_threshold):
    if not question or not question.strip():
        return "Provide a question.", "", ""
    if not docs_store:
        return "No documents ingested yet.", "", ""

    index = _ensure_index()
    if index is None:
        return "No documents ingested yet.", "", ""

    # Build chunks
    chunks = []
    for doc_id, text in docs_store.items():
        for chunk_id, chunk_text in _chunk_text(text):
            chunks.append((doc_id, chunk_id, chunk_text))
    if not chunks:
        return "No chunks available to retrieve.", "", ""

    chunk_texts = [c[2] for c in chunks]
    chunk_embeddings = _embed_texts(chunk_texts)
    question_embedding = _embed_texts([question])[0]
    scores = np.dot(chunk_embeddings, question_embedding)
    order = np.argsort(scores)[::-1]
    top_indices = order[: int(top_k)]

    naive_lines = []
    verified_lines = []
    search_lines = []
    for rank, idx in enumerate(top_indices, start=1):
        doc_id, chunk_id, chunk_text = chunks[int(idx)]
        score = float(scores[int(idx)])
        naive_lines.append(f"- [{rank}] {doc_id}::{chunk_id} score={score:.3f}\n  {chunk_text[:200]}...")

        result = verify_snippet(
            chunk_text,
            index,
            coverage_threshold=coverage_threshold,
            hazard_threshold=hazard_threshold,
            window_bytes=WINDOW_BYTES,
            stride_bytes=STRIDE_BYTES,
            include_reconstruction=True,
            min_text_similarity=0.7,
        )
        total = max(result.total_windows, 1)
        raw_hits = sum(1 for m in result.matches if m.get("matched"))
        hazard_hits = sum(1 for m in result.matches if m.get("hazard_ok"))
        raw_coverage = raw_hits / total
        safe_coverage = hazard_hits / total
        text_sim = result.text_similarity
        status = "Verified" if result.verified else "Unverified"
        recon_preview = (result.reconstruction or chunk_text)[:280]
        search_lines.append(
            f"- [{rank}] {doc_id}::{chunk_id} — Index integrity: {status} "
            f"(hazard-gated & similarity), text_sim={text_sim:.2f}, "
            f"safe_cov={safe_coverage*100:.2f}% (hazard≤{hazard_threshold:.3f})\n"
            f"  Manifold reconstruction (approximate preview; originals preserved in-session): {recon_preview}..."
        )
        verified_lines.append(
            f"- [{rank}] {doc_id}::{chunk_id} raw={raw_coverage*100:.2f}%, safe={safe_coverage*100:.2f}%, "
            f"text_sim={text_sim:.2f}, hazard≤{hazard_threshold:.3f}"
        )
    search_md = "\n".join(search_lines) if search_lines else "_No matches_"
    diagnostics_md = "\n".join(verified_lines) if verified_lines else "_No diagnostics_"
    return "Search results (top-k):", search_md, diagnostics_md


with gr.Blocks(title="Structural Manifold Sidecar") as demo:
    gr.Markdown(
        "**Demo Mode:** Data is not persisted; refresh wipes uploads. Keep total uploads ≲10 MB on free tier to avoid timeouts."
    )
    gr.Markdown(
        "# Searchable Manifold Archive\nIngest large text sets at high compression and search them with on-the-fly reconstruction."
    )

    with gr.Tab("Ingest & Compress"):
        gr.Markdown(
            "Upload one or more documents (PDF/txt/md) or paste text. We build a manifold index, reconstruct a preview, "
            "and show compression + hazard stats."
        )
        gr.Markdown("_Tip: For the demo, limit uploads to ~10 MB total to avoid Space timeouts._")
        file_input = gr.File(
            label="Upload (.pdf, .txt, .md)",
            file_types=[".pdf", ".txt", ".md"],
            file_count="multiple",
        )
        text_input = gr.Textbox(label="Or paste text", lines=6)
        run_btn = gr.Button("Ingest & compress")
        doc_msg = gr.Markdown()
        original_box = gr.Textbox(label="Original preview (first doc)", lines=10)
        recon_box = gr.Textbox(
            label="Manifold reconstruction (approximate preview; originals kept in session only)",
            lines=10,
        )
        stats_box = gr.Markdown(label="Compression stats")
        hazard_plot = gr.Plot(label="Hazard histogram")
        gr.Markdown("### The 'Impossible' Download\nReconstruct the full corpus from the compressed index.")
        reconstruct_btn = gr.Button("Reconstruct & Download ZIP")
        download_output = gr.File(label="Download reconstructed corpus")

    with gr.Tab("Search Compressed Corpus"):
        gr.Markdown(
            "Paste a snippet or query. We search the manifold index, reconstruct matching windows, and report confidence."
        )
        doc_dropdown = gr.Dropdown(
            label="Docs ingested this session",
            choices=list(docs_store.keys()),
            interactive=True,
        )
        snippet_box = gr.Textbox(label="Snippet to verify", lines=6)
        coverage_slider = gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.7,
            step=0.05,
            label="Coverage threshold",
        )
        hazard_slider = gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.8,
            step=0.01,
            label="Hazard gate (raise to be more permissive)",
        )
        verify_btn = gr.Button("Verify")
        verify_status = gr.Markdown()
        verify_matches = gr.Markdown()

    if ENABLE_RETRIEVE:
        with gr.Tab("Retrieve & Verify"):
            gr.Markdown(
                "Chunk-level RAG demo: retrieve top-k chunks via embeddings, then hazard-gate them with manifold verification."
            )
            question_box = gr.Textbox(label="Question / query", lines=3)
            topk_slider = gr.Slider(minimum=1, maximum=10, value=5, step=1, label="Top-k chunks")
            rag_coverage = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.5,
                step=0.05,
                label="Coverage threshold (verification)",
            )
            rag_hazard = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.8,
                step=0.01,
                label="Hazard gate (verification)",
            )
            retrieve_btn = gr.Button("Retrieve & verify")
        retrieve_status = gr.Markdown()
        naive_rag = gr.Markdown(label="Search results (manifold reconstruction)")
        verified_rag = gr.Markdown(label="Diagnostics (coverage, hazard, similarity)")

    run_btn.click(
        handle_compress,
        inputs=[file_input, text_input],
        outputs=[doc_msg, original_box, recon_box, stats_box, hazard_plot, doc_dropdown],
    )
    def bound_verify(snippet, coverage, hazard):
        # stash hazard threshold on the function object so handle_verify can read it without changing signature
        handle_verify.hazard_threshold = hazard  # type: ignore[attr-defined]
        return handle_verify(None, snippet, coverage)

    verify_btn.click(
        bound_verify,
        inputs=[snippet_box, coverage_slider, hazard_slider],
        outputs=[verify_status, verify_matches],
    )
    if ENABLE_RETRIEVE:
        retrieve_btn.click(
            handle_retrieve,
            inputs=[question_box, topk_slider, rag_coverage, rag_hazard],
            outputs=[retrieve_status, naive_rag, verified_rag],
        )
    reconstruct_btn.click(
        handle_reconstruct_download,
        inputs=[],
        outputs=[download_output],
    )


if __name__ == "__main__":
    port = int(os.getenv("PORT", "7860"))
    server = "0.0.0.0" if os.getenv("SYSTEM") == "spaces" else "127.0.0.1"
    demo.queue().launch(
        server_name=server,
        server_port=port,
        share=False,
        inbrowser=False,
        show_error=True,
    )