"""ChromaDB memory store — typed wrapper for vector memory storage."""

import json
import logging
import os
from typing import Optional

import chromadb

from neo_cortex.models import (
    Activity,
    CortexStats,
    MemoryRecord,
    MemoryType,
)

logger = logging.getLogger(__name__)


class MemoryStore:
    """ChromaDB vector store. Reads/writes MemoryRecord objects."""

    def __init__(self, db_path: str, collection_name: str = "conversations_v2"):
        os.makedirs(db_path, exist_ok=True)
        self._client = chromadb.PersistentClient(path=db_path)
        self._collection = self._client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},
        )

    def insert(self, record: MemoryRecord, embedding: list[float]) -> str:
        """Insert a memory with its embedding vector."""
        self._collection.add(
            ids=[record.id],
            embeddings=[embedding],
            documents=[record.document],
            metadatas=[{
                "session_id": record.session_id,
                "timestamp": record.timestamp,
                "turn_number": record.turn_number,
                "question": record.question[:500],
                "answer_preview": record.answer_preview[:500],
                "tools_used": json.dumps(record.tools_used),
                "model": record.model,
                "energy": record.energy,
                "source": record.source,
                "project": record.project,
                "topic": record.topic,
                "activity": record.activity.value,
                "type": record.memory_type.value,
                "observation_level": record.observation_level.value if hasattr(record.observation_level, "value") else "direct",
                "evidence_ids": ",".join(record.evidence_ids) if record.evidence_ids else "",
            }],
        )
        return record.id

    def query(
        self,
        embedding: list[float],
        n: int = 5,
        where: Optional[dict] = None,
    ) -> list[tuple[MemoryRecord, float]]:
        """Vector similarity search. Returns (record, similarity) pairs."""
        kwargs: dict = {
            "query_embeddings": [embedding],
            "n_results": min(n, self._collection.count() or 1),
        }
        if where:
            try:
                kwargs["where"] = where
                results = self._collection.query(**kwargs)
            except Exception:
                # Filter error — retry without filter
                kwargs.pop("where", None)
                results = self._collection.query(**kwargs)
        else:
            results = self._collection.query(**kwargs)

        return self._parse_results(results)

    def get_all_metadata(self) -> list[dict]:
        """Get all metadata dicts (for stats, dream, timeline)."""
        result = self._collection.get(include=["metadatas"])
        ids = result["ids"]
        metadatas = result["metadatas"] or []
        out = []
        for i, meta in enumerate(metadatas):
            row = dict(meta)
            row["id"] = ids[i]
            out.append(row)
        return out

    def update_energy(self, memory_id: str, energy: float) -> None:
        """Update a memory's energy value."""
        existing = self._collection.get(ids=[memory_id], include=["metadatas"])
        if existing["ids"] and existing["metadatas"]:
            meta = dict(existing["metadatas"][0])
            meta["energy"] = energy
            self._collection.update(ids=[memory_id], metadatas=[meta])

    def update(self, memory_id: str, record: MemoryRecord, embedding: list[float]) -> None:
        """Update a memory's content, embedding, and metadata."""
        self._collection.update(
            ids=[memory_id],
            embeddings=[embedding],
            documents=[record.document],
            metadatas=[{
                "session_id": record.session_id,
                "timestamp": record.timestamp,
                "turn_number": record.turn_number,
                "question": record.question[:500],
                "answer_preview": record.answer_preview[:500],
                "tools_used": json.dumps(record.tools_used),
                "model": record.model,
                "energy": record.energy,
                "source": record.source,
                "project": record.project,
                "topic": record.topic,
                "activity": record.activity.value,
                "type": record.memory_type.value,
            }],
        )

    def delete(self, memory_id: str) -> None:
        """Delete a memory from ChromaDB."""
        self._collection.delete(ids=[memory_id])

    def clear(self) -> int:
        """Delete all memories. Returns count of deleted items."""
        n = self._collection.count()
        if n > 0:
            all_ids = self._collection.get()["ids"]
            self._collection.delete(ids=all_ids)
        logger.info("MemoryStore.clear: deleted %d memories", n)
        return n

    def count(self) -> int:
        return self._collection.count()

    def stats(self, dream_count: int) -> CortexStats:
        """Compute cortex statistics from all memories."""
        all_meta = self.get_all_metadata()
        if not all_meta:
            return CortexStats(
                total_memories=0, sessions=0,
                avg_energy=0, max_energy=0, min_energy=0,
                dream_count=dream_count, projects={},
            )

        energies = [m.get("energy", 1.0) for m in all_meta]
        sessions = {m.get("session_id", "") for m in all_meta}
        projects: dict[str, int] = {}
        for m in all_meta:
            p = m.get("project", "general")
            projects[p] = projects.get(p, 0) + 1

        return CortexStats(
            total_memories=len(all_meta),
            sessions=len(sessions),
            avg_energy=round(sum(energies) / len(energies), 3),
            max_energy=round(max(energies), 3),
            min_energy=round(min(energies), 3),
            dream_count=dream_count,
            projects=projects,
        )

    def timeline(self, limit: int = 10, project: Optional[str] = None) -> list[MemoryRecord]:
        """Get memories ordered by timestamp (newest first)."""
        all_meta = self.get_all_metadata()
        if project:
            all_meta = [m for m in all_meta if m.get("project", "general") == project]

        all_meta.sort(key=lambda m: m.get("timestamp", 0), reverse=True)

        return [self._meta_to_record(m) for m in all_meta[:limit]]

    # --- Internal helpers ---

    def _meta_to_record(self, meta: dict) -> MemoryRecord:
        """Convert metadata dict to MemoryRecord."""
        tools = meta.get("tools_used", "[]")
        if isinstance(tools, str):
            try:
                tools = json.loads(tools)
            except Exception:
                tools = []

        activity_raw = meta.get("activity", "discussion")
        try:
            activity = Activity(activity_raw)
        except ValueError:
            activity = Activity.DISCUSSION

        type_raw = meta.get("type", "episodic")
        try:
            memory_type = MemoryType(type_raw)
        except ValueError:
            memory_type = MemoryType.EPISODIC

        raw_evidence = meta.get("evidence_ids", "")
        evidence_ids = [e for e in raw_evidence.split(",") if e] if raw_evidence else []

        return MemoryRecord(
            id=meta.get("id", ""),
            session_id=meta.get("session_id", ""),
            timestamp=meta.get("timestamp", 0),
            turn_number=meta.get("turn_number", 0),
            question=meta.get("question", ""),
            answer_preview=meta.get("answer_preview", ""),
            document=meta.get("document", f"Q: {meta.get('question', '')}\nA: {meta.get('answer_preview', '')}"),
            project=meta.get("project", "general"),
            topic=meta.get("topic", "unknown"),
            activity=activity,
            memory_type=memory_type,
            energy=meta.get("energy", 1.0),
            model=meta.get("model", "opus"),
            source=meta.get("source", "neo-cortex"),
            tools_used=tools,
            evidence_ids=evidence_ids,
        )

    def _parse_results(self, results: dict) -> list[tuple[MemoryRecord, float]]:
        """Parse ChromaDB query results into (MemoryRecord, similarity) pairs."""
        pairs = []
        if not results["ids"] or not results["ids"][0]:
            return pairs

        ids = results["ids"][0]
        distances = results["distances"][0] if results.get("distances") else [0] * len(ids)
        metadatas = results["metadatas"][0] if results.get("metadatas") else [{}] * len(ids)

        for i, mid in enumerate(ids):
            meta = dict(metadatas[i]) if metadatas[i] else {}
            meta["id"] = mid
            similarity = 1.0 - distances[i]  # cosine distance → similarity
            record = self._meta_to_record(meta)
            pairs.append((record, similarity))
        return pairs
