# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Context Engineering for Familiar

Implements Google ADK's three principles:
1. Separate storage from presentation (Session vs Working Context)
2. Explicit transformations (named, ordered processors)
3. Scope by default (minimum context per call)

This transforms raw conversation history and memory into optimized
prompts that fit within token budgets while preserving semantic value.

Usage:
    from familiar.core.context import ContextCompiler, ContextConfig

    compiler = ContextCompiler(config=ContextConfig(max_tokens=8000))
    working_context = compiler.compile(
        memory=agent.memory,
        history=conversation_history,
        tools=available_tools,
        task_hint="email triage"
    )
"""

import hashlib
import json
import logging
import threading
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from datetime import datetime, timedelta
from enum import Enum
from typing import Any, Callable, Dict, List, Optional

logger = logging.getLogger(__name__)


# ============================================================
# ACCURATE TOKEN COUNTING
# ============================================================


class TokenCounter:
    """
    Accurate token counting using tiktoken with graceful fallback.

    Thread-safe singleton pattern ensures encoding is loaded once.

    Usage:
        counter = get_token_counter()
        count = counter.count("Hello world")
        count = counter.count_messages([{"role": "user", "content": "Hi"}])
    """

    _instance = None
    _lock = threading.Lock()
    _encodings: Dict[str, Any] = {}

    def __new__(cls, encoding_name: str = "cl100k_base"):
        with cls._lock:
            if cls._instance is None:
                cls._instance = super().__new__(cls)
                cls._instance._initialized = False
        return cls._instance

    def __init__(self, encoding_name: str = "cl100k_base"):
        if self._initialized:
            return

        self.encoding_name = encoding_name
        self._encoding = self._load_encoding(encoding_name)
        self._initialized = True

    def _load_encoding(self, name: str):
        """Load tiktoken encoding with fallback."""
        if name in self._encodings:
            return self._encodings[name]

        try:
            import tiktoken

            encoding = tiktoken.get_encoding(name)
            self._encodings[name] = encoding
            logger.info(f"TokenCounter: Using tiktoken '{name}' encoding")
            return encoding
        except ImportError:
            logger.warning("TokenCounter: tiktoken not installed. Using heuristic fallback.")
            self._encodings[name] = None
            return None
        except Exception as e:
            logger.warning(f"TokenCounter: Failed to load tiktoken: {e}. Using fallback.")
            self._encodings[name] = None
            return None

    def count(self, text: str) -> int:
        """
        Count tokens in text.

        Uses tiktoken if available, otherwise improved heuristic.
        """
        if not text:
            return 0

        if self._encoding:
            try:
                return len(self._encoding.encode(text))
            except Exception:
                pass  # Fall through to heuristic

        # Improved heuristic fallback
        # JSON/code averages ~3.5 chars/token, prose ~4 chars/token
        # Whitespace-heavy content can be lower
        return self._heuristic_count(text)

    def _heuristic_count(self, text: str) -> int:
        """
        Improved heuristic token estimation.

        Accounts for different content types:
        - JSON/code: ~3.5 chars/token (more punctuation)
        - Prose: ~4 chars/token
        - Whitespace-heavy: ~4.5 chars/token
        """
        if not text:
            return 0

        # Detect content type
        stripped = text.strip()

        # JSON or code detection
        if stripped.startswith(("{", "[", "def ", "class ", "function ", "const ", "let ", "var ")):
            ratio = 3.5
        # XML/HTML detection
        elif stripped.startswith("<") and ">" in stripped[:100]:
            ratio = 3.2
        # Markdown with lots of formatting
        elif text.count("#") > 5 or text.count("*") > 10:
            ratio = 3.8
        else:
            # Standard prose
            ratio = 4.0

        # Adjust for whitespace density
        whitespace_ratio = sum(1 for c in text if c.isspace()) / max(len(text), 1)
        if whitespace_ratio > 0.2:
            ratio += 0.3

        return max(1, int(len(text) / ratio))

    def count_messages(self, messages: List[Dict]) -> int:
        """
        Count tokens in a message list including overhead.

        Accounts for role tokens, separators, and structure.
        """
        if not messages:
            return 0

        total = 0

        for msg in messages:
            # Message structure overhead (~4 tokens per message)
            total += 4

            # Role token
            total += 1

            # Content
            content = msg.get("content", "")
            if isinstance(content, str):
                total += self.count(content)
            elif isinstance(content, list):
                # Multi-part content (text + images)
                for part in content:
                    if isinstance(part, dict):
                        if part.get("type") == "text":
                            total += self.count(part.get("text", ""))
                        elif part.get("type") == "image":
                            # Image tokens vary by size, estimate ~85 for small
                            total += 85

            # Tool calls
            if msg.get("tool_calls"):
                total += self.count(json.dumps(msg["tool_calls"]))

            # Tool results
            if msg.get("tool_call_id"):
                total += 3  # Tool result overhead

        # Conversation overhead
        total += 3

        return total

    def count_tools(self, tools: List[Dict]) -> int:
        """Count tokens for tool definitions."""
        if not tools:
            return 0

        # Tool definitions are JSON, use JSON-aware counting
        return self.count(json.dumps(tools))

    def truncate_to_limit(self, text: str, max_tokens: int) -> str:
        """
        Truncate text to fit within token limit.

        Returns the longest prefix that fits within max_tokens.
        """
        if not text or max_tokens <= 0:
            return ""

        current_count = self.count(text)
        if current_count <= max_tokens:
            return text

        if self._encoding:
            try:
                tokens = self._encoding.encode(text)
                if len(tokens) <= max_tokens:
                    return text
                truncated_tokens = tokens[:max_tokens]
                return self._encoding.decode(truncated_tokens)
            except Exception:
                pass

        # Fallback: binary search for right length
        ratio = max_tokens / max(current_count, 1)
        estimated_chars = int(len(text) * ratio * 0.95)  # Conservative

        # Refine with binary search
        low, high = 0, len(text)
        result = text[:estimated_chars]

        for _ in range(5):  # Max 5 iterations
            mid = (low + high) // 2
            if self.count(text[:mid]) <= max_tokens:
                result = text[:mid]
                low = mid + 1
            else:
                high = mid - 1

        return result

    def fits_in_budget(self, text: str, budget: int) -> bool:
        """Check if text fits within token budget."""
        return self.count(text) <= budget

    def remaining_budget(self, used: int, total: int) -> int:
        """Calculate remaining token budget."""
        return max(0, total - used)


# Global token counter instance
_token_counter: Optional[TokenCounter] = None


def get_token_counter() -> TokenCounter:
    """Get the global token counter instance."""
    global _token_counter
    if _token_counter is None:
        _token_counter = TokenCounter()
    return _token_counter


# ============================================================
# SEMANTIC RETRIEVAL (v1.4.0)
# ============================================================


@dataclass
class Document:
    """
    A document or chunk for semantic retrieval.

    Represents a piece of content that can be embedded and retrieved.
    """

    id: str
    content: str
    metadata: Dict[str, Any] = field(default_factory=dict)
    embedding: Optional[List[float]] = None

    # Retrieval metadata
    score: float = 0.0  # Similarity score from retrieval
    source: str = ""  # Source identifier (file, url, etc.)
    created_at: Optional[str] = None

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "content": self.content,
            "metadata": self.metadata,
            "score": self.score,
            "source": self.source,
        }


@dataclass
class RetrievalConfig:
    """Configuration for semantic retrieval."""

    # Embedding settings
    embedding_model: str = "text-embedding-3-small"  # OpenAI default
    embedding_dimension: int = 1536

    # Chunking settings
    chunk_size: int = 512  # Tokens per chunk
    chunk_overlap: int = 50  # Overlap between chunks

    # Retrieval settings
    top_k: int = 5  # Number of results to return
    min_score: float = 0.0  # Minimum similarity score
    max_tokens: int = 2000  # Max tokens in retrieved context

    # Reranking
    enable_reranking: bool = False
    reranking_model: str = ""

    # Hybrid search
    enable_hybrid: bool = False  # Combine semantic + keyword search
    keyword_weight: float = 0.3  # Weight for keyword results


class EmbeddingProvider(ABC):
    """
    Abstract base class for embedding providers.

    Implementations can use OpenAI, Cohere, local models, etc.
    """

    @abstractmethod
    def embed(self, text: str) -> List[float]:
        """Embed a single text."""
        pass

    @abstractmethod
    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """Embed multiple texts efficiently."""
        pass

    @property
    @abstractmethod
    def dimension(self) -> int:
        """Return the embedding dimension."""
        pass


class OpenAIEmbedding(EmbeddingProvider):
    """
    OpenAI embedding provider.

    Requires: pip install openai

    Usage:
        provider = OpenAIEmbedding(api_key="sk-...")
        embedding = provider.embed("Hello world")
    """

    DIMENSIONS = {
        "text-embedding-3-small": 1536,
        "text-embedding-3-large": 3072,
        "text-embedding-ada-002": 1536,
    }

    def __init__(
        self,
        model: str = "text-embedding-3-small",
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
    ):
        self.model = model
        self._api_key = api_key
        self._base_url = base_url
        self._client = None
        self._dimension = self.DIMENSIONS.get(model, 1536)

    def _get_client(self):
        """Lazy initialization of OpenAI client."""
        if self._client is None:
            try:
                from openai import OpenAI

                kwargs = {}
                if self._api_key:
                    kwargs["api_key"] = self._api_key
                if self._base_url:
                    kwargs["base_url"] = self._base_url
                self._client = OpenAI(**kwargs)
            except ImportError:
                raise ImportError("OpenAI package not available.")
        return self._client

    def embed(self, text: str) -> List[float]:
        """Embed a single text."""
        result = self.embed_batch([text])
        return result[0]

    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """Embed multiple texts."""
        if not texts:
            return []

        client = self._get_client()
        response = client.embeddings.create(model=self.model, input=texts)

        # Sort by index to maintain order
        embeddings = sorted(response.data, key=lambda x: x.index)
        return [e.embedding for e in embeddings]

    @property
    def dimension(self) -> int:
        return self._dimension


class LocalEmbedding(EmbeddingProvider):
    """
    Local embedding provider using sentence-transformers.

    Requires: pip install sentence-transformers

    Usage:
        provider = LocalEmbedding(model_name="all-MiniLM-L6-v2")
        embedding = provider.embed("Hello world")
    """

    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model_name = model_name
        self._model = None
        self._dimension = None

    def _get_model(self):
        """Lazy initialization of sentence-transformers model."""
        if self._model is None:
            try:
                from sentence_transformers import SentenceTransformer

                self._model = SentenceTransformer(self.model_name)
                self._dimension = self._model.get_sentence_embedding_dimension()
            except ImportError:
                raise ImportError("sentence-transformers not available. Check server logs.")
        return self._model

    def embed(self, text: str) -> List[float]:
        """Embed a single text."""
        model = self._get_model()
        embedding = model.encode(text, convert_to_numpy=True)
        return embedding.tolist()

    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        """Embed multiple texts."""
        if not texts:
            return []
        model = self._get_model()
        embeddings = model.encode(texts, convert_to_numpy=True)
        return embeddings.tolist()

    @property
    def dimension(self) -> int:
        if self._dimension is None:
            self._get_model()  # Initialize to get dimension
        return self._dimension


class VectorStore(ABC):
    """
    Abstract base class for vector stores.

    Implementations can use in-memory, ChromaDB, Pinecone, etc.
    """

    @abstractmethod
    def add(self, documents: List[Document]) -> None:
        """Add documents to the store."""
        pass

    @abstractmethod
    def search(
        self, query_embedding: List[float], top_k: int = 5, filter: Optional[Dict] = None
    ) -> List[Document]:
        """Search for similar documents."""
        pass

    @abstractmethod
    def delete(self, ids: List[str]) -> None:
        """Delete documents by ID."""
        pass

    @abstractmethod
    def clear(self) -> None:
        """Clear all documents."""
        pass

    @property
    @abstractmethod
    def count(self) -> int:
        """Return the number of documents."""
        pass


class InMemoryVectorStore(VectorStore):
    """
    Simple in-memory vector store using cosine similarity.

    Good for development and small datasets. For production,
    use ChromaVectorStore or a dedicated vector database.

    Usage:
        store = InMemoryVectorStore()
        store.add([Document(id="1", content="Hello", embedding=[0.1, 0.2, ...])])
        results = store.search(query_embedding, top_k=5)
    """

    def __init__(self):
        self._documents: Dict[str, Document] = {}
        self._lock = threading.RLock()

    def add(self, documents: List[Document]) -> None:
        """Add documents to the store."""
        with self._lock:
            for doc in documents:
                if doc.embedding is None:
                    raise ValueError(f"Document {doc.id} has no embedding")
                self._documents[doc.id] = doc

    def search(
        self, query_embedding: List[float], top_k: int = 5, filter: Optional[Dict] = None
    ) -> List[Document]:
        """Search for similar documents using cosine similarity."""

        with self._lock:
            if not self._documents:
                return []

            # Calculate cosine similarity for all documents
            scores = []
            for doc in self._documents.values():
                # Apply filter if provided
                if filter:
                    match = all(doc.metadata.get(k) == v for k, v in filter.items())
                    if not match:
                        continue

                score = self._cosine_similarity(query_embedding, doc.embedding)
                scores.append((score, doc))

            # Sort by score descending
            scores.sort(key=lambda x: x[0], reverse=True)

            # Return top_k with scores
            results = []
            for score, doc in scores[:top_k]:
                result = Document(
                    id=doc.id,
                    content=doc.content,
                    metadata=doc.metadata,
                    embedding=doc.embedding,
                    score=score,
                    source=doc.source,
                    created_at=doc.created_at,
                )
                results.append(result)

            return results

    def _cosine_similarity(self, a: List[float], b: List[float]) -> float:
        """Calculate cosine similarity between two vectors."""
        import math

        dot_product = sum(x * y for x, y in zip(a, b))
        norm_a = math.sqrt(sum(x * x for x in a))
        norm_b = math.sqrt(sum(x * x for x in b))

        if norm_a == 0 or norm_b == 0:
            return 0.0

        return dot_product / (norm_a * norm_b)

    def delete(self, ids: List[str]) -> None:
        """Delete documents by ID."""
        with self._lock:
            for id in ids:
                self._documents.pop(id, None)

    def clear(self) -> None:
        """Clear all documents."""
        with self._lock:
            self._documents.clear()

    @property
    def count(self) -> int:
        """Return the number of documents."""
        return len(self._documents)

    def get(self, id: str) -> Optional[Document]:
        """Get a document by ID."""
        return self._documents.get(id)

    def get_all(self) -> List[Document]:
        """Get all documents."""
        return list(self._documents.values())


class ChromaVectorStore(VectorStore):
    """
    ChromaDB vector store for persistent storage.

    Requires: pip install chromadb

    Usage:
        store = ChromaVectorStore(collection_name="my_docs", persist_dir="./chroma")
        store.add([Document(id="1", content="Hello", embedding=[...])])
    """

    def __init__(self, collection_name: str = "familiar", persist_dir: Optional[str] = None):
        self.collection_name = collection_name
        self.persist_dir = persist_dir
        self._client = None
        self._collection = None

    def _get_collection(self):
        """Lazy initialization of ChromaDB."""
        if self._collection is None:
            try:
                import chromadb

                if self.persist_dir:
                    self._client = chromadb.PersistentClient(path=self.persist_dir)
                else:
                    self._client = chromadb.Client()

                self._collection = self._client.get_or_create_collection(
                    name=self.collection_name, metadata={"hnsw:space": "cosine"}
                )
            except ImportError:
                raise ImportError("ChromaDB not available.")
        return self._collection

    def add(self, documents: List[Document]) -> None:
        """Add documents to ChromaDB."""
        if not documents:
            return

        collection = self._get_collection()

        ids = [doc.id for doc in documents]
        embeddings = [doc.embedding for doc in documents]
        contents = [doc.content for doc in documents]
        metadatas = [doc.metadata for doc in documents]

        collection.add(ids=ids, embeddings=embeddings, documents=contents, metadatas=metadatas)

    def search(
        self, query_embedding: List[float], top_k: int = 5, filter: Optional[Dict] = None
    ) -> List[Document]:
        """Search ChromaDB."""
        collection = self._get_collection()

        kwargs = {
            "query_embeddings": [query_embedding],
            "n_results": top_k,
            "include": ["documents", "metadatas", "distances", "embeddings"],
        }

        if filter:
            kwargs["where"] = filter

        results = collection.query(**kwargs)

        documents = []
        if results["ids"] and results["ids"][0]:
            for i, id in enumerate(results["ids"][0]):
                # ChromaDB returns distances, convert to similarity
                distance = results["distances"][0][i] if results["distances"] else 0
                score = 1 - distance  # Cosine distance to similarity

                doc = Document(
                    id=id,
                    content=results["documents"][0][i] if results["documents"] else "",
                    metadata=results["metadatas"][0][i] if results["metadatas"] else {},
                    embedding=results["embeddings"][0][i] if results.get("embeddings") else None,
                    score=score,
                )
                documents.append(doc)

        return documents

    def delete(self, ids: List[str]) -> None:
        """Delete documents from ChromaDB."""
        if not ids:
            return
        collection = self._get_collection()
        collection.delete(ids=ids)

    def clear(self) -> None:
        """Clear all documents."""
        if self._client and self._collection:
            self._client.delete_collection(self.collection_name)
            self._collection = None

    @property
    def count(self) -> int:
        """Return document count."""
        collection = self._get_collection()
        return collection.count()


class SemanticRetriever:
    """
    High-level semantic retrieval for RAG applications.

    Combines embedding, chunking, storage, and retrieval into
    a simple interface for adding knowledge and querying it.

    Usage:
        # Initialize
        retriever = SemanticRetriever(
            embedding_provider=OpenAIEmbedding(),
            vector_store=InMemoryVectorStore()
        )

        # Add documents
        retriever.add_text("Document content...", metadata={"source": "manual"})
        retriever.add_texts(["Doc 1", "Doc 2", "Doc 3"])

        # Retrieve relevant context
        results = retriever.retrieve("What is the policy on X?", top_k=5)

        # Get as formatted context string
        context = retriever.get_context("What is X?", max_tokens=1000)
    """

    def __init__(
        self,
        embedding_provider: Optional[EmbeddingProvider] = None,
        vector_store: Optional[VectorStore] = None,
        config: Optional[RetrievalConfig] = None,
    ):
        self.config = config or RetrievalConfig()
        self._embedding = embedding_provider
        self._store = vector_store or InMemoryVectorStore()
        self._token_counter = get_token_counter()
        self._doc_counter = 0
        self._lock = threading.RLock()

    def _get_embedding_provider(self) -> EmbeddingProvider:
        """Get or create embedding provider."""
        if self._embedding is None:
            # Try to create default provider
            try:
                self._embedding = OpenAIEmbedding(model=self.config.embedding_model)
            except ImportError:
                try:
                    self._embedding = LocalEmbedding()
                except ImportError:
                    raise ImportError(
                        "No embedding provider available. Install openai or sentence-transformers."
                    )
        return self._embedding

    def _generate_id(self) -> str:
        """Generate a unique document ID."""
        with self._lock:
            self._doc_counter += 1
            return f"doc_{self._doc_counter}_{datetime.now().strftime('%Y%m%d%H%M%S')}"

    def _chunk_text(self, text: str) -> List[str]:
        """Split text into chunks based on token count."""
        if not text:
            return []

        # Simple sentence-based chunking
        sentences = text.replace("\n", " ").split(". ")

        chunks = []
        current_chunk = []
        current_tokens = 0

        for sentence in sentences:
            sentence = sentence.strip()
            if not sentence:
                continue

            sentence_tokens = self._token_counter.count(sentence)

            if current_tokens + sentence_tokens > self.config.chunk_size:
                # Save current chunk
                if current_chunk:
                    chunks.append(". ".join(current_chunk) + ".")

                # Start new chunk with overlap
                if self.config.chunk_overlap > 0 and len(current_chunk) > 1:
                    # Keep some sentences for overlap
                    overlap_sentences = []
                    overlap_tokens = 0
                    for s in reversed(current_chunk):
                        s_tokens = self._token_counter.count(s)
                        if overlap_tokens + s_tokens > self.config.chunk_overlap:
                            break
                        overlap_sentences.insert(0, s)
                        overlap_tokens += s_tokens
                    current_chunk = overlap_sentences + [sentence]
                    current_tokens = overlap_tokens + sentence_tokens
                else:
                    current_chunk = [sentence]
                    current_tokens = sentence_tokens
            else:
                current_chunk.append(sentence)
                current_tokens += sentence_tokens

        # Don't forget the last chunk
        if current_chunk:
            chunks.append(". ".join(current_chunk) + ".")

        return chunks

    def add_text(
        self, text: str, metadata: Optional[Dict] = None, source: str = "", chunk: bool = True
    ) -> List[str]:
        """
        Add a text document, optionally chunking it.

        Returns list of document IDs created.
        """
        if not text:
            return []

        provider = self._get_embedding_provider()
        metadata = metadata or {}

        if chunk:
            chunks = self._chunk_text(text)
        else:
            chunks = [text]

        if not chunks:
            return []

        # Embed all chunks
        embeddings = provider.embed_batch(chunks)

        # Create documents
        documents = []
        ids = []
        for i, (chunk_text, embedding) in enumerate(zip(chunks, embeddings)):
            doc_id = self._generate_id()
            doc = Document(
                id=doc_id,
                content=chunk_text,
                metadata={**metadata, "chunk_index": i, "total_chunks": len(chunks)},
                embedding=embedding,
                source=source,
                created_at=datetime.now().isoformat(),
            )
            documents.append(doc)
            ids.append(doc_id)

        # Store
        self._store.add(documents)

        logger.debug(f"Added {len(documents)} chunks from text ({len(text)} chars)")
        return ids

    def add_texts(
        self,
        texts: List[str],
        metadatas: Optional[List[Dict]] = None,
        source: str = "",
        chunk: bool = True,
    ) -> List[str]:
        """Add multiple texts."""
        all_ids = []
        metadatas = metadatas or [{}] * len(texts)

        for text, metadata in zip(texts, metadatas):
            ids = self.add_text(text, metadata=metadata, source=source, chunk=chunk)
            all_ids.extend(ids)

        return all_ids

    def add_document(self, document: Document) -> str:
        """Add a pre-built document."""
        if document.embedding is None:
            provider = self._get_embedding_provider()
            document.embedding = provider.embed(document.content)

        self._store.add([document])
        return document.id

    def retrieve(
        self,
        query: str,
        top_k: Optional[int] = None,
        filter: Optional[Dict] = None,
        min_score: Optional[float] = None,
    ) -> List[Document]:
        """
        Retrieve relevant documents for a query.

        Args:
            query: The search query
            top_k: Number of results (default from config)
            filter: Metadata filter
            min_score: Minimum similarity score

        Returns:
            List of Document objects with scores
        """
        if not query:
            return []

        top_k = top_k or self.config.top_k
        min_score = min_score if min_score is not None else self.config.min_score

        # Embed query
        provider = self._get_embedding_provider()
        query_embedding = provider.embed(query)

        # Search
        results = self._store.search(query_embedding, top_k=top_k, filter=filter)

        # Filter by min_score
        if min_score > 0:
            results = [r for r in results if r.score >= min_score]

        return results

    def get_context(
        self,
        query: str,
        max_tokens: Optional[int] = None,
        top_k: Optional[int] = None,
        filter: Optional[Dict] = None,
        format: str = "numbered",
    ) -> str:
        """
        Get retrieved documents as a formatted context string.

        Args:
            query: The search query
            max_tokens: Maximum tokens in result
            top_k: Number of results to consider
            filter: Metadata filter
            format: "numbered", "bullets", or "plain"

        Returns:
            Formatted string of relevant context
        """
        max_tokens = max_tokens or self.config.max_tokens
        top_k = top_k or self.config.top_k

        results = self.retrieve(query, top_k=top_k * 2, filter=filter)  # Get more, then trim

        if not results:
            return ""

        # Build context within token budget
        context_parts = []
        total_tokens = 0

        for i, doc in enumerate(results):
            doc_tokens = self._token_counter.count(doc.content)

            if total_tokens + doc_tokens > max_tokens:
                # Try to fit a truncated version
                remaining = max_tokens - total_tokens
                if remaining > 50:  # Only include if we can fit something meaningful
                    truncated = self._token_counter.truncate_to_limit(doc.content, remaining - 10)
                    context_parts.append(self._format_result(truncated, i + 1, format, doc.score))
                break

            context_parts.append(self._format_result(doc.content, i + 1, format, doc.score))
            total_tokens += doc_tokens

        return "\n\n".join(context_parts)

    def _format_result(self, content: str, index: int, format: str, score: float) -> str:
        """Format a single result."""
        if format == "numbered":
            return f"[{index}] (score: {score:.2f})\n{content}"
        elif format == "bullets":
            return f"• {content}"
        else:
            return content

    def delete(self, ids: List[str]) -> None:
        """Delete documents by ID."""
        self._store.delete(ids)

    def clear(self) -> None:
        """Clear all documents."""
        self._store.clear()

    @property
    def count(self) -> int:
        """Return document count."""
        return self._store.count


# Global retriever instance
_semantic_retriever: Optional[SemanticRetriever] = None


def get_semantic_retriever() -> SemanticRetriever:
    """Get the global semantic retriever instance."""
    global _semantic_retriever
    if _semantic_retriever is None:
        _semantic_retriever = SemanticRetriever()
    return _semantic_retriever


def set_semantic_retriever(retriever: SemanticRetriever):
    """Set the global semantic retriever instance."""
    global _semantic_retriever
    _semantic_retriever = retriever


# ============================================================
# CONFIGURATION
# ============================================================


class CompactionStrategy(str, Enum):
    """How to handle context that exceeds limits."""

    TRUNCATE_OLD = "truncate_old"  # Remove oldest messages
    SUMMARIZE = "summarize"  # LLM summarization of old context
    SLIDING_WINDOW = "sliding_window"  # Keep N recent + summary of older
    IMPORTANCE = "importance"  # Keep high-importance items


@dataclass
class ContextConfig:
    """Configuration for context compilation."""

    # Token budgets (uses tiktoken for accurate counting, falls back to heuristic)
    max_tokens: int = 8000
    max_history_tokens: int = 4000
    max_memory_tokens: int = 1000
    max_tools_tokens: int = 2000

    # History settings
    max_messages: int = 20
    summarize_after: int = 10  # Summarize messages older than this

    # Compaction
    compaction_strategy: CompactionStrategy = CompactionStrategy.SLIDING_WINDOW
    compaction_interval: int = 5  # Summarize every N messages

    # Memory filtering
    memory_recency_days: int = 30  # Prioritize memories from last N days
    memory_relevance_threshold: float = 0.5  # Min relevance score to include

    # Tool filtering
    max_tools: int = 15  # Max tools to include in context
    tool_selection: str = "relevance"  # "all", "relevance", "recent_use"


@dataclass
class WorkingContext:
    """
    The compiled context ready for LLM consumption.

    This is the "view" that gets sent to the model - a projection
    of the underlying storage optimized for the current task.
    """

    system_prompt: str
    messages: List[Dict[str, Any]]
    tools: List[Dict[str, Any]]

    # Metadata for debugging
    compilation_time_ms: float = 0.0
    original_message_count: int = 0
    included_message_count: int = 0
    summarized_message_count: int = 0
    memory_entries_included: int = 0
    tools_included: int = 0
    estimated_tokens: int = 0

    # Cache key for detecting unchanged context
    content_hash: str = ""

    def to_dict(self) -> dict:
        return {
            "system_prompt": self.system_prompt,
            "messages": self.messages,
            "tools": self.tools,
            "metadata": {
                "compilation_time_ms": self.compilation_time_ms,
                "original_messages": self.original_message_count,
                "included_messages": self.included_message_count,
                "summarized_messages": self.summarized_message_count,
                "memory_entries": self.memory_entries_included,
                "tools": self.tools_included,
                "estimated_tokens": self.estimated_tokens,
            },
        }


# ============================================================
# CONTEXT PROCESSORS (Transformations)
# ============================================================


class ContextProcessor(ABC):
    """
    Base class for context transformation processors.

    Processors are applied in order to build the working context.
    Each processor can read from storage and write to the context builder.
    """

    @property
    @abstractmethod
    def name(self) -> str:
        """Processor name for logging."""
        pass

    @abstractmethod
    def process(self, builder: "ContextBuilder", **kwargs) -> None:
        """Transform context. Modifies builder in place."""
        pass


class IdentityProcessor(ContextProcessor):
    """
    Builds the base identity/persona section.

    Phase 1: Uses constitutional personality layer for
    relationship-stage-aware, channel-appropriate identity.
    """

    @property
    def name(self) -> str:
        return "identity"

    def process(self, builder: "ContextBuilder", **kwargs) -> None:
        agent_name = kwargs.get("agent_name", "Familiar")
        agent_persona = kwargs.get("agent_persona", "a helpful AI assistant")
        provider_name = kwargs.get("provider_name", "unknown")
        channel = kwargs.get("channel", "cli")
        user_id = kwargs.get("user_id", "default")
        is_proactive = kwargs.get("is_proactive", False)
        persona_overlay = kwargs.get("persona_overlay", "")

        try:
            from .constitution import get_constitutional_prompt

            identity = get_constitutional_prompt(
                agent_name=agent_name,
                user_id=user_id,
                channel=channel,
                provider_name=provider_name,
                is_proactive=is_proactive,
                persona_overlay=persona_overlay,
            )
        except ImportError:
            # Legacy fallback
            import os

            identity = f"""You are {agent_name}, {agent_persona}.

