"""Context retrieval engine for GreyMatter.

Builds optimized context payloads for agents by combining quality standards,
recent documents, and active patterns. Tracks what context gets served
for self-optimization.
"""
import json
from typing import Any

from .db import Database
from .documents import DocumentStore
from .standards import StandardsManager
from .patterns import PatternManager


class ContextEngine:
    """Builds optimized context payloads for any consumer.

    Combines quality standards, recent documents, and active patterns
    into a token-budgeted context payload. Tracks what context gets served
    for self-optimization via relevance scoring.

    Args:
        db: Database instance for recording context requests.
        doc_store: DocumentStore for retrieving recent documents.
        standards: StandardsManager for loading quality standards.
        patterns: PatternManager for loading active patterns.
    """

    def __init__(
        self,
        db: Database,
        doc_store: DocumentStore,
        standards: StandardsManager,
        patterns: PatternManager,
    ) -> None:
        self.db = db
        self.doc_store = doc_store
        self.standards = standards
        self.patterns = patterns

    def build_session_context(
        self,
        project: str | None = None,
        max_tokens: int = 12000,
    ) -> dict[str, Any]:
        """Build optimized context payload for an agent session.

        Allocates the token budget across quality standards, recent documents,
        and active patterns. Quality standards are loaded first (highest
        priority), with remaining budget split between documents and patterns.

        Args:
            project: Optional project name to scope document retrieval.
            max_tokens: Maximum token budget for the entire payload.

        Returns:
            Dict with keys: ``quality_standards``, ``recent_context``,
            ``active_patterns``.
        """
        quality_payload = self.standards.get_context_payload()
        quality_tokens = len(quality_payload) // 4
        remaining = max_tokens - quality_tokens

        recent_docs = self.doc_store.get_context(
            project=project, max_tokens=max(remaining - 500, 0)
        )

        active_patterns = self.patterns.get_for_session_load(max_count=20)

        return {
            "quality_standards": quality_payload,
            "recent_context": recent_docs,
            "active_patterns": active_patterns,
        }

    def record_request(
        self,
        session_id: str,
        query: str,
        doc_ids: list[str],
        tokens: int,
    ) -> None:
        """Record a context request for tracking and self-optimization.

        Stores which documents were served and how many tokens were used,
        enabling later analysis of retrieval quality via ``evaluate_relevance``.

        Args:
            session_id: The session that requested context.
            query: The query or context string that triggered retrieval.
            doc_ids: List of document IDs that were served.
            tokens: Total tokens served in this request.
        """
        self.db.execute(
            """INSERT INTO context_requests
               (session_id, query, documents_served, tokens_served)
               VALUES (?, ?, ?, ?)""",
            (session_id, query, json.dumps(doc_ids), tokens),
        )

    def evaluate_relevance(self, request_id: int, score: float) -> None:
        """Record a relevance score for a previously served context request.

        Used to close the feedback loop on context quality. Higher scores
        indicate the served context was useful to the agent.

        Args:
            request_id: ID of the context request to score.
            score: Relevance score (0.0 = irrelevant, 1.0 = highly relevant).
        """
        self.db.execute(
            "UPDATE context_requests SET relevance_score = ? WHERE id = ?",
            (score, request_id),
        )
