"""
cortex/core/retrieval_types.py — Data transfer objects and constants for the retrieval pipeline.

Extracted from cortex/core/retrieval.py (#947 — SRP split).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple


# ---------------------------------------------------------------------------
# Sensitivity level ordering (lower index = less sensitive)
# ---------------------------------------------------------------------------

SENSITIVITY_ORDER = ["normal", "private", "secret", "restricted"]
_SENSITIVITY_ALIASES = {"sensitive": "secret", "personal": "private"}


# ---------------------------------------------------------------------------
# Public data-transfer objects
# ---------------------------------------------------------------------------

@dataclass
class RetrievalConstraints:
    """Configurable constraints for a retrieve() call.

    All fields are optional — unset fields use defaults from CortexConfig.

    Attributes:
        memory_types: Filter by item type — 'claim', 'event', 'procedure', etc.
                      None = all types.
        time_range: ISO 8601 (start, end) tuple for temporal filtering.
        entity_ids: Only return items related to these entity UUIDs.
        sensitivity_max: Maximum sensitivity level to return. Items above this
                         level are excluded. Valid: normal < private < secret <
                         restricted. Default: 'secret' (excludes 'restricted').
        token_budget: Override the global retrieval_token_budget for this call.
        top_k: Maximum number of items to return after fusion. Default: 20.
        before_date: ISO date string (YYYY-MM-DD). Exclude claims whose
                     event_date is after this date (inclusive upper bound).
                     Claims without event_date are never excluded.
        after_date: ISO date string (YYYY-MM-DD). Exclude claims whose
                    event_date is before this date (inclusive lower bound).
                    Claims without event_date are never excluded.
    """
    memory_types: Optional[List[str]] = None        # filter by item_type
    time_range: Optional[Tuple[str, str]] = None    # (iso_start, iso_end)
    entity_ids: Optional[List[str]] = None
    entity_id: Optional[str] = None                # filter by subject_entity_id (#923)
    entity_type: Optional[str] = None              # filter by entity_type (#923)
    exclude_episodes: bool = True                   # exclude episode events by default (#919)
    sensitivity_max: str = "secret"                 # excludes 'restricted' by default
    token_budget: Optional[int] = None              # override global budget
    top_k: int = 20
    before_date: Optional[str] = None              # ISO date upper bound for event_date
    after_date: Optional[str] = None               # ISO date lower bound for event_date
    include_expired: bool = False                   # include expired claims (valid_to set) (#925)
    exclude_source_session: Optional[str] = None   # exclude claims from this source session (#1220)


@dataclass
class RetrievedItem:
    """A single memory item returned from retrieval.

    Attributes:
        item_type: Memory layer type — 'claim', 'event', 'procedure', 'brain_graph_node'.
        item_id: UUID of the source record.
        content: Text content to inject into the prompt.
        score: Hybrid/RRF score — higher = more relevant.
        provenance: Optional trace info (retrieval explanation, Phase 2).
        metadata: Optional dict of classification metadata (memory_class, salience, status)
                  populated for claim items so the context compiler can use them.
    """
    item_type: str          # 'claim', 'event', 'procedure', 'brain_graph_node'
    item_id: str
    content: str
    score: float
    provenance: Optional[str] = None
    metadata: Optional[Dict[str, Any]] = None
    source_session_id: Optional[str] = None   # OpenClaw session that created this memory (#954)
    created_at: Optional[str] = None          # ISO 8601 creation timestamp (#954)


@dataclass
class RetrievalResult:
    """Complete result from a retrieve() call.

    Returned by: RetrievalPipeline.retrieve().
    Consumed by: CortexPlugin.retrieve() → returned to caller / MCP / HTTP.

    Attributes:
        items: Ranked list of retrieved memory items.
        retrieval_id: UUID for this retrieval (used in feedback/logging).
        mode: Which retrieval mode was used — 'lightweight', 'agentic', etc.
        latency_ms: End-to-end retrieval time in milliseconds.
        tokens_used: Total tokens consumed by the returned items.
    """
    items: List[RetrievedItem]
    retrieval_id: str
    mode: str
    latency_ms: int
    tokens_used: int
    retrieval_info: Optional[Dict[str, Any]] = None  # Channel health info (#P8-4)


@dataclass
class ContextBlock:
    """Assembled prompt-injection block produced by assemble_context().

    This is the final formatted output ready for system prompt injection.

    Attributes:
        parts: List of (type, text, tokens) tuples in injection order.
        total_tokens: Sum of all part token counts.
        budget_remaining: How many tokens are still available after this block.
        retrieval_id: UUID for tracing/feedback.
    """
    parts: List[Tuple[str, str, int]]  # (type, text, tokens)
    total_tokens: int
    budget_remaining: int
    retrieval_id: str


# ---------------------------------------------------------------------------
# Internal result types that mirror StorageAdapter returns
# (kept here so retrieval modules don't import from adapter directly)
# ---------------------------------------------------------------------------

@dataclass
class _VectorResult:
    item_type: str
    item_id: str
    content: str
    score: float


@dataclass
class _FTSResult:
    item_type: str
    item_id: str
    content: str
    rank: float  # BM25 rank (lower is better in SQLite FTS5 bm25())


# ---------------------------------------------------------------------------
# Retrieval trace (#1347)
# ---------------------------------------------------------------------------

@dataclass
class RetrievalTrace:
    """Structured trace for a single retrieve() call.

    Emitted as a JSON log entry to the ``cortex.retrieval.trace`` logger when
    ``retrieval_trace_enabled`` is True in CortexConfig. Designed for
    Fly.io log aggregation and retrieval quality debugging.

    Fields intentionally flat (no nesting) so they can be parsed with simple
    JSON log processors. All timing values are in milliseconds.
    """
    # Identity
    retrieval_id: str
    owner_id: str
    query: str
    mode: str
    timestamp_utc: str                      # ISO 8601

    # Stage: BM25
    bm25_candidate_count: int = 0
    bm25_top_score: Optional[float] = None  # normalised, after _normalise_bm25

    # Stage: vector search
    vector_candidate_count: int = 0
    vector_top_score: Optional[float] = None

    # Stage: RRF fusion
    fused_candidate_count: int = 0
    fused_top_score: Optional[float] = None

    # Stage: rerank (optional — only populated when rerank is enabled)
    rerank_input_count: int = 0
    rerank_output_count: int = 0
    rerank_top_score: Optional[float] = None
    rerank_latency_ms: int = 0

    # Final selection
    final_count: int = 0
    final_top_score: Optional[float] = None
    final_min_score: Optional[float] = None
    final_item_ids: List[str] = field(default_factory=list)

    # Category distribution across final items (category → count)
    category_distribution: Dict[str, int] = field(default_factory=dict)

    # Timing
    total_latency_ms: int = 0
    bm25_latency_ms: int = 0
    vector_latency_ms: int = 0

    def to_dict(self) -> Dict[str, Any]:
        """Return a plain dict suitable for JSON serialisation."""
        return {
            "event": "retrieval_trace",
            "retrieval_id": self.retrieval_id,
            "owner_id": self.owner_id,
            "query": self.query,
            "mode": self.mode,
            "timestamp_utc": self.timestamp_utc,
            "stages": {
                "bm25": {
                    "candidate_count": self.bm25_candidate_count,
                    "top_score": self.bm25_top_score,
                    "latency_ms": self.bm25_latency_ms,
                },
                "vector": {
                    "candidate_count": self.vector_candidate_count,
                    "top_score": self.vector_top_score,
                    "latency_ms": self.vector_latency_ms,
                },
                "fused": {
                    "candidate_count": self.fused_candidate_count,
                    "top_score": self.fused_top_score,
                },
                "rerank": {
                    "input_count": self.rerank_input_count,
                    "output_count": self.rerank_output_count,
                    "top_score": self.rerank_top_score,
                    "latency_ms": self.rerank_latency_ms,
                },
            },
            "final": {
                "count": self.final_count,
                "top_score": self.final_top_score,
                "min_score": self.final_min_score,
                "item_ids": self.final_item_ids,
                "category_distribution": self.category_distribution,
            },
            "total_latency_ms": self.total_latency_ms,
        }