Current time: {datetime.now().strftime("%Y-%m-%d %H:%M %Z")}
System: {os.uname().nodename} ({os.uname().sysname})
Model: {provider_name}

You have access to tools that let you interact with the local system, manage files,
run commands, and more. Use them when needed to actually accomplish tasks.

Be concise and helpful. When asked to do something, do it - don't just describe how."""

        builder.add_system_section("identity", identity, priority=100)


class MemoryProcessor(ContextProcessor):
    """Injects relevant memories into context."""

    @property
    def name(self) -> str:
        return "memory"

    def process(self, builder: "ContextBuilder", **kwargs) -> None:
        memory = kwargs.get("memory")
        task_hint = kwargs.get("task_hint", "")
        config = kwargs.get("config", ContextConfig())

        if not memory:
            return

        # Get all memories
        all_memories = list(memory.memories.values())

        if not all_memories:
            return

        # Filter by recency
        cutoff = datetime.now() - timedelta(days=config.memory_recency_days)
        recent_memories = []
        older_memories = []

        for mem in all_memories:
            try:
                mem_date = datetime.fromisoformat(mem.updated_at)
                if mem_date > cutoff:
                    recent_memories.append(mem)
                else:
                    older_memories.append(mem)
            except (ValueError, AttributeError):
                older_memories.append(mem)

        # Sort by importance
        recent_memories.sort(key=lambda m: m.importance, reverse=True)
        older_memories.sort(key=lambda m: m.importance, reverse=True)

        # Re-sort by task hint relevance when hint is provided
        if task_hint:
            hint_words = set(task_hint.lower().split())

            def _relevance(mem):
                text = f"{mem.key} {mem.value}".lower()
                return sum(1 for w in hint_words if w in text)

            recent_memories.sort(key=lambda m: (_relevance(m), m.importance), reverse=True)
            older_memories.sort(key=lambda m: (_relevance(m), m.importance), reverse=True)

        # Build memory context with budget
        memory_lines = []
        token_budget = config.max_memory_tokens
        current_tokens = 0
        counter = get_token_counter()

        # Prioritize recent high-importance memories
        for mem in recent_memories + older_memories:
            line = f"- {mem.key}: {mem.value}"
            line_tokens = counter.count(line)

            if current_tokens + line_tokens > token_budget:
                break

            memory_lines.append(line)
            current_tokens += line_tokens

        if memory_lines:
            memory_context = "## What I Know About You\n\n" + "\n".join(memory_lines)
            builder.add_system_section("memory", memory_context, priority=80)
            builder.metadata["memory_entries_included"] = len(memory_lines)


