from __future__ import annotations

import json
import logging
import os
import shutil
import subprocess

from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

logger = logging.getLogger(__name__)

_MAX_EXTRACT_WORKERS = min(8, os.cpu_count() or 4)

from rak.bm25 import BM25Index
from rak.embedder import Embedder
from rak.errors import EmptyLibraryError, ZotNotFoundError
from rak.pdf import chunk_text, extract_file_text, find_attachments
from rak.registry import compute_hash
from rak.store import VectorStore


def parse_zot_items(raw_json: str) -> list[dict]:
    return json.loads(raw_json)


def build_document_text(item: dict, pdf_text: str = "") -> str:
    parts = [item.get("title", "")]
    creators = item.get("creators", [])
    if creators:
        author_names = []
        for c in creators:
            name_parts = [c.get("first_name", ""), c.get("last_name", "")]
            author_names.append(" ".join(p for p in name_parts if p))
        parts.append("Authors: " + ", ".join(author_names))
    abstract = item.get("abstract")
    if abstract:
        parts.append(abstract)
    tags = item.get("tags", [])
    if tags:
        parts.append("Tags: " + ", ".join(tags))
    if pdf_text:
        parts.append(pdf_text)
    return "\n".join(parts)


def _extract_item_text(
    item: dict, storage_dir: Path | None, pdf_provider: str,
) -> tuple[str, int, int]:
    """Extract attachment text for a single item (thread-safe).

    Returns (attachment_text, extraction_attempts, extraction_failures).
    """
    key = item.get("key", "")
    if not key or not storage_dir:
        return "", 0, 0
    attachments = find_attachments(storage_dir, key)
    if not attachments:
        return "", 0, 0
    attempts = len(attachments)
    failures = 0
    texts = []
    for p in attachments:
        t = extract_file_text(p, provider=pdf_provider)
        if t:
            texts.append(t)
        else:
            failures += 1
    return "\n\n".join(texts), attempts, failures


def diff_items(
    items: list[dict], registry: dict[str, str],
    storage_dir: Path | None = None,
    pdf_provider: str = "pymupdf",
) -> tuple[list[dict], list[dict], list[str], dict[str, str]]:
    """Compare items against registry to find changes.

    Returns (to_add, to_update, to_remove, text_cache) where text_cache
    maps item keys to their built document text, avoiding redundant PDF extraction.
    """
    to_add = []
    to_update = []
    text_cache: dict[str, str] = {}
    fetched_keys = set()
    extraction_attempts = 0
    extraction_failures = 0

    # Phase 1: parallel PDF/attachment extraction
    attachment_texts: dict[str, str] = {}
    valid_items = [(item, item.get("key", "")) for item in items]
    valid_items = [(item, key) for item, key in valid_items if key]

    if storage_dir:
        with ThreadPoolExecutor(max_workers=_MAX_EXTRACT_WORKERS) as executor:
            future_to_key = {
                executor.submit(_extract_item_text, item, storage_dir, pdf_provider): key
                for item, key in valid_items
            }
            for future in as_completed(future_to_key):
                key = future_to_key[future]
                att_text, att_attempts, att_failures = future.result()
                extraction_attempts += att_attempts
                extraction_failures += att_failures
                if att_text:
                    attachment_texts[key] = att_text

    # Phase 2: build documents and diff against registry
    for item, key in valid_items:
        fetched_keys.add(key)
        text = build_document_text(item, pdf_text=attachment_texts.get(key, ""))
        if not text.strip():
            continue
        text_cache[key] = text
        content_hash = compute_hash(text)
        if key not in registry:
            to_add.append(item)
        elif registry[key] != content_hash:
            to_update.append(item)

    if extraction_attempts > 0 and extraction_failures > 0:
        rate = extraction_failures / extraction_attempts
        if rate > 0.5:
            logger.warning(
                "%d of %d file extractions failed (%.0f%%). Check PDF provider '%s'.",
                extraction_failures, extraction_attempts, rate * 100, pdf_provider,
            )
        elif extraction_failures > 0:
            logger.info("%d of %d file extractions failed.", extraction_failures, extraction_attempts)
    to_remove = [k for k in registry if k not in fetched_keys]
    return to_add, to_update, to_remove, text_cache


