"""vault/_embeddings.jsonl · chunk 向量 sidecar（lazy write）。"""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any


def chunk_embed_key(vault_id: str, chunk_index: int) -> str:
    return f"{vault_id}:{int(chunk_index)}"


class VaultEmbedIndex:
    def __init__(self, vault_root: Path) -> None:
        self._path = vault_root / "_embeddings.jsonl"
        self._cache: dict[str, dict[str, Any]] | None = None

    @property
    def path(self) -> Path:
        return self._path

    def _load(self) -> dict[str, dict[str, Any]]:
        if self._cache is not None:
            return self._cache
        out: dict[str, dict[str, Any]] = {}
        if self._path.is_file():
            for line in self._path.read_text(encoding="utf-8").splitlines():
                line = line.strip()
                if not line:
                    continue
                try:
                    row = json.loads(line)
                except json.JSONDecodeError:
                    continue
                key = str(row.get("key") or "")
                if key:
                    out[key] = row
        self._cache = out
        return out

    def get_vector(self, key: str) -> list[float] | None:
        row = self._load().get(key)
        if not row:
            return None
        vec = row.get("vector")
        if isinstance(vec, list) and vec:
            return [float(x) for x in vec]
        return None

    def upsert_lazy(
        self,
        *,
        key: str,
        vault_id: str,
        chunk_index: int,
        vector: list[float],
        text_preview: str = "",
    ) -> None:
        if key in self._load():
            return
        self._path.parent.mkdir(parents=True, exist_ok=True)
        row = {
            "key": key,
            "vault_id": vault_id,
            "chunk_index": int(chunk_index),
            "vector": vector,
            "text_preview": str(text_preview)[:120],
        }
        with self._path.open("a", encoding="utf-8") as f:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")
        self._load()[key] = row

    def list_keys_for_vault(self, vault_id: str) -> list[str]:
        return [k for k, row in self._load().items() if str(row.get("vault_id")) == vault_id]

    def dream_batch_backfill_stub(self, *, limit: int = 50) -> int:
        """夜间补全占位：返回已索引条数（真实 embed 由 worker 注入）。"""
        return min(len(self._load()), limit)