class HistoryProcessor(ContextProcessor):
    """Processes conversation history with optional summarization."""

    @property
    def name(self) -> str:
        return "history"

    def process(self, builder: "ContextBuilder", **kwargs) -> None:
        history = kwargs.get("history", [])
        config = kwargs.get("config", ContextConfig())
        summarizer = kwargs.get("summarizer")  # Optional LLM function for summaries

        if not history:
            return

        builder.metadata["original_message_count"] = len(history)

        # Apply compaction strategy
        if len(history) <= config.max_messages:
            # All messages fit
            builder.messages = history.copy()
            builder.metadata["included_message_count"] = len(history)
            return

        if config.compaction_strategy == CompactionStrategy.TRUNCATE_OLD:
            # Simple truncation - keep most recent
            builder.messages = history[-config.max_messages :]
            builder.metadata["included_message_count"] = config.max_messages

        elif config.compaction_strategy == CompactionStrategy.SLIDING_WINDOW:
            # Keep recent messages + summarize older ones
            recent = history[-config.summarize_after :]
            older = history[: -config.summarize_after]

            if older and summarizer:
                # Summarize older messages
                summary = self._summarize_messages(older, summarizer)
                summary_msg = {
                    "role": "system",
                    "content": f"[Earlier conversation summary: {summary}]",
                }
                builder.messages = [summary_msg] + recent
                builder.metadata["summarized_message_count"] = len(older)
            else:
                # No summarizer - just include recent
                builder.messages = recent

            builder.metadata["included_message_count"] = len(builder.messages)

        elif config.compaction_strategy == CompactionStrategy.IMPORTANCE:
            # Keep messages with tool calls and user questions
            important = []
            for msg in history:
                content = msg.get("content", "")
                role = msg.get("role", "")

                # Keep user messages, tool results, and important assistant responses
                is_important = (
                    role == "user"
                    or "tool_result" in str(content)
                    or isinstance(content, list)  # Tool calls
                )

                if is_important:
                    important.append(msg)

            # Ensure we have the most recent exchange
            if history[-2:] not in important[-2:]:
                important = important[:-2] + history[-2:]

            builder.messages = important[-config.max_messages :]
            builder.metadata["included_message_count"] = len(builder.messages)

    def _summarize_messages(self, messages: List[dict], summarizer: Callable) -> str:
        """Use LLM to summarize older messages."""
        # Build a simple text representation
        text_parts = []
        for msg in messages:
            role = msg.get("role", "unknown")
            content = msg.get("content", "")
            if isinstance(content, str):
                text_parts.append(f"{role}: {content[:200]}")

        conversation_text = "\n".join(text_parts)

        prompt = f"""Summarize this conversation excerpt in 2-3 sentences,
focusing on key decisions, facts learned, and tasks discussed:

{conversation_text}

Summary:"""

        try:
            return summarizer(prompt)
        except Exception as e:
            logger.warning(f"Failed to summarize: {e}")
            return "Earlier conversation about various topics."


