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

import json
import logging
import os
from typing import Optional

import lancedb
import pyarrow as pa

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

logger = logging.getLogger(__name__)

# LanceDB schema — all metadata stored alongside vectors
_SCHEMA = pa.schema([
    pa.field("id", pa.utf8()),
    pa.field("vector", pa.list_(pa.float32(), 1024)),
    pa.field("document", pa.utf8()),
    pa.field("session_id", pa.utf8()),
    pa.field("timestamp", pa.float64()),
    pa.field("turn_number", pa.int32()),
    pa.field("question", pa.utf8()),
    pa.field("answer_preview", pa.utf8()),
    pa.field("tools_used", pa.utf8()),
    pa.field("model", pa.utf8()),
    pa.field("energy", pa.float64()),
    pa.field("source", pa.utf8()),
    pa.field("project", pa.utf8()),
    pa.field("topic", pa.utf8()),
    pa.field("activity", pa.utf8()),
    pa.field("type", pa.utf8()),
])


class MemoryStore:
    """LanceDB 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._db = lancedb.connect(db_path)
        self._table_name = collection_name
        try:
            self._table = self._db.open_table(collection_name)
        except Exception:
            self._table = self._db.create_table(
                collection_name, schema=_SCHEMA,
            )

    def insert(self, record: MemoryRecord, embedding: list[float]) -> str:
        """Insert a memory with its embedding vector."""
        self._table.add([{
            "id": record.id,
            "vector": embedding,
            "document": record.document,
            "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,
        }])
        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."""
        q = self._table.search(embedding).metric("cosine").limit(n)

        if where:
            sql_filter = self._where_to_sql(where)
            if sql_filter:
                try:
                    q = q.where(sql_filter)
                except Exception:
                    # Filter error — retry without filter
                    q = self._table.search(embedding).metric("cosine").limit(n)

        results = q.to_list()
        return self._parse_results(results)

    def get_all_ids(self) -> list[str]:
        """Get all memory IDs."""
        return self._table.to_arrow().column("id").to_pylist()

    def get_all_metadata(self) -> list[dict]:
        """Get all metadata dicts (for stats, dream, timeline)."""
        arrow = self._table.to_arrow().select([
            "id", "session_id", "timestamp", "turn_number",
            "question", "answer_preview", "tools_used", "model",
            "energy", "source", "project", "topic", "activity", "type",
        ])
        if arrow.num_rows == 0:
            return []
        cols = arrow.column_names
        result = []
        for i in range(arrow.num_rows):
            row = {c: arrow.column(c)[i].as_py() for c in cols}
            result.append(row)
        return result

    def update_energy(self, memory_id: str, energy: float) -> None:
        """Update a memory's energy value."""
        self._table.update(
            where=f"id = '{memory_id}'",
            values={"energy": energy},
        )

    def count(self) -> int:
        return self._table.count_rows()

    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 ---

    @staticmethod
    def _where_to_sql(where: dict) -> Optional[str]:
        """Convert ChromaDB-style where dict to SQL filter string."""
        if "$and" in where:
            parts = [MemoryStore._where_to_sql(c) for c in where["$and"]]
            parts = [p for p in parts if p]
            return " AND ".join(parts) if parts else None

        for key, val in where.items():
            if key.startswith("$"):
                continue
            if isinstance(val, dict):
                for op, v in val.items():
                    if op == "$eq":
                        return f"{key} = '{v}'" if isinstance(v, str) else f"{key} = {v}"
                    elif op == "$gte":
                        return f"{key} >= {v}"
                    elif op == "$lt":
                        return f"{key} < {v}"
            else:
                return f"{key} = '{val}'" if isinstance(val, str) else f"{key} = {val}"
        return None

    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

        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: list[dict]) -> list[tuple[MemoryRecord, float]]:
        """Parse LanceDB search results into (MemoryRecord, similarity) pairs."""
        pairs = []
        for row in results:
            similarity = 1.0 - row.get("_distance", 0)
            record = self._meta_to_record(row)
            pairs.append((record, similarity))
        return pairs
