from pathlib import Path

import spaces
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch

import json
import random
import re
import time
import uuid
import zipfile

from bs4 import BeautifulSoup
from docx import Document as DocxDocument
from langchain_text_splitters import RecursiveCharacterTextSplitter
from pypdf import PdfReader
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer

SEED = 42
MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
BATCH_SIZE = 8
DEVICE = "cpu"

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
_embedding_model = None


def get_embedding_model():
    global _embedding_model
    if _embedding_model is None:
        _embedding_model = SentenceTransformer(MODEL_ID, device=DEVICE)
    return _embedding_model


def clean_text(value):
    value = value.replace("\x00", " ")
    value = value.replace("\r\n", "\n").replace("\r", "\n")
    value = re.sub(r"[ \t]+", " ", value)
    value = re.sub(r" *\n *", "\n", value)
    value = re.sub(r"\n{3,}", "\n\n", value)
    return value.strip()


def extract_pdf_page(page):
    try:
        return page.extract_text(extraction_mode="layout") or ""
    except TypeError:
        return page.extract_text() or ""


def extract_document(path):
    path = Path(path)
    suffix = path.suffix.lower()
    page_count = None

    if suffix == ".pdf":
        reader = PdfReader(str(path))
        pages = [extract_pdf_page(page) for page in reader.pages]
        raw_text = "\n\n".join(pages)
        page_count = len(pages)
    elif suffix == ".docx":
        document = DocxDocument(str(path))
        raw_text = "\n\n".join(
            paragraph.text
            for paragraph in document.paragraphs
            if paragraph.text.strip()
        )
    elif suffix == ".txt":
        raw_text = path.read_text(encoding="utf-8", errors="ignore")
    elif suffix in {".html", ".htm"}:
        html = path.read_text(encoding="utf-8", errors="ignore")
        soup = BeautifulSoup(html, "lxml")
        for tag in soup(["script", "style", "nav", "footer"]):
            tag.decompose()
        raw_text = soup.get_text("\n")
    else:
        raise ValueError("Upload a PDF, DOCX, TXT, HTML or HTM file.")

    return clean_text(raw_text), page_count, suffix[1:].upper()


def token_count(value):
    return len(
        tokenizer.encode(
            value,
            add_special_tokens=False,
            truncation=False,
            verbose=False,
        )
    )


def split_sentences(value):
    return [
        sentence.strip()
        for sentence in re.split(r"(?<=[.!?])\s+|\n+", value)
        if sentence.strip()
    ]


def fixed_chunks(value, chunk_size, overlap_ratio):
    token_ids = tokenizer.encode(value, add_special_tokens=False)
    overlap = int(chunk_size * overlap_ratio)
    step = max(1, chunk_size - overlap)
    chunks = []

    for start in range(0, len(token_ids), step):
        piece = token_ids[start:start + chunk_size]
        if not piece:
            continue
        chunk = tokenizer.decode(piece, skip_special_tokens=True).strip()
        if chunk:
            chunks.append(chunk)
        if start + chunk_size >= len(token_ids):
            break

    return chunks


def merge_units(units, chunk_size, overlap_ratio):
    chunks = []
    current = []

    for unit in units:
        if token_count(unit) > chunk_size:
            if current:
                chunks.append(" ".join(current).strip())
                current = []
            chunks.extend(fixed_chunks(unit, chunk_size, overlap_ratio))
            continue

        candidate = " ".join(current + [unit]).strip()
        if current and token_count(candidate) > chunk_size:
            chunks.append(" ".join(current).strip())
            overlap_budget = int(chunk_size * overlap_ratio)
            carried = []
            for previous in reversed(current):
                test = " ".join([previous] + carried).strip()
                if token_count(test) <= overlap_budget:
                    carried.insert(0, previous)
                else:
                    break
            current = carried + [unit]
        else:
            current.append(unit)

    if current:
        chunks.append(" ".join(current).strip())
    return [chunk for chunk in chunks if chunk]


