"""Vector KB build helpers with incremental indexing support."""

from __future__ import annotations

import json
import re
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Dict, Literal, Optional

from .chunking import chunk_text_by_chars, chunk_text_by_tokens, split_markdown_sections
from .embedding import resolve_embedding_runtime
from .index_state import IndexStateStore, sha256_file, sha256_text

BuildMode = Literal["rebuild", "append", "sync"]
ChunkMode = Literal["chars", "tokens"]
WriteMode = Literal["upsert", "add"]


@dataclass
class KBBuildResult:
    """Structured build summary."""

    name: str
    chroma_path: str
    mode: str
    total_docs: int
    indexed_docs: int
    skipped_docs: int
    total_chunks: int
    written_chunks: int
    skipped_chunks: int
    deleted_chunks: int
    failed_docs: int
    elapsed_seconds: float
    embedding_model: str
    state_file: str
    errors: list[str] = field(default_factory=list)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "name": self.name,
            "chroma_path": self.chroma_path,
            "mode": self.mode,
            "total_docs": self.total_docs,
            "indexed_docs": self.indexed_docs,
            "skipped_docs": self.skipped_docs,
            "total_chunks": self.total_chunks,
            "written_chunks": self.written_chunks,
            "skipped_chunks": self.skipped_chunks,
            "deleted_chunks": self.deleted_chunks,
            "failed_docs": self.failed_docs,
            "elapsed_seconds": self.elapsed_seconds,
            "embedding_model": self.embedding_model,
            "state_file": self.state_file,
            "errors": list(self.errors),
        }


@dataclass
class _DiscoveredDocument:
    doc_id: str
    source_path: Path
    text_path: Path
    text: str
    metadata: Dict[str, Any]
    file_hash: str


def _slug(text: str, *, max_len: int = 32) -> str:
    value = re.sub(r"[^a-zA-Z0-9_]+", "_", text.strip().lower())
    value = value.strip("_")
    return value[:max_len] or "section"


def _load_metadata(parsed_json_path: Optional[Path]) -> Dict[str, Any]:
    if parsed_json_path is None or not parsed_json_path.exists():
        return {}
    try:
        payload = json.loads(parsed_json_path.read_text(encoding="utf-8"))
    except Exception:
        return {}
    metadata = payload.get("metadata", {})
    return metadata if isinstance(metadata, dict) else {}


def _read_text_file(path: Path) -> str:
    suffix = path.suffix.lower()
    text = path.read_text(encoding="utf-8")
    if suffix == ".html":
        return re.sub(r"<[^>]+>", " ", text)
    return text


def _discover_documents(input_path: Path) -> list[_DiscoveredDocument]:
    if not input_path.exists():
        raise FileNotFoundError(f"输入目录不存在: {input_path}")

    documents: list[_DiscoveredDocument] = []

    # Structured format: each sub-directory contains *_text.md (+ optional *_parsed.json)
    for entry in sorted(input_path.iterdir()):
        if entry.is_dir():
            md_files = sorted(entry.glob("*_text.md"))
            if not md_files:
                continue
            md_path = md_files[0]
            parsed_files = sorted(entry.glob("*_parsed.json"))
            metadata = _load_metadata(parsed_files[0] if parsed_files else None)
            text = _read_text_file(md_path)
            documents.append(
                _DiscoveredDocument(
                    doc_id=entry.name,
                    source_path=md_path,
                    text_path=md_path,
                    text=text,
                    metadata=metadata,
                    file_hash=sha256_file(md_path),
                )
            )
            continue

        # Flat file input for user convenience: .md/.txt/.html
        if entry.is_file() and entry.suffix.lower() in {".md", ".txt", ".html"}:
            text = _read_text_file(entry)
            documents.append(
                _DiscoveredDocument(
                    doc_id=entry.stem,
                    source_path=entry,
                    text_path=entry,
                    text=text,
                    metadata={"title": entry.stem, "source": str(entry)},
                    file_hash=sha256_file(entry),
                )
            )
    return documents