class ToolsProcessor(ContextProcessor):
    """Selects and formats relevant tools."""

    @property
    def name(self) -> str:
        return "tools"

    def process(self, builder: "ContextBuilder", **kwargs) -> None:
        tools = kwargs.get("tools", [])
        task_hint = kwargs.get("task_hint", "")
        config = kwargs.get("config", ContextConfig())
        recent_tools = kwargs.get("recent_tools", [])  # Tools used recently

        if not tools:
            return

        if config.tool_selection == "all" or len(tools) <= config.max_tools:
            builder.tools = tools

        elif config.tool_selection == "recent_use" and recent_tools:
            # Prioritize recently used tools
            recent_set = set(recent_tools)
            prioritized = sorted(tools, key=lambda t: (t["name"] not in recent_set, t["name"]))
            builder.tools = prioritized[: config.max_tools]

        elif config.tool_selection == "relevance" and task_hint:
            # Score tools by relevance to task hint
            scored = []
            task_lower = task_hint.lower()

            for tool in tools:
                score = 0
                name = tool.get("name", "").lower()
                desc = tool.get("description", "").lower()

                # Simple keyword matching
                for word in task_lower.split():
                    if word in name:
                        score += 3
                    if word in desc:
                        score += 1

                scored.append((score, tool))

            scored.sort(key=lambda x: x[0], reverse=True)
            builder.tools = [t for _, t in scored[: config.max_tools]]

        else:
            builder.tools = tools[: config.max_tools]

        builder.metadata["tools_included"] = len(builder.tools)