def sentence_chunks(value, chunk_size, overlap_ratio):
    return merge_units(split_sentences(value), chunk_size, overlap_ratio)


def paragraph_chunks(value, chunk_size, overlap_ratio):
    paragraphs = [
        paragraph.strip()
        for paragraph in re.split(r"\n\s*\n", value)
        if paragraph.strip()
    ]
    return merge_units(paragraphs, chunk_size, overlap_ratio)


def recursive_chunks(value, chunk_size, overlap_ratio):
    splitter = RecursiveCharacterTextSplitter.from_huggingface_tokenizer(
        tokenizer,
        chunk_size=chunk_size,
        chunk_overlap=int(chunk_size * overlap_ratio),
        separators=["\n\n", "\n", ". ", " ", ""],
    )
    return [chunk.strip() for chunk in splitter.split_text(value) if chunk.strip()]


def semantic_chunks(value, target_size, break_percentile=25):
    sentences = split_sentences(value)
    if len(sentences) <= 1:
        return sentences

    vectors = get_embedding_model().encode(
        sentences,
        batch_size=BATCH_SIZE,
        normalize_embeddings=True,
        convert_to_numpy=True,
        show_progress_bar=False,
    )
    similarities = np.sum(vectors[:-1] * vectors[1:], axis=1)
    threshold = np.percentile(similarities, break_percentile)
    chunks = []
    current = [sentences[0]]

    for index in range(1, len(sentences)):
        candidate = " ".join(current + [sentences[index]])
        topic_shift = similarities[index - 1] < threshold
        too_large = token_count(candidate) > target_size
        if topic_shift or too_large:
            chunks.append(" ".join(current).strip())
            current = [sentences[index]]
        else:
            current.append(sentences[index])

    if current:
        chunks.append(" ".join(current).strip())
    return chunks


def create_chunks(text, strategy, chunk_size, overlap_percent):
    overlap_ratio = overlap_percent / 100
    if strategy == "Fixed":
        return fixed_chunks(text, chunk_size, overlap_ratio)
    if strategy == "Sentence":
        return sentence_chunks(text, chunk_size, overlap_ratio)
    if strategy == "Paragraph":
        return paragraph_chunks(text, chunk_size, overlap_ratio)
    if strategy == "Recursive":
        return recursive_chunks(text, chunk_size, overlap_ratio)
    if strategy == "Semantic":
        return semantic_chunks(text, chunk_size)
    raise ValueError(f"Unsupported strategy: {strategy}")


def boundary_score(chunks):
    if not chunks:
        return 0.0
    valid = sum(
        bool(re.search(r"[.!?][\"')\]]?$", chunk.strip()))
        for chunk in chunks
    )
    return valid / len(chunks)


def context_preservation_score(source_text, chunks, sample_size=100):
    candidates = [
        sentence
        for sentence in split_sentences(source_text)
        if token_count(sentence) >= 8
    ]
    if not candidates:
        return 0.0
    sample = random.Random(SEED).sample(
        candidates,
        min(sample_size, len(candidates)),
    )
    preserved = sum(any(sentence in chunk for chunk in chunks) for sentence in sample)
    return preserved / len(sample)


