"""Core cognitive memory system.

Orchestrates storage, embeddings, extraction, and working memory
into a coherent API for human-like memory operations.
"""

from __future__ import annotations

import asyncio
import logging
import math
from typing import Any, cast

from recollect.buffer import WorkingMemory
from recollect.config import MemoryConfig
from recollect.config import config as default_config
from recollect.datetime_utils import memory_timestamp_for_comparison, now_utc
from recollect.embeddings import FastEmbedProvider
from recollect.exceptions import (
    ExtractionError,
    SessionNotFoundError,
    StorageError,
    TraceNotFoundError,
)
from recollect.extraction import PatternExtractor
from recollect.llm.types import ExtractionResult, Message
from recollect.models import (
    Association,
    ConceptEmbedding,
    ConsolidationResult,
    EntityRelation,
    FactCategory,
    FactStatus,
    HealthStatus,
    MemoryStats,
    MemoryTrace,
    PersonaFact,
    Session,
    Thought,
    TraceConcept,
    TraceEntity,
    _clamp_strength,
    apply_activation_boost,
    apply_displacement_decay,
    apply_retrieval_boost,
    apply_time_decay,
)
from recollect.storage_context import StorageContext, create_storage_context

logger = logging.getLogger(__name__)


def _cosine_similarity(a: list[float], b: list[float]) -> float:
    """Compute cosine similarity between two vectors."""
    dot = sum(x * y for x, y in zip(a, b, strict=False))
    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.0 or norm_b == 0.0:
        return 0.0
    return dot / (norm_a * norm_b)


