"""ChromaDB memory store — typed wrapper over conversations_v2 collection."""

import json
import logging
from typing import Optional

import chromadb
from chromadb.config import Settings

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

logger = logging.getLogger(__name__)


class MemoryStore:
    """Typed ChromaDB wrapper. Reads/writes MemoryRecord objects."""

    def __init__(self, db_path: str, collection_name: str = "conversations_v2"):
        self._client = chromadb.PersistentClient(
            path=db_path,
            settings=Settings(
                anonymized_telemetry=False,
                chroma_server_nofile=None,
            ),
        )
        try:
            self._collection = self._client.get_collection(collection_name)
        except Exception:
            self._collection = self._client.create_collection(
                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=[self._record_to_meta(record)],
        )
        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": n,
        }
        if where:
            kwargs["where"] = where

        try:
            results = self._collection.query(**kwargs)
        except Exception:
            # Filter error — retry without where clause
            if where:
                kwargs.pop("where")
                results = self._collection.query(**kwargs)
            else:
                raise

        return self._parse_results(results)

    def get_all_ids(self) -> list[str]:
        """Get all memory IDs."""
        data = self._collection.get()
        return data["ids"]

    def get_all_metadata(self) -> list[dict]:
        """Get all metadata dicts (for stats, dream, timeline)."""
        data = self._collection.get()
        if not data["ids"]:
            return []
        return [
            {"id": data["ids"][i], **data["metadatas"][i]}
            for i in range(len(data["ids"]))
        ]

    def update_energy(self, memory_id: str, energy: float) -> None:
        """Update a memory's energy value."""
        self._collection.update(
            ids=[memory_id],
            metadatas=[{"energy": energy}],
        )

    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 _record_to_meta(self, record: MemoryRecord) -> dict:
        """Convert MemoryRecord to ChromaDB metadata dict."""
        return {
            "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 _meta_to_record(self, meta: dict) -> MemoryRecord:
        """Convert ChromaDB 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

        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=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,
        )

    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

        for i, mem_id in enumerate(results["ids"][0]):
            meta = results["metadatas"][0][i]
            meta["id"] = mem_id
            dist = results["distances"][0][i]
            similarity = 1.0 - dist  # cosine distance → similarity
            record = self._meta_to_record(meta)
            pairs.append((record, similarity))

        return pairs