class SkillsProcessor(ContextProcessor):
    """Injects skill documentation into context."""

    @property
    def name(self) -> str:
        return "skills"

    def process(self, builder: "ContextBuilder", **kwargs) -> None:
        skills = kwargs.get("skills")  # SkillLoader instance
        task_hint = kwargs.get("task_hint", "")

        if not skills:
            return

        skill_docs = skills.get_skill_docs()

        # Re-order skill doc sections by task hint relevance
        if task_hint and skill_docs:
            hint_words = set(task_hint.lower().split())
            sections = skill_docs.split("\n\n")
            sections.sort(
                key=lambda s: sum(1 for w in hint_words if w in s.lower()),
                reverse=True,
            )
            skill_docs = "\n\n".join(sections)

        if skill_docs:
            builder.add_system_section(
                "skills", f"## Available Skills\n\n{skill_docs}", priority=60
            )


# ============================================================
# CONTEXT BUILDER
# ============================================================


class ContextBuilder:
    """
    Accumulates context sections during processing.

    Processors write to this builder, which then compiles
    the final WorkingContext with proper ordering and formatting.
    """

    def __init__(self):
        self.system_sections: List[tuple] = []  # (name, content, priority)
        self.messages: List[Dict[str, Any]] = []
        self.tools: List[Dict[str, Any]] = []
        self.metadata: Dict[str, Any] = {}

    def add_system_section(self, name: str, content: str, priority: int = 50):
        """Add a section to the system prompt. Higher priority = earlier."""
        self.system_sections.append((name, content, priority))

    def compile_system_prompt(self) -> str:
        """Combine all system sections in priority order."""
        sorted_sections = sorted(self.system_sections, key=lambda x: -x[2])
        return "\n\n".join(content for _, content, _ in sorted_sections)

    def estimate_tokens(self) -> int:
        """
        Accurate token estimation using tiktoken.

        Falls back to improved heuristic if tiktoken unavailable.
        """
        counter = get_token_counter()

        # System prompt tokens
        system_tokens = counter.count(self.compile_system_prompt())

        # Message tokens (with overhead)
        message_tokens = counter.count_messages(self.messages)

        # Tool definition tokens
        tool_tokens = counter.count_tools(self.tools)

        return system_tokens + message_tokens + tool_tokens

    def compute_hash(self) -> str:
        """Compute content hash for caching."""
        content = (
            self.compile_system_prompt()
            + json.dumps(self.messages, sort_keys=True)
            + json.dumps(self.tools, sort_keys=True)
        )
        return hashlib.md5(content.encode()).hexdigest()[:12]