def fetch_zot_items(zot_command: str = "zot", limit: int = 5000) -> list[dict]:
    if not shutil.which(zot_command):
        raise ZotNotFoundError(zot_command)
    result = subprocess.run(
        [zot_command, "--json", "--limit", str(limit), "list"],
        capture_output=True, text=True, encoding="utf-8",
    )
    if result.returncode != 0:
        raise RuntimeError(f"zot command failed: {result.stderr}")
    items = parse_zot_items(result.stdout)
    if not items:
        raise EmptyLibraryError()
    return items


def index_items(
    items: list[dict],
    embedder: Embedder,
    vector_store: VectorStore,
    bm25_index: BM25Index,
    on_progress: callable | None = None,
    registry: dict[str, str] | None = None,
    storage_dir: Path | None = None,
    chunk_size: int = 512,
    chunk_overlap: int = 64,
    pdf_provider: str = "pymupdf",
) -> dict:
    """Index items and return a result dict.

    Returns dict with keys: added, updated, removed, unchanged, registry, text_cache.
    For full indexing (registry=None), added=count and updated/removed/unchanged=0.
    """
    if registry is not None:
        return _index_incremental(items, embedder, vector_store, bm25_index, on_progress, registry, storage_dir, chunk_size, chunk_overlap, pdf_provider)
    return _index_full(items, embedder, vector_store, bm25_index, on_progress, storage_dir, chunk_size, chunk_overlap, pdf_provider)


def _delete_chunks(vector_store: VectorStore, key: str) -> None:
    """Delete a document and any associated chunks from the vector store."""
    # Try deleting the base document
    vector_store.delete([key])
    # Delete any chunks (query by parent_key metadata)
    try:
        results = vector_store.get_ids_by_metadata(where={"parent_key": key})
        if results["ids"]:
            vector_store.delete(results["ids"])
    except Exception as exc:
        logger.debug("Failed to delete chunks for %s: %s", key, exc)


def _build_metadata(item: dict) -> dict:
    collections = item.get("collections", [])
    tags = item.get("tags", [])
    metadata = {
        "title": item.get("title", ""),
        "date": item.get("date", ""),
        "item_type": item.get("item_type", ""),
    }
    creators = item.get("creators", [])
    if creators:
        author_parts = [
            " ".join(p for p in [c.get("first_name", ""), c.get("last_name", "")] if p)
            for c in creators
        ]
        metadata["authors"] = ", ".join(author_parts)
    if collections:
        metadata["collections"] = collections
    if tags:
        metadata["tags"] = tags
    return metadata


def _embed_and_store_batch(
    batch_ids: list[str],
    batch_texts: list[str],
    batch_metadatas: list[dict],
    embedder: Embedder,
    vector_store: VectorStore,
) -> None:
    """Embed a batch of texts and store them in the vector store."""
    if not batch_texts:
        return
    embeddings = embedder.embed_batch(batch_texts)
    vector_store.add(ids=batch_ids, embeddings=embeddings, documents=batch_texts, metadatas=batch_metadatas)