def _chunk_document(
    document: _DiscoveredDocument,
    *,
    chunk_size: int,
    chunk_overlap: int,
    chunk_mode: ChunkMode,
) -> tuple[str, str, list[Dict[str, Any]], list[str], list[str]]:
    section_items = split_markdown_sections(document.text)
    chunks: list[Dict[str, Any]] = []
    chunk_ids: list[str] = []
    chunk_hashes: list[str] = []
    chunk_index = 0

    for section_idx, (section_title, section_body) in enumerate(section_items):
        if chunk_mode == "tokens":
            section_chunks = chunk_text_by_tokens(
                section_body,
                max_tokens=chunk_size,
                overlap_tokens=max(0, chunk_overlap),
            )
        else:
            section_chunks = chunk_text_by_chars(
                section_body,
                max_chars=chunk_size,
                overlap_chars=max(0, chunk_overlap),
            )

        for local_idx, chunk_text in enumerate(section_chunks):
            chunk_hash = sha256_text(chunk_text)
            section_slug = _slug(section_title)
            chunk_id = f"{_slug(document.doc_id, max_len=40)}:{section_idx}:{section_slug}:{local_idx}"
            page_value = document.metadata.get("pages", 0)
            try:
                page_number = int(page_value)
            except Exception:
                page_number = 0
            metadata = {
                "doc_id": document.doc_id,
                "section": section_title,
                "source": str(document.source_path),
                "page": page_number,
                "chunk_hash": chunk_hash,
                "version": "v1",
                "lang": document.metadata.get("lang", "unknown"),
                "title": document.metadata.get("title", document.doc_id),
                "chunk_index": chunk_index,
            }
            chunks.append({"id": chunk_id, "text": chunk_text, "metadata": metadata})
            chunk_ids.append(chunk_id)
            chunk_hashes.append(chunk_hash)
            chunk_index += 1

    return document.doc_id, document.file_hash, chunks, chunk_ids, chunk_hashes


def _upsert_with_retry(
    *,
    collection: Any,
    ids: list[str],
    documents: list[str],
    metadatas: list[Dict[str, Any]],
    write_mode: WriteMode,
    retry_limit: int,
) -> None:
    for attempt in range(max(1, retry_limit + 1)):
        try:
            if write_mode == "add":
                collection.add(ids=ids, documents=documents, metadatas=metadatas)
            elif hasattr(collection, "upsert"):
                collection.upsert(ids=ids, documents=documents, metadatas=metadatas)
            else:
                collection.delete(ids=ids)
                collection.add(ids=ids, documents=documents, metadatas=metadatas)
            return
        except Exception:
            if attempt + 1 >= max(1, retry_limit + 1):
                raise


