"""
ChunkStore: Raw Text Chunk Storage with Vector Search

Issue #138: Hybrid Retrieval — Graph + Raw Text Chunks

MCAL's UnifiedGraph captures structured knowledge (goals, decisions, entities,
facts, relationships) but discards the original conversational text. This means
any information not explicitly extracted by the prompt is permanently lost.

ChunkStore preserves raw conversational text as overlapping chunks with
embeddings, enabling vector search over the original content. Combined with
graph search, this "hybrid retrieval" fills factual gaps that the extraction
prompt misses.

Architecture:
    ┌─────────────────────────────────────────────┐
    │  ChunkStore (per-user)                       │
    │  ┌──────────────────────────────────────┐   │
    │  │  TextChunk[]  (sliding window, FIFO) │   │
    │  │  - text, turn_ref, metadata          │   │
    │  │  - embedding (Float16 binary)        │   │
    │  └──────────────────────────────────────┘   │
    │  ┌──────────────────────────────────────┐   │
    │  │  FAISS VectorIndex (cosine sim)      │   │
    │  └──────────────────────────────────────┘   │
    └─────────────────────────────────────────────┘

Performance:
    - Token-based chunking via tiktoken
    - Float16 embeddings (8x compression)
    - FAISS IndexFlatIP for exact search
    - FIFO eviction keeps memory bounded
"""

from __future__ import annotations

import logging
import uuid
from dataclasses import dataclass, field
from typing import Optional

import numpy as np

from .embeddings import (
    EmbeddingService,
    embedding_to_base64,
    base64_to_embedding,
    embedding_to_bytes,
)
from .vector_index import VectorIndex

logger = logging.getLogger(__name__)


# =========================================================================
# Data Model
# =========================================================================

@dataclass
class TextChunk:
    """A chunk of raw conversational text with metadata."""
    id: str
    text: str
    turn_ref: str  # e.g., "turn_5" or "turn_5-turn_8"
    metadata: dict = field(default_factory=dict)
    embedding: Optional[bytes] = None  # Float16 binary

    def to_dict(self) -> dict:
        """Serialize for JSON persistence."""
        d = {
            "id": self.id,
            "text": self.text,
            "turn_ref": self.turn_ref,
            "metadata": self.metadata,
        }
        if self.embedding is not None:
            d["embedding"] = embedding_to_base64(
                np.frombuffer(self.embedding, dtype=np.float16).astype(np.float32)
            )
        return d

    @classmethod
    def from_dict(cls, data: dict) -> "TextChunk":
        """Deserialize from JSON."""
        emb = None
        if "embedding" in data and data["embedding"]:
            emb = embedding_to_bytes(base64_to_embedding(data["embedding"]))
        return cls(
            id=data["id"],
            text=data["text"],
            turn_ref=data.get("turn_ref", ""),
            metadata=data.get("metadata", {}),
            embedding=emb,
        )


# =========================================================================
# ChunkStore
# =========================================================================