def _estimate_tokens(text: str) -> int:
    """Estimate token count from text (chars / 4)."""
    return max(1, len(text) // 4)


def _compute_decay_rate(
    base_rate: float,
    significance: float,
    emotional_valence: float,
    *,
    sig_reduction: float = 0.7,
    val_reduction: float = 0.5,
) -> float:
    """Compute decay rate reduced by significance and emotional intensity.

    High significance and strong emotion make memories last longer.
    """
    sig_factor = 1.0 - (significance * sig_reduction)
    emo_factor = 1.0 - (abs(emotional_valence) * val_reduction)
    return base_rate * sig_factor * emo_factor


def _rrf_fuse(
    ranked_lists: dict[str, list[str]],
    k: int = 60,
) -> dict[str, float]:
    """Reciprocal Rank Fusion across multiple signal sources.

    Each source provides a ranked list of trace IDs (best first).
    Returns {trace_id: fused_score} ordered by score descending.
    """
    scores: dict[str, float] = {}
    for source_ids in ranked_lists.values():
        for rank, trace_id in enumerate(source_ids, start=1):
            scores[trace_id] = scores.get(trace_id, 0.0) + 1.0 / (k + rank)
    return scores


_PREDICATE_ALIASES: dict[str, str] = {
    "started_at": "works_at",
    "employed_at": "works_at",
    "joined": "works_at",
    "has_allergy": "is_allergic_to",
    "allergic_to": "is_allergic_to",
    "likes": "prefers",
    "enjoys": "prefers",
    "lives_in": "located_in",
    "moved_to": "located_in",
}


def _canonicalize_predicate(predicate: str) -> str:
    """Normalize predicate to canonical form via alias lookup."""
    return _PREDICATE_ALIASES.get(predicate, predicate)


_VALID_CATEGORIES: frozenset[str] = frozenset(
    {
        "health",
        "dietary",
        "identity",
        "relationship",
        "preference",
        "schedule",
        "constraint",
        "general",
    }
)


_CATEGORY_SCOPE_MAP: dict[str, str] = {
    "health": "health/safety",
    "dietary": "health/safety",
    "constraint": "health/safety",
    "relationship": "social",
    "identity": "identity",
    "preference": "preference",
    "schedule": "routine",
    "general": "general",
}


def _category_to_scope(category: FactCategory) -> str:
    """Map a fact category to a retrieval scope."""
    return _CATEGORY_SCOPE_MAP.get(category, "general")


_FAST_TRACK_CATEGORIES: frozenset[str] = frozenset(
    {
        "health",
        "dietary",
        "constraint",
    }
)


def _should_fast_track(category: str, confidence: float) -> bool:
    """Check if a fact should skip candidate stage.

    Only safety-critical categories fast-track. The confidence
    parameter is accepted but not used for fast-track decisions.
    """
    return category in _FAST_TRACK_CATEGORIES


def _find_contradicting_fact(
    existing: list[PersonaFact], new_fact: PersonaFact
) -> PersonaFact | None:
    """Find existing fact with same subject+predicate but different object."""
    canonical_new = _canonicalize_predicate(new_fact.predicate)
    for fact in existing:
        canonical_existing = _canonicalize_predicate(fact.predicate)
        if canonical_existing == canonical_new and fact.object != new_fact.object:
            return fact
    return None


def _find_exact_duplicate(
    existing: list[PersonaFact], new_fact: PersonaFact
) -> PersonaFact | None:
    """Find existing fact with same subject+predicate+object."""
    canonical_new = _canonicalize_predicate(new_fact.predicate)
    for fact in existing:
        canonical_existing = _canonicalize_predicate(fact.predicate)
        if canonical_existing == canonical_new and fact.object == new_fact.object:
            return fact
    return None


_QUERY_STOP_WORDS = frozenset(
    {
        "the",
        "a",
        "an",
        "this",
        "that",
        "these",
        "those",
        "what",
        "how",
        "when",
        "where",
        "why",
        "who",
        "which",
        "my",
        "our",
        "his",
        "her",
        "their",
        "its",
        "i",
        "we",
        "it",
        "he",
        "she",
        "they",
        "you",
        "do",
        "does",
        "did",
        "is",
        "are",
        "was",
        "were",
        "can",
        "could",
        "should",
        "would",
        "will",
        "shall",
        "have",
        "has",
        "had",
        "not",
        "no",
        "yes",
        "if",
        "but",
        "and",
        "or",
        "so",
        "yet",
        "for",
        "got",
        "get",
        "go",
        "plan",
        "book",
        "find",
        "need",
        "want",
        "let",
        "tell",
        "show",
        "give",
        "take",
        "make",
        "help",
        "remember",
    }
)


class CognitiveMemory:
    """Human-like memory system with cognitive model.

    Combines working memory buffer, long-term storage with semantic
    search, spreading activation, and strength-based consolidation.
    """

    def __init__(
        self,
        *,
        storage: StorageContext | None = None,
        embeddings: FastEmbedProvider | None = None,
        extractor: PatternExtractor | None = None,
        config: MemoryConfig | None = None,
    ) -> None:
        self._config = config or default_config
        self._storage = storage or create_storage_context()
        self._embeddings = embeddings or FastEmbedProvider()
        self._extractor = extractor
        self._buffer = WorkingMemory(self._config.working_memory_capacity)
        self._connected = False

    # -- Lifecycle --

    async def connect(self, db_url: str | None = None) -> None:
        """Initialize storage connection and schema."""
        if db_url:
            self._storage = create_storage_context(db_url)
        await self._storage.initialize()
        self._connected = True
        logger.info("CognitiveMemory connected")

    async def close(self) -> None:
        """Close storage connection."""
        await self._storage.close()
        self._connected = False
        logger.info("CognitiveMemory closed")

    async def __aenter__(self) -> CognitiveMemory:
        await self.connect()
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: object,
    ) -> None:
        await self.close()

    # -- Core operations --

    async def experience(
        self,
        content: str,
        *,
        context: dict[str, Any] | None = None,
        session_id: str | None = None,
        user_id: str | None = None,
    ) -> MemoryTrace:
        """Encode a new experience as a memory trace.

        1. Extract patterns (concepts, entities, emotion)
        2. Generate embedding vector
        3. Create trace and add to working memory
        4. Handle displacement decay if buffer was full
        5. Store in long-term storage with temporal association
        """
        if not content or not content.strip():
            raise ValueError("Content must be a non-empty string")

        result = await self._extract_pattern(content)
        embedding = await self._embeddings.generate_embedding(content)

        decay_rate = _compute_decay_rate(
            base_rate=float(self._config.get("memory.decay_rate", 0.1)),
            significance=result.significance,
            emotional_valence=result.emotional_valence,
            sig_reduction=float(self._config.get("decay.significance_reduction", 0.7)),
            val_reduction=float(self._config.get("decay.valence_reduction", 0.5)),
        )
        trace = MemoryTrace(
            content=content,
            pattern=result.model_dump(),
            context=context or {},
            embedding=embedding,
            emotional_valence=result.emotional_valence,
            significance=result.significance,
            decay_rate=decay_rate,
            session_id=session_id,
            user_id=user_id,
        )

        if session_id is not None:
            await self._ensure_session(session_id, user_id)

        displaced = self._buffer.add(trace)
        if displaced is not None:
            decayed = apply_displacement_decay(displaced)
            await self._storage.traces.update_trace_strength(
                displaced.id, decayed.strength
            )

        await self._storage.traces.store_trace(trace)

        # Independent operations: temporal link + extraction links + concept embeddings
        await asyncio.gather(
            self._create_temporal_association(trace),
            self._store_extraction_links(trace, result),
            self._embed_trace_concepts(trace, result),
        )

        # Depend on extraction links, but independent of each other
        await asyncio.gather(
            self._create_entity_associations(trace, result),
            self._create_concept_associations(trace, result),
        )

        if self._config.get("persona.auto_extract", True):
            await self._extract_persona_facts(trace, result)

        await self._extract_entity_relations(trace, result)

        logger.debug("Experienced: %s (strength=%.2f)", trace.id[:8], trace.strength)
        return trace

    async def think_about(
        self,
        query: str,
        *,
        token_budget: int = 2000,
        session_id: str | None = None,
        user_id: str | None = None,
    ) -> list[Thought]:
        """Retrieve relevant memories within a token budget.

        1. Generate query embedding
        2. Check working memory and long-term storage
        3. Spread activation from top candidates
        4. Merge, rank, select within budget
        5. Apply retrieval/activation boosts
        """
        if not query or not query.strip():
            raise ValueError("Query must be a non-empty string")
        if token_budget <= 0:
            raise ValueError("Token budget must be positive")

        query_embedding = await self._embeddings.generate_embedding(query)

        search_limit = int(self._config.get("retrieval.search_limit", 10))
        wm_candidates = self._search_working_memory(query_embedding)
        storage_scored = await self._storage.vectors.search_semantic(
            query_embedding,
            limit=search_limit,
            session_id=session_id,
            user_id=user_id,
        )
        storage_candidates = [trace for trace, _ in storage_scored]

        entity_matches = await self._match_query_entities(query)
        if entity_matches:
            entity_trace_ids = [tid for tid, _ in entity_matches]
            entity_candidates = await self._storage.traces.get_traces_bulk(
                entity_trace_ids
            )
            seen_ids = {t.id for t, _ in storage_scored}
            for t in entity_candidates:
                if t.id not in seen_ids:
                    sim = (
                        _cosine_similarity(query_embedding, t.embedding)
                        if t.embedding is not None
                        else 0.0
                    )
                    storage_scored.append((t, sim))
                    storage_candidates.append(t)
                    seen_ids.add(t.id)

        seed_count = int(self._config.get("retrieval.spread_seed_count", 3))
        activated = await self._spread_from_candidates(storage_candidates[:seed_count])

        all_candidates = await self._merge_candidates(
            query_embedding,
            wm_candidates,
            storage_scored,
            activated,
            entity_matches=entity_matches if entity_matches else None,
        )

        if session_id is not None:
            all_candidates = [
                (t, s) for t, s in all_candidates if t.session_id == session_id
            ]

        selected = self._select_within_budget(all_candidates, token_budget)

        thoughts = await self._boost_and_build_thoughts(selected)
        await self._boost_unselected(all_candidates, thoughts)

        persona_facts, semantic_scores = await self._find_relevant_persona_facts(
            query,
            query_embedding,
            user_id=user_id,
        )
        if persona_facts:
            fact_thoughts = self._persona_facts_to_thoughts(
                persona_facts,
                semantic_scores,
            )
            thoughts = thoughts + fact_thoughts

        max_total = int(self._config.get("retrieval.max_retrievals", 7))
        min_traces = int(self._config.get("retrieval.min_trace_slots", 2))
        thoughts = self._assemble_with_trace_guarantee(
            thoughts,
            max_total,
            min_traces,
        )

        logger.debug(
            "think_about: %d thoughts (budget=%d)",
            len(thoughts),
            token_budget,
        )
        return thoughts

    @staticmethod
    def _assemble_with_trace_guarantee(
        thoughts: list[Thought],
        max_total: int,
        min_traces: int,
    ) -> list[Thought]:
        """Assemble final thoughts guaranteeing minimum trace slots.

        Pinned facts get priority but cannot consume more than
        (max_total - min_traces) slots, preserving space for traces.
        Remaining slots filled by relevance regardless of pin status.
        """
        pinned = sorted(
            [t for t in thoughts if t.pinned],
            key=lambda t: -t.relevance,
        )
        unpinned = sorted(
            [t for t in thoughts if not t.pinned],
            key=lambda t: -t.relevance,
        )
        max_pinned = max(0, max_total - min_traces)
        selected = pinned[:max_pinned] + unpinned[:min_traces]
        remaining = pinned[max_pinned:] + unpinned[min_traces:]
        remaining.sort(key=lambda t: -t.relevance)
        selected += remaining[: max_total - len(selected)]
        selected.sort(key=lambda t: -t.relevance)
        return selected

    async def consolidate(self, threshold: float | None = None) -> ConsolidationResult:
        """Consolidate or forget pending traces based on strength.

        Traces above threshold are consolidated to long-term memory.
        Traces past their grace period that remain weak are forgotten.
        Young weak traces are kept pending with updated decay.
        """
        threshold = threshold or self._config.consolidation_threshold
        grace_hours = float(self._config.get("forgetting.grace_period_hours", 24))
        batch_size = int(self._config.get("consolidation.batch_size", 50))
        traces = await self._storage.traces.get_unconsolidated_traces(limit=batch_size)

        consolidated = 0
        forgotten = 0
        still_pending = 0

        for trace in traces:
            decayed = apply_time_decay(trace)
            result = await self._consolidate_one(trace, decayed, threshold, grace_hours)
            if result == "consolidated":
                consolidated += 1
            elif result == "forgotten":
                forgotten += 1
            else:
                still_pending += 1

        logger.info(
            "Consolidation: %d consolidated, %d forgotten, %d pending",
            consolidated,
            forgotten,
            still_pending,
        )
        return ConsolidationResult(
            consolidated=consolidated,
            forgotten=forgotten,
            still_pending=still_pending,
        )

    async def forget(self, trace_id: str) -> bool:
        """Explicitly forget a memory trace and its derived facts."""
        if not trace_id:
            raise ValueError("Trace ID must be a non-empty string")
        deleted = await self._storage.traces.delete_trace(trace_id)
        if not deleted:
            raise TraceNotFoundError(f"Trace {trace_id} not found")
        await self._cleanup_trace_facts(trace_id)
        logger.debug("Forgot trace: %s", trace_id[:8])
        return True

    async def _cleanup_trace_facts(self, trace_id: str) -> None:
        """Remove persona facts and entity relations linked to a trace."""
        try:
            facts = await self._storage.facts.get_persona_facts()
            for fact in facts:
                if fact.source_trace_id == trace_id:
                    await self._storage.concept_embeddings.delete_by_owner(
                        "fact", fact.id
                    )
                    await self._storage.facts.delete_persona_fact(fact.id)
            await self._storage.concept_embeddings.delete_by_owner("trace", trace_id)
            await self._storage.entity_relations.delete_by_trace(trace_id)
        except (StorageError, OSError):
            logger.warning(
                "Cleanup of derived data failed for trace %s",
                trace_id[:8],
            )

    async def reinforce(self, trace_id: str, *, factor: float = 1.1) -> MemoryTrace:
        """Manually strengthen a memory trace."""
        if not trace_id:
            raise ValueError("Trace ID must be a non-empty string")
        if factor <= 0.0:
            raise ValueError("Reinforcement factor must be positive")
        trace = await self._storage.traces.get_trace(trace_id)
        if trace is None:
            raise TraceNotFoundError(f"Trace {trace_id} not found")
        new_strength = _clamp_strength(trace.strength * factor)
        await self._storage.traces.update_trace_strength(trace_id, new_strength)
        return trace.model_copy(update={"strength": new_strength})

    # -- Introspection --

    def active_traces(self) -> list[MemoryTrace]:
        """Get traces currently in working memory."""
        return self._buffer.get_active()

    async def associations(self, trace_id: str) -> list[Association]:
        """Get all associations for a trace."""
        return await self._storage.associations.get_associations(trace_id)

    async def timeline(self, limit: int = 20) -> list[MemoryTrace]:
        """Get most recent memory traces."""
        if limit <= 0:
            raise ValueError("Limit must be positive")
        return await self._storage.traces.get_recent_traces(limit)

    def stats(self) -> MemoryStats:
        """Get memory system statistics."""
        wm_stats = self._buffer.get_stats()
        return MemoryStats(
            working_memory_items=wm_stats["current_items"],
            working_memory_capacity=wm_stats["capacity"],
            working_memory_utilization=wm_stats["utilization"],
            total_seen=wm_stats["total_seen"],
            total_displaced=wm_stats["total_displaced"],
            displacement_rate=wm_stats["displacement_rate"],
            connected=self._connected,
        )

    def health(self) -> HealthStatus:
        """Get system health status."""
        return HealthStatus(
            status="ok" if self._connected else "disconnected",
        )

    async def pin(self, trace_id: str) -> PersonaFact:
        """Create a persona fact from an existing trace."""
        if not trace_id:
            raise ValueError("Trace ID must be a non-empty string")
        trace = await self._storage.traces.get_trace(trace_id)
        if trace is None:
            raise TraceNotFoundError(f"Trace {trace_id} not found")
        fact = PersonaFact(
            subject="user",
            predicate="noted",
            object=trace.content or "",
            content=trace.content or "",
            source_trace_id=trace_id,
            confidence=1.0,
            status="pinned",
        )
        await self._storage.facts.store_persona_fact(fact)
        return fact

    async def unpin(self, fact_id: str) -> bool:
        """Demote a pinned fact back to promoted status."""
        if not fact_id:
            raise ValueError("Fact ID must be a non-empty string")
        try:
            await self._storage.facts.update_fact_status(fact_id, "promoted")
            return True
        except StorageError:
            return False

    async def facts(self, subject: str | None = None) -> list[PersonaFact]:
        """List persona facts, optionally filtered by subject."""
        return await self._storage.facts.get_persona_facts(subject)

    async def start_session(
        self,
        *,
        user_id: str,
        title: str = "",
        session_id: str | None = None,
        metadata: dict[str, Any] | None = None,
    ) -> Session:
        """Start a new conversation session."""
        import uuid as _uuid

        session = Session(
            id=session_id or str(_uuid.uuid4()),
            user_id=user_id,
            title=title,
            metadata=metadata or {},
        )
        await self._storage.sessions.create_session(session)
        return session

    async def end_session(self, session_id: str) -> Session:
        """Mark a session as ended."""
        session = await self._storage.sessions.get_session(session_id)
        if session is None:
            raise SessionNotFoundError(f"Session {session_id} not found")
        await self._storage.sessions.end_session(session_id)
        return session.model_copy(update={"status": "ended", "ended_at": now_utc()})

    async def summarize_session(self, session_id: str) -> MemoryTrace:
        """Generate a summary trace from all memories in a session."""
        session = await self._storage.sessions.get_session(session_id)
        if session is None:
            raise SessionNotFoundError(f"Session {session_id} not found")
        traces = await self._storage.traces.get_traces_by_session(session_id)
        if not traces:
            raise ValueError(f"Session {session_id} has no traces")

        summary_text = await self._generate_session_summary(traces, session)
        embedding = await self._embeddings.generate_embedding(summary_text)
        result = await self._extract_pattern(summary_text)

        summary_trace = MemoryTrace(
            content=summary_text,
            pattern={
                **result.model_dump(),
                "session_summary": True,
                "source_session_id": session_id,
                "trace_count": len(traces),
            },
            embedding=embedding,
            strength=float(self._config.get("session.summary_strength", 0.8)),
            significance=float(self._config.get("session.summary_significance", 0.7)),
            decay_rate=float(self._config.get("session.summary_decay_rate", 0.02)),
            emotional_valence=result.emotional_valence,
            session_id=session_id,
            user_id=session.user_id,
        )

        await self._storage.traces.store_trace(summary_trace)
        await self._storage.sessions.update_session(
            session_id,
            status="summarized",
            summary_trace_id=summary_trace.id,
        )

        await asyncio.gather(
            self._store_extraction_links(summary_trace, result),
            self._embed_trace_concepts(summary_trace, result),
        )

        return summary_trace

    async def get_sessions(
        self,
        user_id: str,
        *,
        limit: int = 50,
    ) -> list[Session]:
        """List sessions for a user."""
        return await self._storage.sessions.get_sessions(user_id, limit=limit)

    # -- Private: experience helpers --

    async def _extract_pattern(self, content: str) -> ExtractionResult:
        """Extract patterns, falling back to empty pattern if no extractor."""
        if self._extractor is None:
            return ExtractionResult()
        try:
            return await self._extractor.extract(content)
        except (ExtractionError, ValueError, KeyError):
            logger.warning("Extraction failed, using empty pattern")
            return ExtractionResult()

    async def _ensure_session(
        self,
        session_id: str,
        user_id: str | None,
    ) -> None:
        """Create session if it does not exist (lazy creation)."""
        existing = await self._storage.sessions.get_session(session_id)
        if existing is None:
            if user_id is None:
                raise ValueError("user_id required when creating a new session")
            session = Session(id=session_id, user_id=user_id)
            await self._storage.sessions.create_session(session)

    async def _generate_session_summary(
        self,
        traces: list[MemoryTrace],
        session: Session,
    ) -> str:
        """Generate session summary via LLM or fallback concatenation."""
        if self._extractor is None:
            contents = [t.content for t in traces if t.content]
            return f"Session ({len(traces)} memories): " + " | ".join(contents[:10])

        trace_texts = [
            f"{i}. {t.content}" for i, t in enumerate(traces, 1) if t.content
        ]
        prompt = (
            "Summarize this conversation session into a single cohesive paragraph. "
            "Capture key topics, decisions, facts learned, and unresolved questions. "
            "Include all proper nouns, numbers, and specific details. "
            "This summary will be embedded for semantic retrieval.\n\n"
            f"Session: {session.title or '(untitled)'}\n" + "\n".join(trace_texts)
        )

        max_tokens = int(self._config.get("session.summary_max_tokens", 1024))
        messages = [
            Message(
                role="system",
                content=(
                    "You are a memory consolidation system. Produce a dense, "
                    "factual summary paragraph. No filler. Include all proper nouns, "
                    "numbers, and specific details."
                ),
            ),
            Message(role="user", content=prompt),
        ]

        try:
            return (
                await self._extractor._provider.complete(
                    messages,
                    max_tokens=max_tokens,
                    temperature=0.0,
                )
            ).strip()
        except (ExtractionError, ValueError, RuntimeError, OSError):
            logger.warning("Session summary LLM failed, using fallback")
            contents = [t.content for t in traces if t.content]
            return f"Session ({len(traces)} memories): " + " | ".join(contents[:10])

    async def _create_temporal_association(self, trace: MemoryTrace) -> None:
        """Create temporal association to the previous trace in buffer."""
        active = self._buffer.get_active()
        if len(active) < 2:
            return
        previous = active[-2]
        weight = float(self._config.get("associations.temporal_weight", 0.5))
        association = Association(
            source_trace_id=previous.id,
            target_trace_id=trace.id,
            association_type="temporal",
            weight=weight,
        )
        try:
            await self._storage.associations.store_association(association)
        except (StorageError, OSError):
            logger.warning(
                "Failed to create temporal association %s -> %s",
                previous.id[:8],
                trace.id[:8],
            )

    async def _store_extraction_links(
        self, trace: MemoryTrace, result: ExtractionResult
    ) -> None:
        """Store entity and concept links from extraction result."""
        entities = [
            TraceEntity(
                entity_name=e.name,
                entity_type=e.entity_type,
                trace_id=trace.id,
            )
            for e in result.entities
        ]
        concepts = [TraceConcept(concept=c, trace_id=trace.id) for c in result.concepts]
        try:
            await self._storage.entities.store_trace_entities(trace.id, entities)
            await self._storage.entities.store_trace_concepts(trace.id, concepts)
        except (StorageError, OSError):
            logger.warning("Failed to store extraction links for %s", trace.id[:8])

    async def _embed_trace_concepts(
        self,
        trace: MemoryTrace,
        result: ExtractionResult,
    ) -> None:
        """Embed each extracted concept individually for concept attention."""
        if not result.concepts:
            return
        try:
            vectors = await self._embeddings.generate_embeddings_batch(result.concepts)
            embeddings = [
                ConceptEmbedding(
                    concept=c,
                    owner_type="trace",
                    owner_id=trace.id,
                    embedding=v,
                )
                for c, v in zip(result.concepts, vectors, strict=True)
            ]
            await self._storage.concept_embeddings.store_concept_embeddings(embeddings)
        except (StorageError, OSError):
            logger.warning("Failed to embed concepts for trace %s", trace.id[:8])

    async def _embed_fact_tags(self, fact: PersonaFact) -> None:
        """Embed each context tag individually for concept attention."""
        if not fact.context_tags:
            return
        try:
            vectors = await self._embeddings.generate_embeddings_batch(
                fact.context_tags
            )
            embeddings = [
                ConceptEmbedding(
                    concept=tag,
                    owner_type="fact",
                    owner_id=fact.id,
                    embedding=v,
                )
                for tag, v in zip(fact.context_tags, vectors, strict=True)
            ]
            await self._storage.concept_embeddings.store_concept_embeddings(embeddings)
        except (StorageError, OSError):
            logger.warning("Failed to embed tags for fact %s", fact.id[:8])

    async def _create_entity_associations(
        self, trace: MemoryTrace, result: ExtractionResult
    ) -> None:
        """Create associations to traces sharing entities."""
        weight = float(self._config.get("associations.entity_weight", 0.7))
        max_links = int(self._config.get("associations.max_links_per_entity", 5))
        for entity in result.entities:
            await self._link_traces_by_key(
                trace, entity.name, "entity", weight, max_links
            )

    async def _create_concept_associations(
        self, trace: MemoryTrace, result: ExtractionResult
    ) -> None:
        """Create associations to traces sharing concepts."""
        weight = float(self._config.get("associations.concept_weight", 0.4))
        max_links = int(self._config.get("associations.max_links_per_entity", 5))
        for concept in result.concepts:
            await self._link_traces_by_key(trace, concept, "concept", weight, max_links)

    async def _link_traces_by_key(
        self,
        trace: MemoryTrace,
        key: str,
        assoc_type: str,
        weight: float,
        max_links: int,
    ) -> None:
        """Create association for a shared key.

        DB pair_key column handles bidirectionality.
        """
        try:
            if assoc_type == "entity":
                existing_ids = await self._storage.entities.get_traces_by_entity(
                    key, limit=max_links
                )
            else:
                existing_ids = await self._storage.entities.get_traces_by_concept(
                    key, limit=max_links
                )
        except (StorageError, OSError):
            logger.warning("Failed to find traces for %s=%s", assoc_type, key)
            return

        for other_id in existing_ids:
            if other_id == trace.id:
                continue
            association = Association(
                source_trace_id=trace.id,
                target_trace_id=other_id,
                association_type=assoc_type,
                weight=weight,
                forward_strength=weight,
                backward_strength=weight,
            )
            try:
                await self._storage.associations.store_association(association)
            except (StorageError, OSError):
                logger.warning(
                    "Failed to create %s association %s <-> %s",
                    assoc_type,
                    trace.id[:8],
                    other_id[:8],
                )

    async def _extract_persona_facts(
        self, trace: MemoryTrace, result: ExtractionResult
    ) -> None:
        """Extract persona facts from semantic-type traces with promotion gate."""
        if result.fact_type != "semantic":
            return
        threshold = float(self._config.get("persona.confidence_threshold", 0.6))
        for rel in result.relations:
            if rel.confidence < threshold:
                continue
            raw_cat = rel.category if hasattr(rel, "category") else "general"
            category: FactCategory = cast(
                FactCategory,
                raw_cat if raw_cat in _VALID_CATEGORIES else "general",
            )
            status: FactStatus = (
                "promoted"
                if _should_fast_track(category, rel.confidence)
                else "candidate"
            )
            content = rel.context or f"{rel.source} {rel.relation} {rel.target}"
            try:
                fact_embedding = await self._embeddings.generate_embedding(content)
            except (ValueError, RuntimeError, OSError):
                fact_embedding = None
            fact = PersonaFact(
                subject=rel.source,
                predicate=_canonicalize_predicate(rel.relation),
                object=rel.target,
                category=category,
                content=content,
                source_trace_id=trace.id,
                confidence=rel.confidence,
                status=status,
                scope=_category_to_scope(category),
                context_tags=rel.context_tags,
                embedding=fact_embedding,
            )
            await self._store_or_promote_fact(fact)
            await self._embed_fact_tags(fact)

    async def _store_or_promote_fact(self, fact: PersonaFact) -> None:
        """Store a new fact or promote an existing candidate."""
        try:
            existing = await self._storage.facts.get_persona_facts(subject=fact.subject)
            duplicate = _find_exact_duplicate(existing, fact)
            if duplicate:
                new_count = await self._storage.facts.increment_mention_count(
                    duplicate.id
                )
                if new_count >= 2 and duplicate.status == "candidate":
                    await self._storage.facts.update_fact_status(
                        duplicate.id, "promoted"
                    )
                    logger.debug(
                        "Promoted fact: %s %s %s (mentions=%d)",
                        duplicate.subject,
                        duplicate.predicate,
                        duplicate.object,
                        new_count,
                    )
                return
            contradicting = _find_contradicting_fact(existing, fact)
            if contradicting:
                await self._check_predicate_similarity(contradicting, fact)
            else:
                await self._storage.facts.store_persona_fact(fact)
        except (StorageError, OSError):
            logger.warning(
                "Failed to store/promote fact: %s %s %s",
                fact.subject,
                fact.predicate,
                fact.object,
            )

    async def _check_predicate_similarity(
        self, old_fact: PersonaFact, new_fact: PersonaFact
    ) -> None:
        """Supersede if predicates match exactly or by embedding similarity."""
        if old_fact.predicate == new_fact.predicate:
            await self._storage.facts.supersede_persona_fact(old_fact.id, new_fact)
            return
        threshold = float(
            self._config.get("persona.predicate_similarity_threshold", 0.8)
        )
        emb_a = await self._embeddings.generate_embedding(old_fact.predicate)
        emb_b = await self._embeddings.generate_embedding(new_fact.predicate)
        if _cosine_similarity(emb_a, emb_b) >= threshold:
            await self._storage.facts.supersede_persona_fact(old_fact.id, new_fact)
        else:
            await self._storage.facts.store_persona_fact(new_fact)

    async def _extract_entity_relations(
        self, trace: MemoryTrace, result: ExtractionResult
    ) -> None:
        """Create entity-to-entity edges from extracted relations."""
        for rel in result.relations:
            if not rel.source or not rel.target:
                continue
            entity_rel = EntityRelation(
                source_entity=rel.source,
                relation=rel.relation,
                target_entity=rel.target,
                context=f"{rel.source} {rel.relation} {rel.target}",
                weight=rel.confidence,
                source_trace_id=trace.id,
            )
            try:
                await self._storage.entity_relations.store_relation(entity_rel)
            except (StorageError, OSError):
                logger.warning(
                    "Failed to store entity relation: %s -> %s",
                    rel.source,
                    rel.target,
                )

    # -- Private: think_about helpers --

    def _search_working_memory(
        self, query_embedding: list[float]
    ) -> list[tuple[MemoryTrace, float]]:
        """Find relevant traces in working memory by embedding similarity."""
        wm_threshold = float(self._config.get("retrieval.wm_similarity_threshold", 0.3))
        results: list[tuple[MemoryTrace, float]] = []
        for trace in self._buffer.get_active():
            if trace.embedding is not None:
                sim = _cosine_similarity(query_embedding, trace.embedding)
                if sim > wm_threshold:
                    results.append((trace, sim))
        results.sort(key=lambda x: x[1], reverse=True)
        return results

    async def _spread_from_candidates(
        self, candidates: list[MemoryTrace]
    ) -> list[tuple[MemoryTrace, float]]:
        """Spread activation from top candidates."""
        activated: list[tuple[MemoryTrace, float]] = []
        seen_ids: set[str] = set()
        for candidate in candidates:
            spreads = await self._storage.vectors.spread_activation(candidate.id)
            for trace, level in spreads:
                if trace.id not in seen_ids:
                    activated.append((trace, level))
                    seen_ids.add(trace.id)
        return activated

    async def _merge_candidates(
        self,
        query_embedding: list[float],
        wm: list[tuple[MemoryTrace, float]],
        storage: list[tuple[MemoryTrace, float]],
        activated: list[tuple[MemoryTrace, float]],
        entity_matches: list[tuple[str, float]] | None = None,
    ) -> list[tuple[MemoryTrace, float]]:
        """Merge candidates from all sources using weighted max-score fusion.

        Takes the maximum similarity score a trace receives from any source,
        then adds small bonuses for entity matches and spreading activation.
        Entity bonus is gated by concept similarity.
        """
        entity_bonus = float(self._config.get("retrieval.entity_match_bonus", 0.1))
        spread_bonus = float(self._config.get("activation.spread_weight_factor", 0.1))
        significance_weight = float(
            self._config.get("retrieval.significance_weight", 0.15)
        )
        valence_weight = float(self._config.get("retrieval.valence_weight", 0.05))

        traces_by_id: dict[str, MemoryTrace] = {}
        scores_by_id: dict[str, float] = {}

        for trace, sim in wm:
            traces_by_id[trace.id] = trace
            scores_by_id[trace.id] = max(scores_by_id.get(trace.id, 0.0), sim)

        for trace, sim in storage:
            traces_by_id[trace.id] = trace
            scores_by_id[trace.id] = max(scores_by_id.get(trace.id, 0.0), sim)

        activation_levels = self._collect_activation_levels(
            activated, traces_by_id, scores_by_id, spread_bonus
        )

        entity_sims: dict[str, float] = {}
        if entity_matches:
            for tid, sim in entity_matches:
                entity_sims[tid] = sim

        concept_weight = float(
            self._config.get("retrieval.concept_attention_weight", 0.70)
        )
        concept_sims: dict[str, float] = {}
        if concept_weight > 0.0 and traces_by_id:
            try:
                concept_sims = (
                    await self._storage.concept_embeddings.get_max_sim_per_owner(
                        query_embedding,
                        owner_type="trace",
                        owner_ids=list(traces_by_id.keys()),
                    )
                )
            except (StorageError, OSError):
                logger.warning("Concept attention lookup failed for traces")

        return self._compute_fused_scores(
            traces_by_id,
            scores_by_id,
            activation_levels,
            entity_sims,
            spread_bonus,
            entity_bonus,
            significance_weight,
            valence_weight,
            concept_sims=concept_sims,
            concept_weight=concept_weight,
        )

    @staticmethod
    def _collect_activation_levels(
        activated: list[tuple[MemoryTrace, float]],
        traces_by_id: dict[str, MemoryTrace],
        scores_by_id: dict[str, float],
        spread_bonus: float,
    ) -> dict[str, float]:
        """Collect activation levels, using fraction as base for unseen traces."""
        activation_levels: dict[str, float] = {}
        for trace, level in activated:
            traces_by_id[trace.id] = trace
            activation_levels[trace.id] = max(
                activation_levels.get(trace.id, 0.0), level
            )
            # Traces found only via activation get a fraction as base score
            if trace.id not in scores_by_id:
                scores_by_id[trace.id] = level * spread_bonus
        return activation_levels

    @staticmethod
    def _compute_fused_scores(
        traces_by_id: dict[str, MemoryTrace],
        scores_by_id: dict[str, float],
        activation_levels: dict[str, float],
        entity_sims: dict[str, float],
        spread_bonus: float,
        entity_bonus: float,
        significance_weight: float = 0.15,
        valence_weight: float = 0.05,
        *,
        concept_sims: dict[str, float] | None = None,
        concept_weight: float = 0.0,
    ) -> list[tuple[MemoryTrace, float]]:
        """Compute fused scores with concept-primary blend, clamped to [0, 1].

        When concept embeddings exist for a trace, base cosine similarity
        is blended with concept max-sim (weight*concept + (1-weight)*base).
        This mirrors attention: concept scores modulate rather than add.
        Entity bonus remains gated by concept similarity.
        """
        if concept_sims is None:
            concept_sims = {}
        result: list[tuple[MemoryTrace, float]] = []
        for trace_id, base in scores_by_id.items():
            trace = traces_by_id.get(trace_id)
            if trace is None:
                continue
            # Concept-primary blend: concept attention modulates base
            # similarity instead of adding to it. When concept embeddings
            # exist, effective_sim = weight*concept + (1-weight)*base.
            concept_sim = concept_sims.get(trace_id, 0.0)
            if concept_sim > 0.0 and concept_weight > 0.0:
                effective_sim = (
                    concept_weight * concept_sim + (1.0 - concept_weight) * base
                )
            else:
                effective_sim = base
            significance_boost = trace.significance * significance_weight
            valence_boost = abs(trace.emotional_valence) * valence_weight
            score = effective_sim + significance_boost + valence_boost
            if trace_id in activation_levels:
                score += activation_levels[trace_id] * spread_bonus
            if trace_id in entity_sims:
                score += (
                    entity_sims[trace_id]
                    * entity_bonus
                    * trace.significance
                    * concept_sim
                )
            score = min(score, 1.0)
            result.append((trace, score))
        result.sort(key=lambda x: x[1], reverse=True)
        return result

    def _select_within_budget(
        self,
        candidates: list[tuple[MemoryTrace, float]],
        token_budget: int,
    ) -> list[tuple[MemoryTrace, float]]:
        """Select top candidates that fit within token budget."""
        selected: list[tuple[MemoryTrace, float]] = []
        tokens_used = 0
        max_retrievals = int(self._config.get("retrieval.max_retrievals", 3))

        for trace, score in candidates:
            if len(selected) >= max_retrievals:
                break
            tokens = _estimate_tokens(trace.content or "")
            if tokens_used + tokens > token_budget:
                break
            selected.append((trace, score))
            tokens_used += tokens

        return selected

    async def _boost_and_build_thoughts(
        self, selected: list[tuple[MemoryTrace, float]]
    ) -> list[Thought]:
        """Apply retrieval boost to selected traces and build Thoughts."""
        thoughts: list[Thought] = []
        for trace, relevance in selected:
            in_wm = self._buffer.find(lambda t, _id=trace.id: t.id == _id) is not None
            boosted = apply_retrieval_boost(trace, from_working_memory=in_wm)
            await self._storage.traces.update_trace_strength(trace.id, boosted.strength)
            await self._storage.traces.mark_retrieved(trace.id)

            content = trace.content or str(trace.pattern)
            thoughts.append(
                Thought(
                    trace=trace,
                    relevance=relevance,
                    token_count=_estimate_tokens(content),
                    reconstructed=trace.content is None,
                    reconstruction=content,
                )
            )
        return thoughts

    async def _boost_unselected(
        self,
        all_candidates: list[tuple[MemoryTrace, float]],
        thoughts: list[Thought],
    ) -> None:
        """Apply small activation boost to candidates not selected."""
        selected_ids = {t.trace.id for t in thoughts}
        for trace, _ in all_candidates:
            if trace.id not in selected_ids:
                boosted = apply_activation_boost(trace)
                await self._storage.traces.update_trace_strength(
                    trace.id, boosted.strength
                )
                await self._storage.traces.mark_activated(trace.id)

    async def _match_query_entities(self, query: str) -> list[tuple[str, float]]:
        """Find trace IDs for entities mentioned in the query.

        Extracts candidate entity names from query text, then uses
        trigram similarity matching in storage.
        """
        names = self._extract_entity_names_from_query(query)
        if not names:
            return []
        try:
            return await self._storage.entities.match_entities(names)
        except (StorageError, OSError):
            logger.warning("Entity matching failed for query")
            return []

    async def _find_relevant_persona_facts(
        self,
        query: str,
        query_embedding: list[float],
        *,
        user_id: str | None = None,
    ) -> tuple[list[PersonaFact], dict[str, float]]:
        """Find persona facts via semantic search + entity supplement.

        Scoring uses concept attention (ColBERT-style max-sim on LLM-extracted
        tags, 0.7 weight) blended with bi-encoder similarity (0.3 weight).
        Cross-encoder re-ranking was removed: MS MARCO models cannot perform
        the causal inference this retrieval requires, and the LLM already
        reasons about relevance at extraction time via query-predictive tags.

        Primary: embed(query) vs embed(fact.context) via pgvector.
        Supplement: entity name match adds structurally-related facts.
        """
        try:
            entity_names = self._extract_entity_names_from_query(query)
            max_facts = int(self._config.get("persona.max_facts_per_query", 5))
            sim_threshold = float(self._config.get("persona.similarity_threshold", 0.3))

            # Primary: semantic search on fact content embeddings
            semantic_results = await self._storage.facts.search_facts_semantic(
                query_embedding,
                limit=max_facts * 2,
                user_id=user_id,
            )
            semantic_facts = [f for f, sim in semantic_results if sim >= sim_threshold]
            semantic_scores: dict[str, float] = {
                f.id: sim for f, sim in semantic_results if sim >= sim_threshold
            }

            # Supplement: entity name match (adds facts for named entities)
            entity_facts = await self._storage.facts.get_persona_facts_by_entities(
                entity_names,
                user_id=user_id,
            )

            all_facts = self._deduplicate_facts(semantic_facts + entity_facts)

            # Score entity-only facts via embedding
            for fact in all_facts:
                if fact.id not in semantic_scores and fact.embedding:
                    sim = _cosine_similarity(query_embedding, fact.embedding)
                    if sim >= sim_threshold:
                        semantic_scores[fact.id] = sim

            # Filter: only keep facts that passed the similarity threshold
            all_facts = [f for f in all_facts if f.id in semantic_scores]

            # Concept attention: blend concept and bi-encoder scores
            # Concept is primary (0.7) because LLM wrote tags as query predictions.
            # Bi-encoder is fallback (0.3) for general semantic similarity.
            try:
                fact_ids = [f.id for f in all_facts]
                if fact_ids:
                    fact_concept_sims = (
                        await self._storage.concept_embeddings.get_max_sim_per_owner(
                            query_embedding,
                            owner_type="fact",
                            owner_ids=fact_ids,
                        )
                    )
                    for fid, csim in fact_concept_sims.items():
                        bi_sim = semantic_scores.get(fid, 0.0)
                        semantic_scores[fid] = 0.7 * csim + 0.3 * bi_sim
            except (StorageError, OSError):
                logger.warning("Fact concept attention lookup failed")

            return (
                self._rank_and_limit_facts(all_facts, max_facts, semantic_scores),
                semantic_scores,
            )
        except (StorageError, OSError):
            logger.warning("Persona fact lookup failed")
            return [], {}

    @staticmethod
    def _deduplicate_facts(facts: list[PersonaFact]) -> list[PersonaFact]:
        """Remove duplicate facts by ID, preserving order."""
        seen: set[str] = set()
        unique: list[PersonaFact] = []
        for fact in facts:
            if fact.id not in seen:
                seen.add(fact.id)
                unique.append(fact)
        return unique

    @staticmethod
    def _rank_and_limit_facts(
        facts: list[PersonaFact],
        limit: int,
        semantic_scores: dict[str, float] | None = None,
    ) -> list[PersonaFact]:
        """Rank facts by semantic similarity to the query, then limit."""

        def sort_key(f: PersonaFact) -> tuple[float, float]:
            sim = (semantic_scores or {}).get(f.id, 0.0)
            return (-sim, -f.confidence)

        return sorted(facts, key=sort_key)[:limit]


    def _extract_entity_names_from_query(self, query: str) -> list[str]:
        """Extract potential entity names from query text.

        Capitalized words not in common stop words are treated as
        potential entity names (proper nouns).
        """
        words = query.split()
        entities: list[str] = []
        for word in words:
            cleaned = word.strip(".,!?;:'\"")
            if (
                cleaned
                and cleaned[0].isupper()
                and cleaned.lower() not in _QUERY_STOP_WORDS
            ):
                entities.append(cleaned)
        return entities

    def _persona_facts_to_thoughts(
        self,
        facts: list[PersonaFact],
        semantic_scores: dict[str, float] | None = None,
    ) -> list[Thought]:
        """Convert persona facts to Thought objects using ranking strategy."""
        strategy = str(self._config.get("persona.ranking_strategy", "hybrid"))
        threshold = float(self._config.get("persona.confidence_threshold", 0.6))
        thoughts: list[Thought] = []
        for fact in facts:
            relevance = self._compute_fact_relevance(
                fact,
                (semantic_scores or {}).get(fact.id, 0.0),
            )
            triple = f"{fact.subject} {fact.predicate} {fact.object}"
            if fact.content and fact.content != triple:
                content = f"[IMPORTANT CONTEXT] {triple} -- {fact.content}"
            else:
                content = f"[IMPORTANT CONTEXT] {triple}"
            trace = MemoryTrace(
                content=content,
                pattern={"persona_fact": True, "category": fact.category},
                strength=1.0,
                significance=1.0,
            )
            if strategy == "pinned":
                pinned = True
            elif strategy == "relevance":
                pinned = False
            else:  # hybrid
                pinned = relevance >= threshold
            thoughts.append(
                Thought(
                    trace=trace,
                    relevance=relevance,
                    token_count=_estimate_tokens(content),
                    reconstructed=True,
                    reconstruction=content,
                    pinned=pinned,
                )
            )
        return thoughts

    @staticmethod
    def _compute_fact_relevance(
        fact: PersonaFact,
        semantic_similarity: float = 0.0,
    ) -> float:
        """Compute fact relevance blending confidence with semantic similarity."""
        if semantic_similarity > 0.0:
            return 0.3 * fact.confidence + 0.7 * semantic_similarity
        return fact.confidence

    # -- Private: consolidation helpers --

    async def _consolidate_one(
        self,
        trace: MemoryTrace,
        decayed: MemoryTrace,
        threshold: float,
        grace_hours: float,
    ) -> str:
        """Consolidate a single trace.

        Returns 'consolidated', 'forgotten', or 'pending'.
        """
        if decayed.strength >= threshold:
            await self._storage.traces.update_trace_strength(trace.id, decayed.strength)
            await self._storage.traces.mark_consolidated(trace.id)
            return "consolidated"

        now = memory_timestamp_for_comparison(now_utc())
        created = memory_timestamp_for_comparison(trace.created_at)
        age_hours = (now - created).total_seconds() / 3600.0

        if age_hours >= grace_hours:
            await self._storage.traces.delete_trace(trace.id)
            return "forgotten"

        await self._storage.traces.update_trace_strength(trace.id, decayed.strength)
        return "pending"