def create_vector_kb(
    name: str,
    input_dir: str,
    output_dir: str,
    overwrite: bool = False,
    embedding_model: Optional[str] = "all-MiniLM-L6-v2",
    *,
    chunk_size: int = 800,
    chunk_overlap: int = 80,
    batch_size: int = 100,
    collection_metadata: Optional[Dict[str, Any]] = None,
    mode: BuildMode = "rebuild",
    max_workers: int = 1,
    retry_limit: int = 1,
    return_result: bool = False,
    chunk_mode: ChunkMode = "chars",
    write_mode: WriteMode = "upsert",
    state_file: Optional[str] = None,
    embedding_cache_path: Optional[str] = None,
) -> str | KBBuildResult:
    """
    Build vector KB from markdown/structured docs.

    Compatible with existing usage: returns chroma path string by default.
    """
    started_at = time.time()
    if mode not in {"rebuild", "append", "sync"}:
        raise ValueError(f"mode 无效: {mode}")
    if chunk_mode not in {"chars", "tokens"}:
        raise ValueError(f"chunk_mode 无效: {chunk_mode}")
    if write_mode not in {"upsert", "add"}:
        raise ValueError(f"write_mode 无效: {write_mode}")

    input_path = Path(input_dir)
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)

    chroma_path = output_path / "chroma_store"
    if mode == "rebuild" and chroma_path.exists() and not overwrite:
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        chroma_path = output_path / f"chroma_store_{timestamp}"
    chroma_path.mkdir(parents=True, exist_ok=True)

    documents = _discover_documents(input_path)
    if not documents:
        raise ValueError(
            f"输入目录中未发现可索引文档: {input_path}，需包含目录化 *_text.md 或平铺 .md/.txt/.html 文件。"
        )

    state_path = Path(state_file) if state_file else output_path / f".{name}_index_state.json"
    state_store = IndexStateStore(state_path)
    if mode == "rebuild" and overwrite:
        state_store._state["docs"] = {}

    skipped_docs = 0
    skipped_chunks = 0
    indexed_docs = 0
    deleted_chunks = 0
    failed_docs = 0
    errors: list[str] = []
    removed_chunk_ids: list[str] = []

    pending_docs: list[_DiscoveredDocument] = []
    all_doc_ids = {doc.doc_id for doc in documents}
    if mode == "sync":
        removed = state_store.remove_docs_not_in(all_doc_ids)
        for item in removed.values():
            chunk_ids = item.get("chunk_ids", []) or []
            removed_chunk_ids.extend(str(chunk_id) for chunk_id in chunk_ids)
        deleted_chunks = len(removed_chunk_ids)

    for doc in documents:
        if mode in {"append", "sync"} and state_store.should_skip_doc(doc.doc_id, doc.file_hash):
            skipped_docs += 1
            skipped_chunks += state_store.get_chunk_count(doc.doc_id)
            continue
        pending_docs.append(doc)

    processed_chunks: list[Dict[str, Any]] = []
    state_updates: dict[str, Dict[str, Any]] = {}

    def _run_chunk(doc: _DiscoveredDocument):
        return _chunk_document(doc, chunk_size=chunk_size, chunk_overlap=chunk_overlap, chunk_mode=chunk_mode)

    if pending_docs:
        if max_workers > 1:
            with ThreadPoolExecutor(max_workers=max_workers) as executor:
                for doc_id, file_hash, chunks, chunk_ids, chunk_hashes in executor.map(_run_chunk, pending_docs):
                    indexed_docs += 1
                    processed_chunks.extend(chunks)
                    state_updates[doc_id] = {
                        "file_hash": file_hash,
                        "chunk_ids": chunk_ids,
                        "chunk_hashes": chunk_hashes,
                    }
        else:
            for doc in pending_docs:
                try:
                    doc_id, file_hash, chunks, chunk_ids, chunk_hashes = _run_chunk(doc)
                    indexed_docs += 1
                    processed_chunks.extend(chunks)
                    state_updates[doc_id] = {
                        "file_hash": file_hash,
                        "chunk_ids": chunk_ids,
                        "chunk_hashes": chunk_hashes,
                    }
                except Exception as exc:
                    failed_docs += 1
                    errors.append(f"{doc.doc_id}: {exc}")

    if failed_docs and indexed_docs == 0:
        raise RuntimeError("所有文档分块失败，无法构建向量库。")

    import chromadb

    runtime = resolve_embedding_runtime(
        embedding_model=embedding_model,
        cache_path=embedding_cache_path,
    )
    client = chromadb.PersistentClient(path=str(chroma_path))

    collection_meta = {"description": f"kbase vector collection: {name}"}
    if collection_metadata:
        collection_meta.update(collection_metadata)

    if mode == "rebuild" and overwrite:
        try:
            client.delete_collection(name)
        except Exception:
            pass

    collection = client.get_or_create_collection(
        name=name,
        metadata=collection_meta,
        embedding_function=runtime.embedding_function,
    )

    if mode == "sync" and removed_chunk_ids:
        collection.delete(ids=removed_chunk_ids)

    written_chunks = 0
    if processed_chunks:
        for idx in range(0, len(processed_chunks), max(1, batch_size)):
            batch = processed_chunks[idx : idx + max(1, batch_size)]
            ids = [item["id"] for item in batch]
            texts = [item["text"] for item in batch]
            metadatas = [item["metadata"] for item in batch]
            try:
                _upsert_with_retry(
                    collection=collection,
                    ids=ids,
                    documents=texts,
                    metadatas=metadatas,
                    write_mode=write_mode,
                    retry_limit=retry_limit,
                )
                written_chunks += len(batch)
            except Exception as exc:
                errors.append(f"batch {idx // max(1, batch_size)} 写入失败: {exc}")
                raise RuntimeError(f"Chroma 写入失败: {exc}") from exc

    for doc in pending_docs:
        update = state_updates.get(doc.doc_id)
        if not update:
            continue
        state_store.update_doc(
            doc.doc_id,
            source=str(doc.source_path),
            file_hash=update["file_hash"],
            chunk_ids=update["chunk_ids"],
            chunk_hashes=update["chunk_hashes"],
            model_version=runtime.model_id,
        )
    state_store.save()
    runtime.close()

    elapsed = round(time.time() - started_at, 4)
    result = KBBuildResult(
        name=name,
        chroma_path=str(chroma_path),
        mode=mode,
        total_docs=len(documents),
        indexed_docs=indexed_docs,
        skipped_docs=skipped_docs,
        total_chunks=written_chunks + skipped_chunks,
        written_chunks=written_chunks,
        skipped_chunks=skipped_chunks,
        deleted_chunks=deleted_chunks,
        failed_docs=failed_docs,
        elapsed_seconds=elapsed,
        embedding_model=runtime.model_id,
        state_file=str(state_path),
        errors=errors,
    )

    if return_result:
        return result
    return str(chroma_path)


__all__ = ["KBBuildResult", "create_vector_kb"]
