"""
cortex/identity/cornerstone_injector.py — Pure injection layer for cornerstones.

Handles prompt formatting, category-block ordering, and token budget calculation
for CornerstoneMemory records. This module has no storage or CRUD concerns.

Extracted from cornerstone.py as part of issue #711 (v1.3.0 architecture).

Design principles:
    - Pure injection: no DB access, no LLM calls, no CRUD
    - Stateless: takes cornerstones as input, returns formatted strings
    - Category-block ordering per Liu et al. 2023 "Lost in the Middle"
    - Independently testable: only requires CornerstoneMemory dataclasses

Used by:
    - CornerstoneGuardian (cortex.identity.cornerstone) — facade wiring
    - CortexPlugin (M13) — wake() calls format_injection_block()
    - TokenBudgetManager (M9) — reserves cornerstone_token_budget
"""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING, Dict, List, Optional

if TYPE_CHECKING:
    from cortex.config import CortexConfig

from cortex.identity.cornerstone import CornerstoneMemory

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Default category block ordering
# (positional-bias-aware per Liu et al. 2023 "Lost in the Middle")
# ---------------------------------------------------------------------------

#: Default block ordering. Categories align with VALID_CORNERSTONE_CATEGORIES
#: in cortex.validation. Override via config.cornerstone_block_ordering.
DEFAULT_BLOCK_ORDER: List[str] = [
    "memory",      # autobiographical/episodic identity (#1–#4)
    "anchor",      # grounding facts — soul, relationship, continuity (#5–#7)
    "meta_story",  # narrative identity and origin story (#9–#10)
    "guardrail",   # behavioural constraints and refusals (#8, #11–#12)
    "other",       # uncategorized / custom
]


class CornerstoneInjector:
    """
    Pure injection layer for cornerstones.

    Formats CornerstoneMemory records into prompt injection blocks, respecting
    category-based ordering for optimal LLM positional attention.

    This class is stateless with respect to storage — it never reads from or
    writes to the database. Pass pre-loaded cornerstones to format_block().

    Example usage::

        injector = CornerstoneInjector(config)
        cornerstones = await store.get_all(owner_id="eva")
        block = injector.format_block(cornerstones)

    """

    def __init__(self, config: "CortexConfig") -> None:
        self.config = config
        self._block_order: List[str] = (
            getattr(config, "cornerstone_block_ordering", None)
            or DEFAULT_BLOCK_ORDER
        )

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def format_block(self, cornerstones: List[CornerstoneMemory]) -> str:
        """
        Format a list of cornerstones into a prompt injection block.

        When cornerstones have ``category`` fields set, they are grouped into
        named blocks in the order defined by ``config.cornerstone_block_ordering``
        (defaults to :data:`DEFAULT_BLOCK_ORDER`). This exploits LLM positional
        bias deliberately — see *Lost in the Middle* (Liu et al. 2023).

        Backward-compatible: if ALL cornerstones lack a ``category`` value, the
        output falls back to the flat-list format used prior to v1.3.0::

            [CORNERSTONES — WHO I AM]
            Label: content
            ...
            [END CORNERSTONES — N tokens]

        When category blocks are active the format is::

            [CORNERSTONES — WHO I AM]

            ## Memory
            Label: content
            ...

            ## Guardrail
            Label: content
            ...

            [END CORNERSTONES — N tokens]

        Args:
            cornerstones: Pre-loaded, pre-sorted list of CornerstoneMemory records.
                          Caller is responsible for filtering to active records.

        Returns:
            Formatted injection block string, or ``""`` when list is empty.
        """
        if not cornerstones:
            return ""

        total_tokens = sum(c.token_count for c in cornerstones)

        # Backward-compat: if no cornerstone has a non-None category, use flat list
        # (pre-v1.3.0 / pre-migration rows).
        has_any_category = any(
            getattr(cs, "category", None) is not None
            for cs in cornerstones
        )
        if not has_any_category:
            lines = ["[CORNERSTONES — WHO I AM]"]
            for cs in cornerstones:
                lines.append(f"{cs.label}: {cs.content}")
            lines.append(f"[END CORNERSTONES — {total_tokens} tokens]")
            return "\n".join(lines)

        # Category-block mode
        groups = self._group_by_category(cornerstones)
        ordered_cats = self._ordered_categories(groups)

        lines = ["[CORNERSTONES — WHO I AM]"]
        for cat in ordered_cats:
            lines.append("")
            lines.append(f"## {cat.capitalize()}")
            for cs in groups[cat]:
                lines.append(f"{cs.label}: {cs.content}")

        lines.append("")
        lines.append(f"[END CORNERSTONES — {total_tokens} tokens]")
        return "\n".join(lines)

    def total_tokens(self, cornerstones: List[CornerstoneMemory]) -> int:
        """
        Sum of token counts across the given cornerstones.

        Args:
            cornerstones: List of CornerstoneMemory records.

        Returns:
            Total token count (int).
        """
        return sum(c.token_count for c in cornerstones)

    def is_within_budget(
        self,
        cornerstones: List[CornerstoneMemory],
        token_budget: int,
    ) -> bool:
        """
        Check whether all cornerstones fit within the given token budget.

        Args:
            cornerstones: Cornerstones to check.
            token_budget: Available token budget.

        Returns:
            True if total tokens <= budget.
        """
        return self.total_tokens(cornerstones) <= token_budget

    def filter_to_budget(
        self,
        cornerstones: List[CornerstoneMemory],
        token_budget: int,
    ) -> List[CornerstoneMemory]:
        """
        Return a subset of cornerstones that fit within the token budget.

        Preserves injection_order. Includes as many as possible before the
        budget is exhausted. Always-loaded cornerstones (lower injection_order)
        take priority.

        Args:
            cornerstones: Pre-sorted list of cornerstones (by injection_order).
            token_budget: Maximum tokens available for injection.

        Returns:
            Subset of cornerstones fitting within the budget.
        """
        result: List[CornerstoneMemory] = []
        used = 0
        for cs in cornerstones:
            if used + cs.token_count <= token_budget:
                result.append(cs)
                used += cs.token_count
            else:
                logger.debug(
                    "Token budget exceeded: skipping cornerstone %r (%d tokens, "
                    "budget remaining=%d)",
                    cs.label, cs.token_count, token_budget - used,
                )
        return result

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _group_by_category(
        self, cornerstones: List[CornerstoneMemory]
    ) -> Dict[str, List[CornerstoneMemory]]:
        """Group cornerstones by category; uncategorized → 'other'."""
        groups: Dict[str, List[CornerstoneMemory]] = {}
        for cs in cornerstones:
            cat = (getattr(cs, "category", None) or "other").lower().strip()
            groups.setdefault(cat, []).append(cs)
        return groups

    def _ordered_categories(
        self, groups: Dict[str, List[CornerstoneMemory]]
    ) -> List[str]:
        """
        Build the ordered category list for block rendering.

        Applies block_order, then appends any unlisted categories, with 'other'
        always last.
        """
        ordered: List[str] = []
        for cat in self._block_order:
            if cat in groups:
                ordered.append(cat)
        # Append any categories present but not in block_order (before "other")
        for cat in groups:
            if cat not in ordered and cat != "other":
                ordered.append(cat)
        if "other" in groups and "other" not in ordered:
            ordered.append("other")
        return ordered