def run_benchmark(uploaded_file, strategy, chunk_size, overlap_percent, maximum_chunks):
    if uploaded_file is None:
        raise gr.Error("Upload one document first.")

    run_directory = Path("/tmp") / f"benchmark_{uuid.uuid4().hex[:8]}"
    run_directory.mkdir(parents=True, exist_ok=True)

    try:
        text, page_count, document_type = extract_document(uploaded_file)
    except Exception as error:
        raise gr.Error(f"Document extraction failed: {error}")

    if len(text) < 100:
        raise gr.Error("Very little text was extracted. A scanned PDF may require OCR.")

    document_summary = pd.DataFrame([{
        "File": Path(uploaded_file).name,
        "Type": document_type,
        "Pages": page_count,
        "Characters": len(text),
        "Words": len(text.split()),
        "Tokens": token_count(text),
    }])

    started = time.perf_counter()
    chunks = create_chunks(text, strategy, int(chunk_size), int(overlap_percent))
    chunk_seconds = time.perf_counter() - started
    if not chunks:
        raise gr.Error("The selected configuration produced no chunks.")

    lengths = [token_count(chunk) for chunk in chunks]
    duplicates = len(chunks) - len(set(chunks))
    effective_overlap = 0 if strategy == "Semantic" else int(overlap_percent)

    chunk_summary = pd.DataFrame([{
        "Strategy": strategy,
        "Target tokens": int(chunk_size),
        "Overlap": f"{effective_overlap}%",
        "Chunks": len(chunks),
        "Average tokens": round(float(np.mean(lengths)), 2),
        "Minimum tokens": int(min(lengths)),
        "Maximum tokens": int(max(lengths)),
        "Boundary score": round(boundary_score(chunks), 4),
        "Context preservation": round(context_preservation_score(text, chunks), 4),
        "Duplicate ratio": round(duplicates / max(len(chunks), 1), 4),
        "Processing seconds": round(chunk_seconds, 4),
    }])

    preview = "\n\n".join(
        f"CHUNK {index}\n{'=' * 60}\n{chunk[:1200]}"
        for index, chunk in enumerate(chunks[:3], start=1)
    )

    selected_chunks = chunks[:int(maximum_chunks)]
    model = get_embedding_model()
    started = time.perf_counter()
    vectors = model.encode(
        selected_chunks,
        batch_size=BATCH_SIZE,
        normalize_embeddings=True,
        convert_to_numpy=True,
        show_progress_bar=False,
    )
    embedding_seconds = time.perf_counter() - started

    test_pairs = {
        "Related": [
            "Chunking divides a document into smaller retrieval units.",
            "Documents are split into smaller passages before retrieval.",
        ],
        "Unrelated": [
            "Chunking divides a document into smaller retrieval units.",
            "The football match was postponed because of heavy rain.",
        ],
    }
    similarity_rows = []
    for pair_name, pair in test_pairs.items():
        pair_vectors = model.encode(
            pair,
            normalize_embeddings=True,
            convert_to_numpy=True,
            show_progress_bar=False,
        )
        similarity_rows.append({
            "Pair": pair_name,
            "Cosine similarity": round(float(np.dot(pair_vectors[0], pair_vectors[1])), 4),
        })
    similarity_df = pd.DataFrame(similarity_rows)

    embedding_summary = pd.DataFrame([{
        "Model": "all-MiniLM-L6-v2",
        "Device": DEVICE,
        "Chunks embedded": len(selected_chunks),
        "Vector dimension": int(vectors.shape[1]),
        "Output shape": str(tuple(vectors.shape)),
        "Embedding seconds": round(embedding_seconds, 3),
        "Chunks per second": round(len(selected_chunks) / max(embedding_seconds, 1e-9), 3),
        "Vector storage MB": round(vectors.nbytes / (1024 ** 2), 4),
    }])

    figure, axes = plt.subplots(1, 2, figsize=(13, 5))
    axes[0].hist(lengths, bins=min(15, max(5, len(lengths))), color="#4F46E5", edgecolor="white")
    axes[0].set_title("Chunk token distribution")
    axes[0].set_xlabel("Tokens per chunk")
    axes[0].set_ylabel("Number of chunks")
    axes[0].grid(alpha=0.2)
    axes[1].bar(
        similarity_df["Pair"],
        similarity_df["Cosine similarity"],
        color=["#16A34A", "#DC2626"],
    )
    axes[1].set_title("Semantic similarity test")
    axes[1].set_xlabel("Sentence-pair type")
    axes[1].set_ylabel("Cosine similarity")
    axes[1].grid(axis="y", alpha=0.2)
    figure.tight_layout()

    document_summary.to_csv(run_directory / "document_summary.csv", index=False)
    chunk_summary.to_csv(run_directory / "chunk_summary.csv", index=False)
    pd.DataFrame({
        "chunk_id": range(len(chunks)),
        "tokens": lengths,
        "text": chunks,
    }).to_csv(run_directory / "chunks.csv", index=False)
    embedding_summary.to_csv(run_directory / "embedding_summary.csv", index=False)
    similarity_df.to_csv(run_directory / "similarity_results.csv", index=False)
    np.save(run_directory / "vectors.npy", vectors)
    figure.savefig(run_directory / "benchmark_plots.png", dpi=180, bbox_inches="tight")

    settings = {
        "strategy": strategy,
        "chunk_size": int(chunk_size),
        "overlap_percent": effective_overlap,
        "model": MODEL_ID,
        "maximum_chunks": int(maximum_chunks),
        "device": DEVICE,
        "seed": SEED,
    }
    with open(run_directory / "settings.json", "w") as file:
        json.dump(settings, file, indent=2)

    archive_path = run_directory / "benchmark_results.zip"
    with zipfile.ZipFile(archive_path, "w", zipfile.ZIP_DEFLATED) as archive:
        for output_path in run_directory.iterdir():
            if output_path.is_file() and output_path != archive_path:
                archive.write(output_path, arcname=output_path.name)

    status = (
        "### Benchmark completed\n"
        f"Created **{len(chunks)} chunks** and embedded "
        f"**{len(selected_chunks)} chunks** using **all-MiniLM-L6-v2**."
    )
    return (
        status,
        document_summary,
        chunk_summary,
        preview,
        embedding_summary,
        figure,
        str(archive_path),
    )