class ChunkStore:
    """
    Raw text chunk storage with vector search (Issue #138).

    Stores conversational text as overlapping token-based chunks with
    embeddings for semantic retrieval. Complements UnifiedGraph by
    preserving information that the extraction prompt doesn't capture.

    Features:
        - Token-based chunking (not character-based) for accurate budgets
        - Overlapping windows prevent information loss at boundaries
        - FIFO eviction keeps memory bounded  
        - Float16 embeddings for 8x compression
        - FAISS-based cosine similarity search

    Usage:
        store = ChunkStore(chunk_size=512, chunk_overlap=64, max_chunks=200)
        store.add("User said they love hiking and their sister is a nurse.",
                   turn_ref="turn_5")
        results = store.search("What does the user's sister do?", k=3)
        for chunk, score in results:
            print(f"{score:.3f}: {chunk.text[:80]}...")
    """

    def __init__(
        self,
        chunk_size: int = 512,
        chunk_overlap: int = 64,
        max_chunks: int = 200,
        tokenizer_model: str = "gpt-4o",
    ):
        """
        Initialize ChunkStore.

        Args:
            chunk_size: Maximum tokens per chunk (default: 512)
            chunk_overlap: Overlapping tokens between consecutive chunks (default: 64)
            max_chunks: Maximum chunks to retain (FIFO eviction, default: 200)
            tokenizer_model: tiktoken model for tokenization (default: gpt-4o)
        """
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.max_chunks = max_chunks
        self._tokenizer_model = tokenizer_model

        self.chunks: list[TextChunk] = []
        self._index: Optional[VectorIndex] = None
        self._index_dirty: bool = True
        self._embedding_service = EmbeddingService()

    # =========================================================================
    # Add
    # =========================================================================

    def add(
        self,
        text: str,
        turn_ref: str = "",
        metadata: Optional[dict] = None,
    ) -> int:
        """
        Split text into overlapping chunks, embed, and store.

        Args:
            text: Raw conversational text to store
            turn_ref: Turn reference (e.g., "turn_5" or "turn_5-turn_8")
            metadata: Optional metadata dict (session_id, etc.)

        Returns:
            Number of chunks created
        """
        if not text or not text.strip():
            return 0

        metadata = metadata or {}
        token_chunks = self._split_into_chunks(text)

        if not token_chunks:
            return 0

        # Create TextChunk objects
        new_chunks = []
        for chunk_text in token_chunks:
            chunk = TextChunk(
                id=f"chunk_{uuid.uuid4().hex[:12]}",
                text=chunk_text,
                turn_ref=turn_ref,
                metadata=metadata.copy(),
            )
            new_chunks.append(chunk)

        # Batch embed
        texts = [c.text for c in new_chunks]
        embeddings = self._embedding_service.embed_texts(texts)
        for chunk, emb in zip(new_chunks, embeddings):
            chunk.embedding = emb

        # Append and enforce FIFO limit
        self.chunks.extend(new_chunks)
        evicted = 0
        if len(self.chunks) > self.max_chunks:
            evicted = len(self.chunks) - self.max_chunks
            self.chunks = self.chunks[evicted:]

        self._index_dirty = True

        logger.debug(
            f"ChunkStore: added {len(new_chunks)} chunks "
            f"(total={len(self.chunks)}, evicted={evicted})"
        )
        return len(new_chunks)

    # =========================================================================
    # Search
    # =========================================================================

    def search(self, query: str, k: int = 5) -> list[tuple[TextChunk, float]]:
        """
        Vector search over raw chunks.

        Args:
            query: Search query text
            k: Number of results to return

        Returns:
            List of (TextChunk, similarity_score) tuples, sorted by score desc
        """
        if not self.chunks:
            return []

        self._ensure_index()

        query_emb = self._embedding_service.embed_text(query)
        raw_results = self._index.search(query_emb, k=k)

        # Map chunk IDs back to TextChunk objects
        chunk_map = {c.id: c for c in self.chunks}
        results = []
        for chunk_id, score in raw_results:
            if chunk_id in chunk_map:
                results.append((chunk_map[chunk_id], score))

        return results

    # =========================================================================
    # Index Management
    # =========================================================================

    def _ensure_index(self) -> None:
        """Rebuild FAISS index if dirty."""
        if not self._index_dirty and self._index is not None:
            return

        self._index = VectorIndex()
        nodes_with_embeddings = [
            (chunk.id, chunk.embedding)
            for chunk in self.chunks
            if chunk.embedding is not None
        ]
        if nodes_with_embeddings:
            self._index.add_batch(nodes_with_embeddings)
        self._index_dirty = False

    # =========================================================================
    # Chunking
    # =========================================================================

    def _split_into_chunks(self, text: str) -> list[str]:
        """
        Split text into overlapping token-based chunks.

        Uses tiktoken for accurate token counting. Falls back to
        word-based splitting if tiktoken is unavailable.

        Args:
            text: Raw text to split

        Returns:
            List of chunk strings
        """
        try:
            import tiktoken
            enc = tiktoken.encoding_for_model(self._tokenizer_model)
            tokens = enc.encode(text)

            if len(tokens) <= self.chunk_size:
                return [text]

            chunks = []
            start = 0
            stride = self.chunk_size - self.chunk_overlap

            while start < len(tokens):
                end = min(start + self.chunk_size, len(tokens))
                chunk_tokens = tokens[start:end]
                chunk_text = enc.decode(chunk_tokens)
                chunks.append(chunk_text)

                if end >= len(tokens):
                    break
                start += stride

            return chunks

        except ImportError:
            logger.warning("tiktoken not available, using word-based chunking")
            return self._split_by_words(text)

    def _split_by_words(self, text: str) -> list[str]:
        """Fallback: split by words (~0.75 tokens per word heuristic)."""
        words = text.split()
        approx_words_per_chunk = int(self.chunk_size * 0.75)
        overlap_words = int(self.chunk_overlap * 0.75)

        if len(words) <= approx_words_per_chunk:
            return [text]

        chunks = []
        start = 0
        stride = approx_words_per_chunk - overlap_words

        while start < len(words):
            end = min(start + approx_words_per_chunk, len(words))
            chunk_text = " ".join(words[start:end])
            chunks.append(chunk_text)

            if end >= len(words):
                break
            start += stride

        return chunks

    # =========================================================================
    # Persistence
    # =========================================================================

    def to_dict(self) -> dict:
        """Serialize for JSON persistence."""
        return {
            "version": "chunk_store_v1",
            "chunk_size": self.chunk_size,
            "chunk_overlap": self.chunk_overlap,
            "max_chunks": self.max_chunks,
            "chunks": [c.to_dict() for c in self.chunks],
        }

    @classmethod
    def from_dict(cls, data: dict) -> "ChunkStore":
        """Deserialize from JSON persistence."""
        store = cls(
            chunk_size=data.get("chunk_size", 512),
            chunk_overlap=data.get("chunk_overlap", 64),
            max_chunks=data.get("max_chunks", 200),
        )
        store.chunks = [TextChunk.from_dict(c) for c in data.get("chunks", [])]
        store._index_dirty = True
        return store

    # =========================================================================
    # Utilities
    # =========================================================================

    def clear(self) -> None:
        """Remove all chunks and reset index."""
        self.chunks = []
        self._index = None
        self._index_dirty = True

    def total_tokens(self) -> int:
        """Estimate total tokens stored across all chunks."""
        try:
            import tiktoken
            enc = tiktoken.encoding_for_model(self._tokenizer_model)
            return sum(len(enc.encode(c.text)) for c in self.chunks)
        except ImportError:
            # Rough estimate: 1.3 tokens per word
            return sum(int(len(c.text.split()) * 1.3) for c in self.chunks)

    def __len__(self) -> int:
        """Return number of chunks stored."""
        return len(self.chunks)

    def __repr__(self) -> str:
        return f"ChunkStore(chunks={len(self.chunks)}, max={self.max_chunks})"
