"""
Map text chunks to PDF page numbers for Chroma/BM25 metadata.
"""
from __future__ import annotations

import os
from typing import Any, Dict, List, Optional, Sequence, Tuple

# Must match ContentExtractor / ChromaVectorDB fixed-window join between pages.
PDF_PAGE_SEPARATOR = "\n\n"

_STATIC_PREFIXES = (
    "extracted_images_",
    "document_images_",
    "uploaded_files_",
)

EMPTY_PAGE_PLACEHOLDER = "[empty page]"


def join_page_texts(page_texts: Sequence[str]) -> Tuple[str, List[Tuple[int, int]]]:
    """
    Join page texts with PDF_PAGE_SEPARATOR and record (start, end) char span of each page's
    text in the returned ``full_text`` (separator is not part of any page span).
    """
    pages = [((t or "").strip()) for t in page_texts]
    if not pages:
        return "", []

    parts: List[str] = []
    spans: List[Tuple[int, int]] = []
    pos = 0
    for i, p in enumerate(pages):
        start = pos
        parts.append(p)
        pos += len(p)
        spans.append((start, pos))
        if i < len(pages) - 1:
            parts.append(PDF_PAGE_SEPARATOR)
            pos += len(PDF_PAGE_SEPARATOR)
    return "".join(parts), spans


def pages_intersecting_span(
    span_start: int,
    span_end: int,
    page_spans: Sequence[Tuple[int, int]],
) -> List[int]:
    """Return sorted unique 1-based page numbers whose text span intersects [span_start, span_end)."""
    if span_start < 0:
        span_start = 0
    if span_end < span_start:
        return []
    if not page_spans:
        return []
    out: List[int] = []
    for i, (a, b) in enumerate(page_spans):
        if max(span_start, a) < min(span_end, b):
            out.append(i + 1)
    return out


def to_static_relative_path(abs_path: str) -> str:
    """Reduce an absolute path to a server-relative prefix for static file routes."""
    if not abs_path:
        return ""
    norm = abs_path.replace("\\", "/")
    for pref in _STATIC_PREFIXES:
        j = norm.find(pref)
        if j >= 0:
            return norm[j:]
    return os.path.basename(norm)


def build_page_chunks(
    file_id: str,
    page_texts: Sequence[str],
    base_metadata: dict,
    *,
    empty_placeholder: str = EMPTY_PAGE_PLACEHOLDER,
) -> Tuple[List[str], List[str], List[dict]]:
    """
    One Chroma/BM25 chunk per page with stable ids ``{file_id}_page_{n}`` (1-based n).
    """
    ids: List[str] = []
    documents: List[str] = []
    metadatas: List[dict] = []
    total = len(page_texts)

    for n in range(1, total + 1):
        raw = page_texts[n - 1] if n - 1 < total else ""
        text = (raw or "").strip()
        if not text:
            text = empty_placeholder

        chunk_id = f"{file_id}_page_{n}"
        ids.append(chunk_id)
        documents.append(text)

        chunk_metadata = base_metadata.copy() if isinstance(base_metadata, dict) else {}
        chunk_metadata.update(
            {
                "chunk_index": n - 1,
                "total_chunks": total,
                "file_id": file_id,
                "page_number": n,
            }
        )
        metadatas.append(chunk_metadata)

    return ids, documents, metadatas


def page_number_from_metadata(meta: dict) -> Optional[int]:
    """Resolve singular page_number from chunk metadata (new or legacy)."""
    if not meta:
        return None
    pn = meta.get("page_number")
    if pn is not None:
        try:
            return int(pn)
        except (TypeError, ValueError):
            pass
    raw = meta.get("page_numbers")
    if raw:
        first = str(raw).split(",")[0].strip()
        if first.isdigit():
            return int(first)
    return None


def best_effort_span_for_semantic_chunk(
    chunk_text: str,
    full_text: str,
) -> Optional[Tuple[int, int]]:
    """Find the first occurrence of chunk_text in full_text for page mapping."""
    if not chunk_text or not full_text:
        return None
    needle = chunk_text.strip()
    if len(needle) < 20:
        return None
    idx = full_text.find(needle)
    if idx < 0:
        return None
    idx2 = full_text.find(needle, idx + 1)
    if idx2 >= 0 and idx2 < idx + len(needle):
        return None
    return idx, idx + len(needle)


def build_page_images_for_find(
    extracted_rows: Sequence,
    matching_pages: Optional[Sequence[int]] = None,
) -> Dict[str, List[Dict[str, Any]]]:
    """
    Build ``page_images`` for @find from embedded images (not page snapshots).
    """
    matching = {int(p) for p in (matching_pages or []) if p is not None}
    by_page: Dict[int, List[Dict[str, Any]]] = {}
    for row in extracted_rows or []:
        if not row or len(row) < 4:
            continue
        img_type = (row[2] or "").strip()
        if img_type == "page_snapshot":
            continue
        rel_path = to_static_relative_path(row[1] or "")
        if not rel_path:
            continue
        page = int(row[3] or 0)
        if page < 1:
            continue
        if matching and page not in matching:
            continue
        by_page.setdefault(page, []).append(
            {
                "path": rel_path,
                "type": img_type or "embedded_image",
                "is_matching": bool(matching and page in matching),
            }
        )
    out: Dict[str, List[Dict[str, Any]]] = {}
    for page in sorted(by_page.keys()):
        seen: set[tuple[str, str]] = set()
        items: List[Dict[str, Any]] = []
        for img in by_page[page]:
            key = (img["path"], img["type"])
            if key in seen:
                continue
            seen.add(key)
            items.append(dict(img))
        out[str(page)] = items
    return out