with gr.Blocks(title="Chunking and Embedding Benchmark") as demo:
    gr.Markdown(
        """
        # Chunking and Embedding Benchmark

        Upload a document and test a chunking configuration. The application
        measures chunk quality, generates MiniLM embeddings and exports the results.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            uploaded_file = gr.File(
                label="Upload document",
                file_types=[".pdf", ".docx", ".txt", ".html", ".htm"],
                type="filepath",
            )
            strategy = gr.Dropdown(
                choices=["Fixed", "Sentence", "Paragraph", "Recursive", "Semantic"],
                value="Recursive",
                label="Chunking strategy",
            )
            chunk_size = gr.Dropdown(
                choices=[128, 256, 512, 1024],
                value=256,
                label="Target chunk size",
            )
            overlap_percent = gr.Dropdown(
                choices=[0, 10, 20, 30],
                value=10,
                label="Overlap percentage",
            )
            maximum_chunks = gr.Slider(
                minimum=5,
                maximum=300,
                value=100,
                step=5,
                label="Maximum chunks to embed",
            )
            run_button = gr.Button("Run Benchmark", variant="primary")

        with gr.Column(scale=2):
            status = gr.Markdown("Upload a document and run the benchmark.")
            document_output = gr.Dataframe(label="Document summary", interactive=False)
            chunk_output = gr.Dataframe(label="Chunking results", interactive=False)

    with gr.Tab("Chunk preview"):
        preview_output = gr.Textbox(label="First three chunks", lines=22, interactive=False)

    with gr.Tab("Embedding results"):
        embedding_output = gr.Dataframe(label="Embedding measurements", interactive=False)
        plot_output = gr.Plot(label="Benchmark plots")

    with gr.Tab("Download"):
        download_output = gr.File(label="Download benchmark results")

    run_button.click(
        fn=run_benchmark,
        inputs=[uploaded_file, strategy, chunk_size, overlap_percent, maximum_chunks],
        outputs=[
            status,
            document_output,
            chunk_output,
            preview_output,
            embedding_output,
            plot_output,
            download_output,
        ],
    )


if __name__ == "__main__":
    demo.queue().launch(
        theme=gr.themes.Soft(),
        ssr_mode=False,
        server_name="0.0.0.0",
        server_port=7860,
    )