# ============================================================
# CONTEXT COMPILER
# ============================================================


class ContextCompiler:
    """
    Main entry point for context engineering.

    Orchestrates processors to transform raw storage into
    optimized working context for LLM consumption.

    Usage:
        compiler = ContextCompiler()
        context = compiler.compile(
            memory=agent.memory,
            history=conversation.messages,
            tools=tool_registry.get_schemas(),
            task_hint="help with email"
        )
    """

    def __init__(
        self,
        config: Optional[ContextConfig] = None,
        processors: Optional[List[ContextProcessor]] = None,
    ):
        self.config = config or ContextConfig()

        # Default processor pipeline
        self.processors = processors or [
            IdentityProcessor(),
            MemoryProcessor(),
            SkillsProcessor(),
            HistoryProcessor(),
            ToolsProcessor(),
        ]

    def compile(
        self,
        memory=None,
        history: Optional[List[Dict]] = None,
        tools: Optional[List[Dict]] = None,
        skills=None,
        task_hint: str = "",
        summarizer: Optional[Callable] = None,
        **kwargs,
    ) -> WorkingContext:
        """
        Compile storage into working context.

        Args:
            memory: Memory instance with stored facts
            history: Conversation message history
            tools: Available tool schemas
            skills: SkillLoader with skill documentation
            task_hint: Hint about current task for relevance scoring
            summarizer: Optional function(prompt) -> str for summarization
            **kwargs: Additional context (agent_name, provider_name, etc.)

        Returns:
            WorkingContext ready for LLM consumption
        """
        import time

        start = time.perf_counter()

        builder = ContextBuilder()

        # Run processors in order
        processor_kwargs = {
            "memory": memory,
            "history": history or [],
            "tools": tools or [],
            "skills": skills,
            "task_hint": task_hint,
            "config": self.config,
            "summarizer": summarizer,
            **kwargs,
        }

        for processor in self.processors:
            try:
                processor.process(builder, **processor_kwargs)
                logger.debug(f"Processor '{processor.name}' completed")
            except Exception as e:
                logger.error(f"Processor '{processor.name}' failed: {e}")

        # Build final context
        compilation_time = (time.perf_counter() - start) * 1000

        context = WorkingContext(
            system_prompt=builder.compile_system_prompt(),
            messages=builder.messages,
            tools=builder.tools,
            compilation_time_ms=round(compilation_time, 2),
            original_message_count=builder.metadata.get("original_message_count", 0),
            included_message_count=builder.metadata.get("included_message_count", 0),
            summarized_message_count=builder.metadata.get("summarized_message_count", 0),
            memory_entries_included=builder.metadata.get("memory_entries_included", 0),
            tools_included=builder.metadata.get("tools_included", 0),
            estimated_tokens=builder.estimate_tokens(),
            content_hash=builder.compute_hash(),
        )

        logger.info(
            f"Context compiled: {context.estimated_tokens} tokens, "
            f"{context.included_message_count}/{context.original_message_count} messages, "
            f"{compilation_time:.1f}ms"
        )

        return context

    def add_processor(self, processor: ContextProcessor, index: int = -1):
        """Add a custom processor to the pipeline."""
        if index == -1:
            self.processors.append(processor)
        else:
            self.processors.insert(index, processor)

    def remove_processor(self, name: str):
        """Remove a processor by name."""
        self.processors = [p for p in self.processors if p.name != name]

    async def compile_streaming(
        self,
        memory=None,
        history: Optional[List[Dict]] = None,
        tools: Optional[List[Dict]] = None,
        skills=None,
        task_hint: str = "",
        summarizer: Optional[Callable] = None,
        yield_messages: bool = True,
        yield_tools: bool = True,
        **kwargs,
    ):
        """
        Compile context with streaming output (v1.4.0 Phase 3).

        Yields ContextChunk objects as context is assembled, enabling:
        - Faster time-to-first-token
        - Progressive context delivery
        - Real-time progress monitoring

        Args:
            memory: Memory instance
            history: Conversation history
            tools: Tool schemas
            skills: Skill loader
            task_hint: Task hint for relevance
            summarizer: Summarization function
            yield_messages: Yield individual messages
            yield_tools: Yield individual tools
            **kwargs: Additional context

        Yields:
            ContextChunk objects as context is assembled

        Usage:
            async for chunk in compiler.compile_streaming(memory=mem, history=hist):
                if chunk.event == StreamingContextEvent.SYSTEM_COMPLETE:
                    # System prompt ready, can start warming LLM
                    pass
                elif chunk.event == StreamingContextEvent.COMPLETE:
                    context = chunk.data  # Final WorkingContext
        """
        # Delegate to StreamingContextCompiler with our config and processors
        streaming_compiler = StreamingContextCompiler(
            config=self.config, processors=self.processors
        )

        async for chunk in streaming_compiler.compile_streaming(
            memory=memory,
            history=history,
            tools=tools,
            skills=skills,
            task_hint=task_hint,
            summarizer=summarizer,
            yield_messages=yield_messages,
            yield_tools=yield_tools,
            **kwargs,
        ):
            yield chunk


