"""
cortex/identity/cornerstone_store.py — Pure storage layer for cornerstones.

Handles all CRUD operations, validation, token counting, and hash computation
for CornerstoneMemory records. This module has no prompt-formatting concerns.

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

Design principles:
    - Pure storage: no prompt formatting, no LLM calls
    - Validates input before persisting
    - Hashes the STORED content (not raw input — storage may normalise)
    - Independently testable: only depends on StorageAdapter + CortexConfig

Used by:
    - CornerstoneGuardian (cortex.identity.cornerstone) — facade wiring
    - CornerstoneInjector (cortex.identity.cornerstone_injector) — reads for injection
    - DriftCalculator (M11) — reads cornerstones, calls update_drift()
"""

from __future__ import annotations

import hashlib
import logging
from datetime import datetime, timezone
from typing import TYPE_CHECKING, List, Optional

try:
    import tiktoken
    _TIKTOKEN_AVAILABLE = True
except ImportError:  # pragma: no cover
    _TIKTOKEN_AVAILABLE = False

if TYPE_CHECKING:
    from cortex.storage.adapter import StorageAdapter
    from cortex.config import CortexConfig

from cortex.identity.cornerstone import (
    CornerstoneMemory,
    DriftReport,
    DEFAULT_MAX_CORNERSTONES,
    DEFAULT_TIKTOKEN_MODEL,
    _FALLBACK_CHARS_PER_TOKEN,
    _sha256,
    _count_tokens,
    _resolve_tiktoken_model,
)

logger = logging.getLogger(__name__)


