"""Vector memory over the OpenSearch k-NN index."""
import json
import uuid
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

from ..stores import opensearch
from ..stores.embeddings import embed, embed_batch, provider_name


def _now() -> str:
    return datetime.now(timezone.utc).isoformat()


def _memory_index() -> str:
    return f"{opensearch.settings.os_index_prefix}-memory"


async def remember(text: str, metadata: Optional[dict] = None) -> dict:
    text = (text or "").strip()
    if not text:
        raise ValueError("text must not be empty")
    doc_id = str(uuid.uuid4())
    doc = {
        "doc_id": doc_id,
        "text": text[:8000],
        "metadata": metadata or {},
        "embedding": embed(text),
        "created_at": _now(),
    }
    await opensearch.store.index_doc("memory", doc)
    return {"doc_id": doc_id, "provider": provider_name()}


async def remember_batch(texts: List[str], metadata: Optional[dict] = None) -> List[dict]:
    texts = [t.strip()[:8000] for t in texts if t and t.strip()]
    if not texts:
        return []
    vectors = embed_batch(texts)
    docs = []
    for text, vector in zip(texts, vectors):
        docs.append({
            "doc_id": str(uuid.uuid4()),
            "text": text,
            "metadata": metadata or {},
            "embedding": vector,
            "created_at": _now(),
        })
    lines: List[str] = []
    for d in docs:
        lines.append(json.dumps({"index": {"_index": _memory_index(), "_id": d["doc_id"]}}))
        lines.append(json.dumps(d))
    res = await opensearch.store.client().post(
        "/_bulk",
        content=("\n".join(lines) + "\n").encode(),
        headers={"Content-Type": "application/x-ndjson"},
    )
    res.raise_for_status()
    await opensearch.store.refresh("memory")
    return [{"doc_id": d["doc_id"]} for d in docs]


async def recall(query: str, k: int = 8) -> List[dict]:
    query = (query or "").strip()
    if not query:
        return []
    k = max(1, min(k, 50))
    body = {
        "size": k,
        "query": {"knn": {"embedding": {"vector": embed(query), "k": k}}},
        "_source": ["text", "metadata", "created_at"],
    }
    return await opensearch.store.search("memory", body)