"""Index downloaded PMC articles into a ChromaDB vector store.

Documents are read from the JSON files produced by ``pubmed-stream`` (or any
compatible source), split into overlapping text chunks, and upserted into a
persistent ChromaDB collection.  Already-indexed chunks are silently skipped
on subsequent runs (ChromaDB upsert semantics).
"""

from __future__ import annotations

import json
import logging
import re
from pathlib import Path
from typing import Any, Callable, List, Optional

from .types import IndexStats

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Defaults
# ---------------------------------------------------------------------------

DEFAULT_CHUNK_SIZE = 2000    # characters ≈ 500 tokens  (1 token ≈ 4 chars)
DEFAULT_CHUNK_OVERLAP = 200  # character overlap between adjacent chunks


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------


def _chunk_text(text: str, chunk_size: int, overlap: int) -> List[str]:
    """Split *text* into overlapping chunks of at most *chunk_size* characters.

    Tries to break on sentence boundaries (``. ``) within the last half of
    each window so chunks remain semantically coherent.
    """
    if not text:
        return []

    chunks: List[str] = []
    start = 0

    while start < len(text):
        end = min(start + chunk_size, len(text))

        # Prefer a sentence boundary in the second half of the window
        if end < len(text):
            boundary = text.rfind(". ", start + chunk_size // 2, end)
            if boundary != -1:
                end = boundary + 1  # include the period

        chunk = text[start:end].strip()
        if chunk:
            chunks.append(chunk)

        if end >= len(text):
            break
        start = end - overlap

    return chunks


def _safe_str(value: Any, max_len: int = 500) -> str:
    """Convert *value* to a non-None string, truncated to *max_len* chars."""
    return str(value or "")[:max_len]


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def index_directory(
    papers_dir: str | Path,
    db_path: str | Path,
    collection_name: str = "publications",
    chunk_size: int = DEFAULT_CHUNK_SIZE,
    chunk_overlap: int = DEFAULT_CHUNK_OVERLAP,
    embedding_function: Optional[Callable] = None,
    batch_size: int = 100,
    recursive: bool = True,
) -> IndexStats:
    """Read every ``*.json`` file under *papers_dir* and upsert its text into ChromaDB.

    Each article is split into overlapping text chunks. Chunks are upserted
    with their PMC metadata so they can be cited by the RAG pipeline.

    Parameters
    ----------
    papers_dir:
        Root directory of downloaded JSON files (as produced by *pubmed-stream*).
    db_path:
        Filesystem path for the persistent ChromaDB database.
    collection_name:
        Name of the ChromaDB collection (default: ``"publications"``).
    chunk_size:
        Characters per chunk (default 2 000 ≈ 500 tokens).
    chunk_overlap:
        Character overlap between consecutive chunks (default: 200).
    embedding_function:
        A ChromaDB-compatible embedding function.  Defaults to
        ``DefaultEmbeddingFunction`` (``all-MiniLM-L6-v2``).  Pass a custom
        function to use biomedical embeddings such as MedEmbed.
    batch_size:
        Documents per ChromaDB ``upsert`` call.  Larger batches are faster
        but use more RAM.
    recursive:
        Search *papers_dir* recursively for ``*.json`` files (default: ``True``).

    Returns
    -------
    :class:`~neuromind.alzheimers.types.IndexStats`
    """
    try:
        import chromadb
        from chromadb.utils.embedding_functions import DefaultEmbeddingFunction
    except ImportError as exc:
        raise ImportError(
            "chromadb is required for indexing.\n"
            "Install with:  pip install chromadb"
        ) from exc

    papers_dir = Path(papers_dir)
    db_path = Path(db_path)
    db_path.mkdir(parents=True, exist_ok=True)

    ef = embedding_function if embedding_function is not None else DefaultEmbeddingFunction()

    chroma = chromadb.PersistentClient(path=str(db_path))
    collection = chroma.get_or_create_collection(
        name=collection_name,
        embedding_function=ef,
    )

    pattern = "**/*.json" if recursive else "*.json"
    json_files = sorted(papers_dir.glob(pattern))
    logger.info("Found %d JSON files under %s", len(json_files), papers_dir)

    stats = IndexStats(total_files=len(json_files))

    # Accumulator for batch upserts
    docs_buf:  List[str]  = []
    metas_buf: List[dict] = []
    ids_buf:   List[str]  = []

    def _flush() -> None:
        if not docs_buf:
            return
        collection.upsert(documents=docs_buf, metadatas=metas_buf, ids=ids_buf)
        stats.total_chunks += len(docs_buf)
        docs_buf.clear()
        metas_buf.clear()
        ids_buf.clear()

    for file_path in json_files:
        try:
            with open(file_path, encoding="utf-8") as fh:
                article = json.load(fh)
        except (json.JSONDecodeError, OSError) as exc:
            logger.warning("Skipping %s: %s", file_path.name, exc)
            stats.failed += 1
            continue

        pmcid: str = article.get("pmcid") or file_path.stem
        text: str = article.get("text", "") or ""

        if not text.strip():
            logger.debug("Skipping %s: no text content", pmcid)
            stats.skipped += 1
            continue

        meta_raw = article.get("metadata") or {}

        # Build a flat metadata dict (ChromaDB requires string/int/float values)
        base_meta: dict = {
            "pmcid":    _safe_str(pmcid),
            "pmc_id":   re.sub(r"^PMC", "", _safe_str(pmcid)),
            "title":    _safe_str(meta_raw.get("title"), 500),
            "journal":  _safe_str(meta_raw.get("journal"), 200),
            "year":     _safe_str(meta_raw.get("year"), 10),
            "doi":      _safe_str(meta_raw.get("doi"), 200),
            "authors":  "; ".join(meta_raw.get("authors") or [])[:500],
            "abstract": _safe_str(meta_raw.get("abstract"), 1000),
            "keywords": "; ".join(meta_raw.get("keywords") or [])[:300],
        }

        chunks = _chunk_text(text, chunk_size, chunk_overlap)
        if not chunks:
            stats.skipped += 1
            continue

        for idx, chunk in enumerate(chunks):
            chunk_id = f"{pmcid}_chunk_{idx}"
            docs_buf.append(chunk)
            metas_buf.append({**base_meta, "chunk_index": idx, "total_chunks": len(chunks)})
            ids_buf.append(chunk_id)

            if len(docs_buf) >= batch_size:
                _flush()

        stats.indexed += 1

    _flush()

    logger.info(
        "Indexing complete — indexed=%d skipped=%d failed=%d chunks=%d",
        stats.indexed,
        stats.skipped,
        stats.failed,
        stats.total_chunks,
    )
    return stats