class CornerstoneStore:
    """
    Pure storage layer for cornerstone memories.

    Provides CRUD operations with validation and hash management. All prompt
    formatting and injection concerns are handled by CornerstoneInjector.

    This class is independently testable with a mock StorageAdapter.
    """

    def __init__(
        self,
        storage: "StorageAdapter",
        config: "CortexConfig",
    ) -> None:
        self.storage = storage
        self.config = config
        self._tiktoken_model: str = _resolve_tiktoken_model(
            getattr(config, "embedding_model", None)
        )

    @property
    def _max_cornerstones(self) -> int:
        """Read max cornerstones from config dynamically (supports runtime changes)."""
        return getattr(
            self.config, "max_cornerstone_count",
            getattr(self.config, "max_cornerstones", DEFAULT_MAX_CORNERSTONES)
        )

    # ------------------------------------------------------------------
    # Read operations
    # ------------------------------------------------------------------

    async def get_all(self, owner_id: str = "default") -> List[CornerstoneMemory]:
        """
        Return all ACTIVE cornerstones for an owner, sorted by injection_order.

        Args:
            owner_id: Owner scope.

        Returns:
            List[CornerstoneMemory] sorted by injection_order ascending.
        """
        all_cs = await self.storage.get_cornerstones(owner_id=owner_id)
        active = [c for c in all_cs if c.status == "active"]
        return sorted(active, key=lambda c: c.injection_order)

    async def get_by_id(
        self, cs_id: str, owner_id: Optional[str] = None
    ) -> Optional[CornerstoneMemory]:
        """
        Return a cornerstone by ID regardless of status, or None.

        Args:
            cs_id:    Cornerstone ID.
            owner_id: Optional owner scope filter.

        Returns:
            CornerstoneMemory or None.
        """
        if hasattr(self.storage, "get_cornerstone_by_id"):
            return await self.storage.get_cornerstone_by_id(cs_id, owner_id=owner_id)
        # Fallback: scan all (for adapters that don't implement get_cornerstone_by_id)
        all_cs = await self.storage.get_cornerstones(owner_id=owner_id or "")
        return next((c for c in all_cs if c.id == cs_id), None)

    async def get_total_token_count(self, owner_id: str = "default") -> int:
        """
        Sum of token counts across all active cornerstones.

        Used by the token budget manager (M9) to reserve space for cornerstone
        injection at session start.

        Args:
            owner_id: Owner scope.

        Returns:
            Total token count (int).
        """
        cornerstones = await self.get_all(owner_id=owner_id)
        return sum(c.token_count for c in cornerstones)

    # ------------------------------------------------------------------
    # Write operations
    # ------------------------------------------------------------------

    async def create(
        self,
        label: str,
        content: str,
        owner_id: str = "default",
        sealed_by: str = "operator",
        category: Optional[str] = "memory",
    ) -> CornerstoneMemory:
        """
        Create and persist a new cornerstone.

        Enforces the max cornerstone cap. Hashes the STORED content (not raw
        input) to guard against storage-layer normalisation.

        Args:
            label:     Human-readable identifier.
            content:   Identity text to protect.
            owner_id:  Owner scope.
            sealed_by: Who sealed it.
            category:  Category for block-ordered injection.

        Returns:
            CornerstoneMemory with status='active'.

        Raises:
            ValueError: If the cornerstone cap would be exceeded.
        """
        existing = await self.get_all(owner_id=owner_id)
        if len(existing) >= self._max_cornerstones:
            raise ValueError(
                f"Cornerstone cap reached ({self._max_cornerstones}). "
                "Unseal an existing cornerstone before sealing a new one."
            )

        # Validate category
        from cortex.validation import validate_cornerstone_category
        category = validate_cornerstone_category(category)

        # Persist; hash what is ACTUALLY stored
        cs = await self.storage.seal_cornerstone(
            label=label,
            content=content,
            owner_id=owner_id,
            sealed_by=sealed_by,
            category=category,
        )

        # Back-fill computed fields
        stored_hash = _sha256(cs.content)
        token_count = _count_tokens(cs.content, model=self._tiktoken_model)
        injection_order = len(existing)

        cs.golden_copy = cs.content
        cs.golden_hash = stored_hash
        cs.current_hash = stored_hash
        cs.token_count = token_count
        cs.injection_order = injection_order

        logger.info(
            "Cornerstone stored: id=%s label=%r tokens=%d hash=%s",
            cs.id, cs.label, cs.token_count, cs.golden_hash[:12],
        )
        return cs

    async def retire(self, cs_id: str) -> CornerstoneMemory:
        """
        Retire (unseal) a cornerstone — removes it from the active set.

        Args:
            cs_id: Cornerstone ID to retire.

        Returns:
            Updated CornerstoneMemory with status='retired'.

        Raises:
            ValueError: If not found or already retired.
        """
        target = await self.get_by_id(cs_id)
        if target is None:
            raise ValueError(f"Cornerstone not found: {cs_id}")
        if target.status == "retired":
            raise ValueError(f"Cornerstone {cs_id!r} is already retired.")

        if hasattr(self.storage, "retire_cornerstone"):
            return await self.storage.retire_cornerstone(cs_id)
        # Fallback for adapters without retire_cornerstone
        updated = await self.storage.restore_cornerstone(cs_id)
        updated.status = "retired"
        return updated

    async def update_content(
        self,
        cs_id: str,
        new_content: str,
        reason: str,
        revised_by: str = "operator",
        revision: Optional[int] = None,
        revision_history: Optional[list] = None,
    ) -> CornerstoneMemory:
        """
        Update cornerstone content (for explicit operator revisions).

        Recomputes hash and token count. Does NOT auto-apply drift corrections.

        Args:
            cs_id:            Cornerstone ID.
            new_content:      Replacement content.
            reason:           Why the revision was made.
            revised_by:       Who authorised the change.
            revision:         New revision number (auto-increments if None).
            revision_history: Updated revision history list.

        Returns:
            Updated CornerstoneMemory.

        Raises:
            ValueError: If not found or is retired.
        """
        target = await self.get_by_id(cs_id)
        if target is None:
            raise ValueError(f"Cornerstone not found: {cs_id}")
        if target.status == "retired":
            raise ValueError(f"Cannot revise a retired cornerstone: {cs_id!r}")

        new_revision = revision if revision is not None else (target.revision + 1)
        history = revision_history if revision_history is not None else list(
            target.revision_history
        ) + [{
            "version": target.revision,
            "text": target.golden_copy,
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "reason": reason,
            "revised_by": revised_by,
        }]

        revised = await self.storage.seal_cornerstone(
            label=target.label,
            content=new_content,
            owner_id=target.owner_id,
            sealed_by=revised_by,
            cs_id=cs_id,
            revision=new_revision,
            revision_history=history,
        )

        # Back-fill computed fields
        new_hash = _sha256(new_content)
        revised.golden_copy = new_content
        revised.golden_hash = new_hash
        revised.current_hash = new_hash
        revised.token_count = _count_tokens(new_content, model=self._tiktoken_model)
        revised.revision = new_revision
        revised.revision_history = history

        logger.info(
            "Cornerstone content updated: id=%s label=%r revision=%d",
            cs_id, revised.label, revised.revision,
        )
        return revised

    async def update_drift(self, cs_id: str, drift_score: float) -> None:
        """
        Record drift score for a cornerstone.

        Args:
            cs_id:       Cornerstone ID.
            drift_score: Float drift score (0.0 = clean, 1.0 = full drift).
        """
        await self.storage.update_cornerstone_drift(cs_id, drift_score=drift_score)

    # ------------------------------------------------------------------
    # Validation helpers
    # ------------------------------------------------------------------

    def compute_hash(self, text: str) -> str:
        """Compute SHA-256 hex digest for the given text."""
        return _sha256(text)

    def count_tokens(self, text: str) -> int:
        """Count tokens for the given text using configured tokenizer."""
        return _count_tokens(text, model=self._tiktoken_model)

    def check_cap(self, existing_count: int) -> None:
        """
        Raise ValueError if adding one more cornerstone would exceed the cap.

        Args:
            existing_count: Current number of active cornerstones.

        Raises:
            ValueError: If cap would be exceeded.
        """
        if existing_count >= self._max_cornerstones:
            raise ValueError(
                f"Cornerstone cap reached ({self._max_cornerstones}). "
                "Unseal an existing cornerstone before sealing a new one."
            )