# ============================================================
# CONTEXT CACHE
# ============================================================


class ContextCache:
    """
    Caches compiled contexts to avoid redundant processing.

    Uses content hash to detect when recompilation is needed.
    """

    def __init__(self, max_size: int = 100):
        self.max_size = max_size
        self._cache: Dict[str, WorkingContext] = {}
        self._access_order: List[str] = []

    def get(self, cache_key: str) -> Optional[WorkingContext]:
        """Get cached context if available."""
        if cache_key in self._cache:
            # Move to end of access order (LRU)
            self._access_order.remove(cache_key)
            self._access_order.append(cache_key)
            return self._cache[cache_key]
        return None

    def put(self, cache_key: str, context: WorkingContext):
        """Cache a compiled context."""
        if len(self._cache) >= self.max_size:
            # Evict oldest
            oldest = self._access_order.pop(0)
            del self._cache[oldest]

        self._cache[cache_key] = context
        self._access_order.append(cache_key)

    def invalidate(self, cache_key: str = None):
        """Invalidate specific or all cached contexts."""
        if cache_key:
            self._cache.pop(cache_key, None)
            if cache_key in self._access_order:
                self._access_order.remove(cache_key)
        else:
            self._cache.clear()
            self._access_order.clear()


# Singleton cache instance
_context_cache = ContextCache()


def get_context_cache() -> ContextCache:
    """Get the global context cache."""
    return _context_cache


# ============================================================
# STREAMING CONTEXT ASSEMBLY (v1.4.0 Phase 3)
# ============================================================


class StreamingContextEvent(Enum):
    """Events emitted during streaming context assembly."""

    START = "start"
    SYSTEM_SECTION = "system_section"
    SYSTEM_COMPLETE = "system_complete"
    MESSAGE = "message"
    MESSAGES_COMPLETE = "messages_complete"
    TOOL = "tool"
    TOOLS_COMPLETE = "tools_complete"
    COMPLETE = "complete"
    ERROR = "error"


@dataclass
class ContextChunk:
    """A chunk of context being streamed."""

    event: StreamingContextEvent
    data: Any = None
    section_name: str = ""
    index: int = 0
    total: int = 0
    cumulative_tokens: int = 0
    elapsed_ms: float = 0.0

    def to_dict(self) -> dict:
        return {
            "event": self.event.value,
            "data": self.data,
            "section_name": self.section_name,
            "index": self.index,
            "total": self.total,
            "cumulative_tokens": self.cumulative_tokens,
            "elapsed_ms": self.elapsed_ms,
        }


class StreamingContextBuilder:
    """
    Builds context incrementally, yielding chunks as they become available.

    This enables faster time-to-first-token by allowing the LLM to start
    processing the system prompt while messages and tools are still being assembled.

    Usage:
        async for chunk in compiler.compile_streaming(memory=mem, history=hist):
            if chunk.event == StreamingContextEvent.SYSTEM_COMPLETE:
                # System prompt ready, can start warming up LLM
                pass
            elif chunk.event == StreamingContextEvent.MESSAGE:
                # Each message as it's processed
                pass
    """

    def __init__(self, config: Optional[ContextConfig] = None):
        self.config = config or ContextConfig()
        self._counter = get_token_counter()
        self._cumulative_tokens = 0
        self._start_time = None

    def _elapsed_ms(self) -> float:
        if self._start_time is None:
            return 0.0
        import time

        return (time.perf_counter() - self._start_time) * 1000

    async def _process_system_sections(self, builder: ContextBuilder, **kwargs):
        """Process system sections and yield chunks."""
        # Sort by priority
        sorted_sections = sorted(builder.system_sections, key=lambda x: -x[2])

        for i, (name, content, priority) in enumerate(sorted_sections):
            tokens = self._counter.count(content)
            self._cumulative_tokens += tokens

            yield ContextChunk(
                event=StreamingContextEvent.SYSTEM_SECTION,
                data=content,
                section_name=name,
                index=i,
                total=len(sorted_sections),
                cumulative_tokens=self._cumulative_tokens,
                elapsed_ms=self._elapsed_ms(),
            )

        yield ContextChunk(
            event=StreamingContextEvent.SYSTEM_COMPLETE,
            data=builder.compile_system_prompt(),
            cumulative_tokens=self._cumulative_tokens,
            elapsed_ms=self._elapsed_ms(),
        )

    async def _process_messages(self, messages: List[Dict[str, Any]], config: ContextConfig):
        """Process messages and yield chunks."""
        total = len(messages)

        for i, message in enumerate(messages):
            # Count tokens for this message
            msg_tokens = self._counter.count_messages([message])
            self._cumulative_tokens += msg_tokens

            yield ContextChunk(
                event=StreamingContextEvent.MESSAGE,
                data=message,
                index=i,
                total=total,
                cumulative_tokens=self._cumulative_tokens,
                elapsed_ms=self._elapsed_ms(),
            )

        yield ContextChunk(
            event=StreamingContextEvent.MESSAGES_COMPLETE,
            data={"count": total},
            total=total,
            cumulative_tokens=self._cumulative_tokens,
            elapsed_ms=self._elapsed_ms(),
        )

    async def _process_tools(self, tools: List[Dict[str, Any]]):
        """Process tools and yield chunks."""
        total = len(tools)

        for i, tool in enumerate(tools):
            tool_tokens = self._counter.count(json.dumps(tool))
            self._cumulative_tokens += tool_tokens

            yield ContextChunk(
                event=StreamingContextEvent.TOOL,
                data=tool,
                section_name=tool.get("function", {}).get("name", tool.get("name", f"tool_{i}")),
                index=i,
                total=total,
                cumulative_tokens=self._cumulative_tokens,
                elapsed_ms=self._elapsed_ms(),
            )

        yield ContextChunk(
            event=StreamingContextEvent.TOOLS_COMPLETE,
            data={"count": total},
            total=total,
            cumulative_tokens=self._cumulative_tokens,
            elapsed_ms=self._elapsed_ms(),
        )