def _index_full(
    items: list[dict],
    embedder: Embedder,
    vector_store: VectorStore,
    bm25_index: BM25Index,
    on_progress: callable | None = None,
    storage_dir: Path | None = None,
    chunk_size: int = 512,
    chunk_overlap: int = 64,
    pdf_provider: str = "pymupdf",
) -> dict:
    """Index all items from scratch. Returns result dict."""
    count = 0
    text_cache: dict[str, str] = {}
    batch_ids: list[str] = []
    batch_texts: list[str] = []
    batch_metadatas: list[dict] = []
    embed_batch_size = embedder.batch_size
    extraction_attempts = 0
    extraction_failures = 0

    # Phase 1: parallel PDF/attachment extraction
    valid_items = [(item, item.get("key", "")) for item in items]
    valid_items = [(item, key) for item, key in valid_items if key]

    attachment_texts: dict[str, str] = {}
    if storage_dir:
        with ThreadPoolExecutor(max_workers=_MAX_EXTRACT_WORKERS) as executor:
            future_to_key = {
                executor.submit(_extract_item_text, item, storage_dir, pdf_provider): key
                for item, key in valid_items
            }
            for future in as_completed(future_to_key):
                key = future_to_key[future]
                att_text, att_attempts, att_failures = future.result()
                extraction_attempts += att_attempts
                extraction_failures += att_failures
                if att_text:
                    attachment_texts[key] = att_text

    # Phase 2: sequential embedding and indexing
    for i, (item, key) in enumerate(valid_items):
        text = build_document_text(item, attachment_texts.get(key, ""))
        if not text.strip():
            continue
        text_cache[key] = text
        metadata = _build_metadata(item)
        # Store full document for BM25
        bm25_index.add(key, text)
        # Accumulate chunks for batch embedding
        chunks = chunk_text(text, chunk_size=chunk_size, overlap=chunk_overlap)
        if len(chunks) <= 1:
            batch_ids.append(key)
            batch_texts.append(text)
            batch_metadatas.append(metadata)
        else:
            for ci, chunk in enumerate(chunks):
                chunk_id = f"{key}_chunk_{ci}"
                chunk_meta = {**metadata, "parent_key": key, "chunk_index": ci}
                batch_ids.append(chunk_id)
                batch_texts.append(chunk)
                batch_metadatas.append(chunk_meta)
        # Flush batch when full
        if len(batch_texts) >= embed_batch_size:
            _embed_and_store_batch(batch_ids, batch_texts, batch_metadatas, embedder, vector_store)
            batch_ids, batch_texts, batch_metadatas = [], [], []
        count += 1
        if on_progress:
            on_progress(i + 1, len(valid_items))

    # Flush remaining
    _embed_and_store_batch(batch_ids, batch_texts, batch_metadatas, embedder, vector_store)
    if extraction_attempts > 0 and extraction_failures > 0:
        rate = extraction_failures / extraction_attempts
        if rate > 0.5:
            logger.warning(
                "%d of %d file extractions failed (%.0f%%). Check PDF provider '%s'.",
                extraction_failures, extraction_attempts, rate * 100, pdf_provider,
            )
        elif extraction_failures > 0:
            logger.info("%d of %d file extractions failed.", extraction_failures, extraction_attempts)
    return {
        "added": count,
        "updated": 0,
        "removed": 0,
        "unchanged": 0,
        "registry": {k: compute_hash(v) for k, v in text_cache.items()},
        "text_cache": text_cache,
    }


def _index_incremental(
    items: list[dict],
    embedder: Embedder,
    vector_store: VectorStore,
    bm25_index: BM25Index,
    on_progress: callable | None = None,
    registry: dict[str, str] | None = None,
    storage_dir: Path | None = None,
    chunk_size: int = 512,
    chunk_overlap: int = 64,
    pdf_provider: str = "pymupdf",
) -> dict[str, int]:
    to_add, to_update, to_remove, text_cache = diff_items(items, registry, storage_dir=storage_dir, pdf_provider=pdf_provider)
    new_registry = dict(registry)
    added = 0
    updated = 0

    work_items = [(item, "add") for item in to_add] + [(item, "update") for item in to_update]
    batch_ids: list[str] = []
    batch_texts: list[str] = []
    batch_metadatas: list[dict] = []
    embed_batch_size = embedder.batch_size

    for i, (item, action) in enumerate(work_items):
        key = item["key"]
        text = text_cache.get(key, "")
        if not text.strip():
            continue
        metadata = _build_metadata(item)
        if action == "update":
            bm25_index.delete(key)
            _delete_chunks(vector_store, key)
        bm25_index.add(key, text)
        chunks = chunk_text(text, chunk_size=chunk_size, overlap=chunk_overlap)
        if len(chunks) <= 1:
            batch_ids.append(key)
            batch_texts.append(text)
            batch_metadatas.append(metadata)
        else:
            for ci, chunk in enumerate(chunks):
                chunk_id = f"{key}_chunk_{ci}"
                chunk_meta = {**metadata, "parent_key": key, "chunk_index": ci}
                batch_ids.append(chunk_id)
                batch_texts.append(chunk)
                batch_metadatas.append(chunk_meta)
        if len(batch_texts) >= embed_batch_size:
            _embed_and_store_batch(batch_ids, batch_texts, batch_metadatas, embedder, vector_store)
            batch_ids, batch_texts, batch_metadatas = [], [], []
        new_registry[key] = compute_hash(text)
        if action == "add":
            added += 1
        else:
            updated += 1
        if on_progress:
            on_progress(i + 1, len(work_items))

    _embed_and_store_batch(batch_ids, batch_texts, batch_metadatas, embedder, vector_store)

    for key in to_remove:
        _delete_chunks(vector_store, key)
        bm25_index.delete(key)
        del new_registry[key]

    return {
        "added": added,
        "updated": updated,
        "removed": len(to_remove),
        "unchanged": len(items) - len(to_add) - len(to_update),
        "registry": new_registry,
        "text_cache": text_cache,
    }
