# Copyright (C) 2026 GoodData Corporation

"""Embedding initialization for rule retrieval."""

from __future__ import annotations

import logging
from contextlib import suppress
from pathlib import Path
from typing import TYPE_CHECKING, Any

from aida_mcp.tools.rules.embeddings.chunker import RuleChunker
from aida_mcp.tools.rules.embeddings.model import get_embedding_model, is_model_cached
from aida_mcp.tools.rules.embeddings.store import VectorStore
from aida_mcp.tools.rules.store import get_rule_store

if TYPE_CHECKING:
    from mcp.server.fastmcp import Context as MCPContext

    Context = MCPContext[Any, Any, Any]

logger = logging.getLogger(__name__)

# Process-wide cache of embedding initialization state keyed by materialized rules directory.
_VECTOR_STORES: dict[Path, VectorStore | None] = {}
_EMBEDDINGS_INITIALIZED: set[Path] = set()


def get_vector_store(rules_dir: Path) -> VectorStore | None:
    """Get the vector store for a materialized rules directory, if initialized."""
    return _VECTOR_STORES.get(rules_dir.resolve())


async def initialize_embeddings_async(
    rules_dir: Path,
    ctx: Context | None = None,
    *,
    allow_download: bool = False,
    suppress_errors: bool = True,
) -> bool:
    """
    Initialize the embedding index for all rules with progress reporting.

    By default, this function will **not** download the embedding model. This keeps `get_rules`
    deterministic and avoids surprise network calls. Use `allow_download=True` from an explicit
    warmup tool or installer workflow when you want to populate caches.

    Args:
        rules_dir: Path to the rules directory
        ctx: Optional MCP context for progress reporting
        allow_download: If True, allow fetching missing model files from HuggingFace Hub
        suppress_errors: If True, swallow initialization errors and fall back to BM25-only retrieval

    Returns:
        True if embeddings were initialized and are available, False otherwise.
    """
    key = rules_dir.resolve()
    if key in _EMBEDDINGS_INITIALIZED:
        return _VECTOR_STORES.get(key) is not None

    logger.info(f"Initializing embeddings for rules in {key}")

    # Total steps: load rules, chunk, load model
    # Note: "generate embeddings" and "build index" steps are not notified
    # to avoid race condition with tool completion (client invalidates progress
    # token as soon as it receives tool result)
    total_steps = 3

    # Step 1: Loading rules
    if ctx:
        with suppress(Exception):
            await ctx.report_progress(1, total_steps, "🔍 Loading rules...")

    store = get_rule_store(key)
    rules = store.get_all_rules()

    if not rules:
        logger.warning("No rules found, skipping embedding initialization")
        _EMBEDDINGS_INITIALIZED.add(key)
        _VECTOR_STORES[key] = None
        return False

    # Step 2: Chunking
    if ctx:
        with suppress(Exception):
            await ctx.report_progress(2, total_steps, f"📄 Chunking {len(rules)} rules for semantic search...")

    chunker = RuleChunker()
    chunks = chunker.chunk_rules(rules)

    if not allow_download and not is_model_cached():
        logger.info(
            "Embedding model not cached; skipping semantic index init (BM25-only). "
            "Run warmup tool to download model and enable semantic retrieval."
        )
        return False

    try:
        # Step 3: Loading embedding model (may download ~90MB on first run)
        # Important: Avoid report_progress(3/3). Some MCP clients invalidate progress tokens early,
        # which can surface as noisy "unknown token" errors during first-time model load.
        model = get_embedding_model()
        # Force model loading now (triggers download if needed)
        _ = model.dimension

        # Steps 4-5: Generating embeddings and building index
        # Note: We intentionally do NOT send progress notifications for these final steps.
        texts = [chunk.text for chunk in chunks]
        embeddings = model.embed(texts)

        vector_store = VectorStore()
        vector_store.add(chunks, embeddings)

        _VECTOR_STORES[key] = vector_store
        _EMBEDDINGS_INITIALIZED.add(key)

        logger.info(f"Initialized embeddings: {len(rules)} rules, {len(chunks)} chunks")
        return True
    except Exception:
        if not suppress_errors:
            raise
        logger.exception("Failed to initialize embeddings; falling back to BM25-only retrieval")
        _VECTOR_STORES[key] = None
        _EMBEDDINGS_INITIALIZED.add(key)
        return False


def initialize_embeddings(
    rules_dir: Path,
    *,
    allow_download: bool = False,
    suppress_errors: bool = True,
) -> bool:
    """
    Synchronous embedding initialization (for backward compatibility/testing).

    For production use with progress reporting, use initialize_embeddings_async().
    """
    key = rules_dir.resolve()
    if key in _EMBEDDINGS_INITIALIZED:
        return _VECTOR_STORES.get(key) is not None

    logger.info(f"Initializing embeddings for rules in {key}")

    store = get_rule_store(key)
    rules = store.get_all_rules()

    if not rules:
        logger.warning("No rules found, skipping embedding initialization")
        _EMBEDDINGS_INITIALIZED.add(key)
        _VECTOR_STORES[key] = None
        return False

    chunker = RuleChunker()
    chunks = chunker.chunk_rules(rules)

    if not allow_download and not is_model_cached():
        logger.info(
            "Embedding model not cached; skipping semantic index init (BM25-only). "
            "Run warmup tool to download model and enable semantic retrieval."
        )
        return False

    try:
        model = get_embedding_model()
        texts = [chunk.text for chunk in chunks]
        embeddings = model.embed(texts)

        vector_store = VectorStore()
        vector_store.add(chunks, embeddings)

        _VECTOR_STORES[key] = vector_store
        _EMBEDDINGS_INITIALIZED.add(key)

        logger.info(f"Initialized embeddings: {len(rules)} rules, {len(chunks)} chunks")
        return True
    except Exception:
        if not suppress_errors:
            raise
        logger.exception("Failed to initialize embeddings; falling back to BM25-only retrieval")
        _VECTOR_STORES[key] = None
        _EMBEDDINGS_INITIALIZED.add(key)
        return False