class StreamingContextCompiler:
    """
    Async streaming version of ContextCompiler.

    Yields context chunks as they become available, enabling:
    - Faster time-to-first-token
    - Progressive context delivery
    - Real-time progress monitoring
    - Parallel processing of context sections

    Usage:
        compiler = StreamingContextCompiler()

        async for chunk in compiler.compile_streaming(
            memory=agent.memory,
            history=conversation_history,
            tools=tool_schemas
        ):
            print(f"{chunk.event.value}: {chunk.cumulative_tokens} tokens")

            if chunk.event == StreamingContextEvent.COMPLETE:
                context = chunk.data  # Final WorkingContext
    """

    def __init__(
        self,
        config: Optional[ContextConfig] = None,
        processors: Optional[List[ContextProcessor]] = None,
    ):
        self.config = config or ContextConfig()
        self.processors = processors or [
            IdentityProcessor(),
            MemoryProcessor(),
            SkillsProcessor(),
            HistoryProcessor(),
            ToolsProcessor(),
        ]
        self._counter = get_token_counter()

    async def compile_streaming(
        self,
        memory=None,
        history: Optional[List[Dict]] = None,
        tools: Optional[List[Dict]] = None,
        skills=None,
        task_hint: str = "",
        summarizer: Optional[Callable] = None,
        yield_messages: bool = True,
        yield_tools: bool = True,
        **kwargs,
    ):
        """
        Compile context with streaming output.

        Args:
            memory: Memory instance
            history: Conversation history
            tools: Tool schemas
            skills: Skill loader
            task_hint: Task hint for relevance
            summarizer: Summarization function
            yield_messages: Yield individual messages (vs batch)
            yield_tools: Yield individual tools (vs batch)
            **kwargs: Additional context

        Yields:
            ContextChunk objects as context is assembled
        """
        import time

        start_time = time.perf_counter()
        cumulative_tokens = 0

        def elapsed_ms():
            return (time.perf_counter() - start_time) * 1000

        # Yield start event
        yield ContextChunk(event=StreamingContextEvent.START, elapsed_ms=0.0)

        try:
            # Build context using processors (sync, but we'll stream the results)
            builder = ContextBuilder()

            processor_kwargs = {
                "memory": memory,
                "history": history or [],
                "tools": tools or [],
                "skills": skills,
                "task_hint": task_hint,
                "config": self.config,
                "summarizer": summarizer,
                **kwargs,
            }

            # Run processors
            for processor in self.processors:
                try:
                    processor.process(builder, **processor_kwargs)
                except Exception as e:
                    logger.error(f"Processor '{processor.name}' failed: {e}")
                    yield ContextChunk(
                        event=StreamingContextEvent.ERROR,
                        data={"processor": processor.name, "error": str(e)},
                        elapsed_ms=elapsed_ms(),
                    )

            # Stream system sections
            sorted_sections = sorted(builder.system_sections, key=lambda x: -x[2])
            for i, (name, content, priority) in enumerate(sorted_sections):
                tokens = self._counter.count(content)
                cumulative_tokens += tokens

                yield ContextChunk(
                    event=StreamingContextEvent.SYSTEM_SECTION,
                    data=content,
                    section_name=name,
                    index=i,
                    total=len(sorted_sections),
                    cumulative_tokens=cumulative_tokens,
                    elapsed_ms=elapsed_ms(),
                )

            # Yield system complete
            system_prompt = builder.compile_system_prompt()
            yield ContextChunk(
                event=StreamingContextEvent.SYSTEM_COMPLETE,
                data=system_prompt,
                cumulative_tokens=cumulative_tokens,
                elapsed_ms=elapsed_ms(),
            )

            # Stream messages
            if yield_messages:
                for i, message in enumerate(builder.messages):
                    msg_tokens = self._counter.count_messages([message])
                    cumulative_tokens += msg_tokens

                    yield ContextChunk(
                        event=StreamingContextEvent.MESSAGE,
                        data=message,
                        index=i,
                        total=len(builder.messages),
                        cumulative_tokens=cumulative_tokens,
                        elapsed_ms=elapsed_ms(),
                    )
            else:
                cumulative_tokens += self._counter.count_messages(builder.messages)

            yield ContextChunk(
                event=StreamingContextEvent.MESSAGES_COMPLETE,
                data={"count": len(builder.messages)},
                total=len(builder.messages),
                cumulative_tokens=cumulative_tokens,
                elapsed_ms=elapsed_ms(),
            )

            # Stream tools
            if yield_tools:
                for i, tool in enumerate(builder.tools):
                    tool_tokens = self._counter.count(json.dumps(tool))
                    cumulative_tokens += tool_tokens

                    tool_name = tool.get("function", {}).get("name", tool.get("name", f"tool_{i}"))
                    yield ContextChunk(
                        event=StreamingContextEvent.TOOL,
                        data=tool,
                        section_name=tool_name,
                        index=i,
                        total=len(builder.tools),
                        cumulative_tokens=cumulative_tokens,
                        elapsed_ms=elapsed_ms(),
                    )
            else:
                cumulative_tokens += self._counter.count_tools(builder.tools)

            yield ContextChunk(
                event=StreamingContextEvent.TOOLS_COMPLETE,
                data={"count": len(builder.tools)},
                total=len(builder.tools),
                cumulative_tokens=cumulative_tokens,
                elapsed_ms=elapsed_ms(),
            )

            # Build final context
            compilation_time = elapsed_ms()
            context = WorkingContext(
                system_prompt=system_prompt,
                messages=builder.messages,
                tools=builder.tools,
                compilation_time_ms=round(compilation_time, 2),
                original_message_count=builder.metadata.get("original_message_count", 0),
                included_message_count=builder.metadata.get("included_message_count", 0),
                summarized_message_count=builder.metadata.get("summarized_message_count", 0),
                memory_entries_included=builder.metadata.get("memory_entries_included", 0),
                tools_included=builder.metadata.get("tools_included", 0),
                estimated_tokens=cumulative_tokens,
                content_hash=builder.compute_hash(),
            )

            # Yield complete event with final context
            yield ContextChunk(
                event=StreamingContextEvent.COMPLETE,
                data=context,
                cumulative_tokens=cumulative_tokens,
                elapsed_ms=compilation_time,
            )

        except Exception as e:
            logger.error(f"Streaming compilation failed: {e}")
            yield ContextChunk(
                event=StreamingContextEvent.ERROR, data={"error": str(e)}, elapsed_ms=elapsed_ms()
            )

    def compile(
        self,
        memory=None,
        history: Optional[List[Dict]] = None,
        tools: Optional[List[Dict]] = None,
        skills=None,
        task_hint: str = "",
        summarizer: Optional[Callable] = None,
        **kwargs,
    ) -> WorkingContext:
        """
        Synchronous compile (non-streaming).

        For streaming, use compile_streaming() instead.
        """
        import asyncio

        async def collect():
            context = None
            async for chunk in self.compile_streaming(
                memory=memory,
                history=history,
                tools=tools,
                skills=skills,
                task_hint=task_hint,
                summarizer=summarizer,
                yield_messages=False,
                yield_tools=False,
                **kwargs,
            ):
                if chunk.event == StreamingContextEvent.COMPLETE:
                    context = chunk.data
            return context

        # Run async in sync context
        try:
            loop = asyncio.get_event_loop()
            if loop.is_running():
                # Already in async context, use nest_asyncio or return coroutine
                import concurrent.futures

                with concurrent.futures.ThreadPoolExecutor() as pool:
                    future = pool.submit(asyncio.run, collect())
                    return future.result()
            else:
                return loop.run_until_complete(collect())
        except RuntimeError:
            return asyncio.run(collect())


def compile_context_streaming(
    memory=None,
    history: Optional[List[Dict]] = None,
    tools: Optional[List[Dict]] = None,
    config: Optional[ContextConfig] = None,
    **kwargs,
):
    """
    Convenience function for streaming context compilation.

    Usage:
        async for chunk in compile_context_streaming(memory=mem, history=hist):
            if chunk.event == StreamingContextEvent.SYSTEM_COMPLETE:
                print(f"System prompt ready: {chunk.cumulative_tokens} tokens")
    """
    compiler = StreamingContextCompiler(config=config)
    return compiler.compile_streaming(memory=memory, history=history, tools=tools, **kwargs)


# Global streaming compiler instance
_streaming_compiler: Optional[StreamingContextCompiler] = None


def get_streaming_compiler(config: Optional[ContextConfig] = None) -> StreamingContextCompiler:
    """Get or create the global streaming context compiler."""
    global _streaming_compiler
    if _streaming_compiler is None or config is not None:
        _streaming_compiler = StreamingContextCompiler(config=config)
    return _streaming_compiler
