"""Internal published knowledge index registry.

Code outside ``mindroom.knowledge`` should use package facades such as
``mindroom.knowledge.status`` or ``mindroom.knowledge.utils`` instead of
importing this module directly.
"""

from __future__ import annotations

import asyncio
from dataclasses import dataclass, replace
from datetime import UTC, datetime
from pathlib import Path
from typing import TYPE_CHECKING, Literal, ParamSpec, Protocol, TypeVar, cast

from mindroom.embedding_factory import create_configured_embedder, embedder_client_signature
from mindroom.knowledge.availability import KnowledgeAvailability
from mindroom.knowledge.index_metadata import (
    PublishedIndexState,
    load_published_index_state,
    save_published_index_state,
)
from mindroom.knowledge.indexing_config import (
    IndexingSettings,
    chroma_collection_exists,
    indexing_settings_key,
    storage_key_for_base,
)
from mindroom.logging_config import get_logger
from mindroom.runtime_resolution import resolve_knowledge_binding
from mindroom.strict_knowledge import StrictSearchKnowledge

if TYPE_CHECKING:
    from collections.abc import Callable

    from agno.knowledge.knowledge import Knowledge

    from mindroom.config.main import Config
    from mindroom.constants import RuntimePaths
    from mindroom.runtime_resolution import ResolvedKnowledgeBinding
    from mindroom.tool_system.worker_routing import ToolExecutionIdentity

logger = get_logger(__name__)
# Identity levels:
# - KnowledgeSourceRoot: one physical source root. It gates source mutation locks and alias fanout.
# - KnowledgeRefreshTarget: one refresh target. It coalesces background work for a source and base ID.
# - PublishedIndexKey: one published, query-compatible index. It includes indexing settings for read paths.


@dataclass(frozen=True)
class PublishedIndexKey:
    """Stable key for one configured knowledge binding."""

    base_id: str
    storage_root: str
    knowledge_path: str
    indexing_settings: IndexingSettings


@dataclass(frozen=True)
class KnowledgeRefreshTarget:
    """Stable key for refresh work for one physical knowledge binding."""

    base_id: str
    storage_root: str
    knowledge_path: str


@dataclass(frozen=True)
class KnowledgeSourceRoot:
    """Stable key for source filesystem mutations shared by aliases."""

    storage_root: str
    knowledge_path: str


@dataclass(frozen=True)
class _PublishedIndexHandle:
    """Read handle for the published knowledge index."""

    key: PublishedIndexKey
    knowledge: Knowledge
    state: PublishedIndexState
    metadata_path: Path
    embedder_client_signature: str | None = None


@dataclass(frozen=True)
class PublishedIndexResolution:
    """Result of resolving the published index for one knowledge base."""

    key: PublishedIndexKey
    index: _PublishedIndexHandle | None
    state: PublishedIndexState | None
    availability: KnowledgeAvailability
    schedule_refresh_on_access: bool = False


class _PublishedIndexVectorDb(Protocol):
    client: object | None
    collection_name: str

    def exists(self) -> bool:
        """Return whether the collection exists."""
        ...


_published_indexes: dict[PublishedIndexKey, _PublishedIndexHandle] = {}
_CONSECUTIVE_REFRESH_FAILURE_ALERT_THRESHOLD = 3
_PRIVATE_KNOWLEDGE_BASE_ID_PREFIX = "__agent_private__:"
_MAX_PRIVATE_PUBLISHED_INDEXES = 128
_P = ParamSpec("_P")
_T = TypeVar("_T")


def _utc_now() -> str:
    return datetime.now(tz=UTC).isoformat()


async def _run_to_thread_to_completion_on_cancel(
    func: Callable[_P, _T],
    *args: _P.args,
    **kwargs: _P.kwargs,
) -> _T:
    thread_task = asyncio.create_task(asyncio.to_thread(func, *args, **kwargs))
    try:
        return await asyncio.shield(thread_task)
    except asyncio.CancelledError:
        await asyncio.shield(thread_task)
        raise


def _published_index_key_from_binding(
    base_id: str,
    binding: ResolvedKnowledgeBinding,
    *,
    config: Config,
) -> PublishedIndexKey:
    storage_root = binding.storage_root.expanduser().resolve()
    knowledge_path = binding.knowledge_path.resolve()
    return PublishedIndexKey(
        base_id=base_id,
        storage_root=str(storage_root),
        knowledge_path=str(knowledge_path),
        indexing_settings=indexing_settings_key(
            config,
            storage_root,
            base_id,
            knowledge_path,
        ),
    )


def _resolve_published_index_key_and_binding(
    base_id: str,
    *,
    config: Config,
    runtime_paths: RuntimePaths,
    execution_identity: ToolExecutionIdentity | None = None,
    create: bool = False,
) -> tuple[PublishedIndexKey, ResolvedKnowledgeBinding]:
    binding = resolve_knowledge_binding(
        base_id,
        config,
        runtime_paths,
        execution_identity=execution_identity,
        start_watchers=False,
        create=create,
    )
    return _published_index_key_from_binding(base_id, binding, config=config), binding


def resolve_published_index_key(
    base_id: str,
    *,
    config: Config,
    runtime_paths: RuntimePaths,
    execution_identity: ToolExecutionIdentity | None = None,
    create: bool = False,
) -> PublishedIndexKey:
    """Resolve one base ID to its current published index key."""
    key, _binding = _resolve_published_index_key_and_binding(
        base_id,
        config=config,
        runtime_paths=runtime_paths,
        execution_identity=execution_identity,
        create=create,
    )
    return key


def refresh_target_for_published_index_key(key: PublishedIndexKey) -> KnowledgeRefreshTarget:
    """Return the refresh target for one published index key."""
    return KnowledgeRefreshTarget(
        base_id=key.base_id,
        storage_root=key.storage_root,
        knowledge_path=key.knowledge_path,
    )


def source_root_for_refresh_target(key: KnowledgeRefreshTarget) -> KnowledgeSourceRoot:
    """Return the physical source root for one refresh target."""
    return KnowledgeSourceRoot(storage_root=key.storage_root, knowledge_path=key.knowledge_path)


def source_root_for_published_index_key(key: PublishedIndexKey) -> KnowledgeSourceRoot:
    """Return the physical source root for one published index key."""
    return KnowledgeSourceRoot(storage_root=key.storage_root, knowledge_path=key.knowledge_path)


def resolve_refresh_target(
    base_id: str,
    *,
    config: Config,
    runtime_paths: RuntimePaths,
    execution_identity: ToolExecutionIdentity | None = None,
    create: bool = False,
) -> KnowledgeRefreshTarget:
    """Resolve one base ID to its refresh target."""
    return refresh_target_for_published_index_key(
        resolve_published_index_key(
            base_id,
            config=config,
            runtime_paths=runtime_paths,
            execution_identity=execution_identity,
            create=create,
        ),
    )


def published_index_storage_path(key: PublishedIndexKey) -> Path:
    """Return the private storage directory backing one resolved knowledge base."""
    knowledge_path = Path(key.knowledge_path)
    return (Path(key.storage_root) / "knowledge_db" / storage_key_for_base(key.base_id, knowledge_path)).resolve()


def published_index_metadata_path(key: PublishedIndexKey) -> Path:
    """Return the single persisted state file for one knowledge base."""
    return published_index_storage_path(key) / "indexing_settings.json"


def published_index_refresh_state(
    state: PublishedIndexState | None,
    *,
    metadata_exists: bool = False,
) -> Literal["none", "stale", "refreshing", "refresh_failed"]:
    """Return the UI refresh state derived from the single metadata file."""
    if state is None:
        return "refresh_failed" if metadata_exists else "none"
    if state.status == "failed" or state.refresh_job == "failed" or state.last_error is not None:
        return "refresh_failed"
    if state.refresh_job == "running":
        return "refreshing"
    if state.refresh_job == "pending" or state.reason is not None:
        return "stale"
    return "none"


def _state_with_refresh_fields(
    key: PublishedIndexKey,
    *,
    refresh_job: Literal["idle", "pending", "running", "failed"],
    status_when_missing: Literal["indexing", "failed"],
    reason: str | None = None,
    last_error: str | None = None,
    clear_error: bool = False,
) -> PublishedIndexState:
    current = load_published_index_state(published_index_metadata_path(key))
    now = _utc_now()
    if current is None:
        return PublishedIndexState(
            settings=key.indexing_settings,
            status=status_when_missing,
            refresh_job=refresh_job,
            reason=reason,
            last_error=last_error,
            updated_at=now,
            last_refresh_at=now if refresh_job in {"idle", "failed"} else None,
        )
    return replace(
        current,
        refresh_job=refresh_job,
        reason=reason,
        last_error=None if clear_error else last_error,
        updated_at=now,
        last_refresh_at=now if refresh_job in {"idle", "failed"} else current.last_refresh_at,
    )


def mark_published_index_stale(
    key: PublishedIndexKey,
    *,
    reason: str,
    refresh_job: Literal["idle", "pending", "running", "failed"] = "pending",
) -> None:
    """Mark the published index stale without changing the last queryable index."""
    save_published_index_state(
        published_index_metadata_path(key),
        _state_with_refresh_fields(
            key,
            refresh_job=refresh_job,
            status_when_missing="indexing",
            reason=reason,
        ),
    )


def mark_published_index_refresh_running(key: PublishedIndexKey, *, reason: str = "refreshing") -> None:
    """Mark refresh work running while keeping the last queryable index readable."""
    save_published_index_state(
        published_index_metadata_path(key),
        _state_with_refresh_fields(
            key,
            refresh_job="running",
            status_when_missing="indexing",
            reason=reason,
        ),
    )


def mark_published_index_refresh_failed_preserving_last_good(key: PublishedIndexKey, *, error: str) -> None:
    """Record refresh failure while keeping any last queryable index readable."""
    current = load_published_index_state(published_index_metadata_path(key))
    state = _state_with_refresh_fields(
        key,
        refresh_job="failed",
        status_when_missing="failed",
        reason="refresh_failed",
        last_error=error,
    )
    consecutive_refresh_failures = (current.consecutive_refresh_failures if current is not None else 0) + 1
    state = replace(state, consecutive_refresh_failures=consecutive_refresh_failures)
    if current is not None and current.status == "complete":
        state = replace(state, status="complete")
    save_published_index_state(published_index_metadata_path(key), state)
    if consecutive_refresh_failures >= _CONSECUTIVE_REFRESH_FAILURE_ALERT_THRESHOLD:
        logger.error(
            "knowledge_refresh_failing_repeatedly",
            base_id=key.base_id,
            consecutive_refresh_failures=consecutive_refresh_failures,
            last_error=error,
        )


def mark_published_index_refresh_succeeded(key: PublishedIndexKey) -> None:
    """Clear refresh status after a successful publish."""
    state = load_published_index_state(published_index_metadata_path(key))
    if state is None:
        return
    save_published_index_state(
        published_index_metadata_path(key),
        replace(
            state,
            refresh_job="idle",
            reason=None,
            last_error=None,
            updated_at=_utc_now(),
            last_refresh_at=_utc_now(),
            consecutive_refresh_failures=0,
        ),
    )


def _state_collection_name(state: PublishedIndexState) -> str:
    if state.collection is None:
        msg = "Published knowledge metadata is missing a collection name"
        raise ValueError(msg)
    return state.collection


def _build_published_index_vector_db(
    key: PublishedIndexKey,
    state: PublishedIndexState,
    *,
    config: Config,
    runtime_paths: RuntimePaths,
) -> _PublishedIndexVectorDb:
    from agno.vectordb.chroma import ChromaDb  # noqa: PLC0415

    return cast(
        "_PublishedIndexVectorDb",
        ChromaDb(
            collection=_state_collection_name(state),
            path=str(published_index_storage_path(key)),
            persistent_client=True,
            embedder=create_configured_embedder(config, runtime_paths),
        ),
    )


def _build_published_index_knowledge(
    key: PublishedIndexKey,
    state: PublishedIndexState,
    *,
    config: Config,
    runtime_paths: RuntimePaths,
) -> Knowledge:
    return StrictSearchKnowledge(
        vector_db=_build_published_index_vector_db(key, state, config=config, runtime_paths=runtime_paths),
    )


def published_index_collection_exists_for_state(key: PublishedIndexKey, state: PublishedIndexState) -> bool:
    """Return whether persisted metadata points at an existing Chroma collection."""
    if state.status != "complete" or state.collection is None:
        return False
    try:
        return chroma_collection_exists(published_index_storage_path(key), state.collection)
    except Exception:
        logger.warning(
            "Published knowledge collection existence check failed",
            base_id=key.base_id,
            collection=state.collection,
            exc_info=True,
        )
        return False


def _indexing_settings_query_compatible(
    published_settings: IndexingSettings,
    current_settings: IndexingSettings,
) -> bool:
    """Return whether current queries can use a collection from published settings."""
    return published_settings.query_compatibility_key() == current_settings.query_compatibility_key()


def published_index_settings_compatible(
    published_settings: IndexingSettings,
    current_settings: IndexingSettings,
) -> bool:
    """Return whether a published index can be queried under the current config."""
    return (
        _indexing_settings_query_compatible(
            published_settings,
            current_settings,
        )
        and published_settings.corpus_compatibility_key() == current_settings.corpus_compatibility_key()
    )


def _published_index_state_queryable(key: PublishedIndexKey, state: PublishedIndexState) -> bool:
    return (
        state.status == "complete"
        and state.collection is not None
        and published_index_settings_compatible(state.settings, key.indexing_settings)
    )


def _published_index_availability(
    *,
    key: PublishedIndexKey,
    state: PublishedIndexState | None,
    metadata_exists: bool = False,
) -> KnowledgeAvailability:
    refresh_state = published_index_refresh_state(state, metadata_exists=metadata_exists)
    if state is None:
        availability = (
            KnowledgeAvailability.REFRESH_FAILED
            if refresh_state == "refresh_failed"
            else KnowledgeAvailability.INITIALIZING
        )
    elif state.collection is None and refresh_state == "refresh_failed":
        availability = KnowledgeAvailability.REFRESH_FAILED
    elif (
        not published_index_settings_compatible(
            state.settings,
            key.indexing_settings,
        )
        or state.settings != key.indexing_settings
    ):
        availability = KnowledgeAvailability.CONFIG_MISMATCH
    elif refresh_state == "refresh_failed":
        availability = KnowledgeAvailability.REFRESH_FAILED
    elif state.status != "complete":
        availability = KnowledgeAvailability.INITIALIZING
    elif refresh_state in {"stale", "refreshing"}:
        availability = KnowledgeAvailability.STALE
    else:
        availability = KnowledgeAvailability.READY
    return availability


def published_index_availability_for_state(
    *,
    key: PublishedIndexKey,
    state: PublishedIndexState | None,
    metadata_exists: bool = False,
) -> KnowledgeAvailability:
    """Return the public availability value for published index state."""
    return _published_index_availability(
        key=key,
        state=state,
        metadata_exists=metadata_exists,
    )


def _cached_index_still_queryable(index: _PublishedIndexHandle) -> bool:
    if not _published_index_state_queryable(index.key, index.state):
        return False
    vector_db = cast("_PublishedIndexVectorDb | None", index.knowledge.vector_db)
    return vector_db is not None and vector_db.exists()


def _cached_index_matches_persisted_state(
    index: _PublishedIndexHandle,
    state: PublishedIndexState,
) -> bool:
    """Return whether a process-local handle still points at persisted metadata."""
    return (
        index.state.settings == state.settings
        and index.state.status == state.status
        and index.state.collection == state.collection
        and index.state.last_published_at == state.last_published_at
        and index.state.published_revision == state.published_revision
        and index.state.indexed_count == state.indexed_count
        and index.state.source_signature == state.source_signature
    )


def _load_queryable_index_from_state(
    key: PublishedIndexKey,
    state: PublishedIndexState,
    *,
    config: Config,
    runtime_paths: RuntimePaths,
) -> Knowledge | None:
    if not _published_index_state_queryable(key, state):
        return None
    if not published_index_collection_exists_for_state(key, state):
        return None
    return _build_published_index_knowledge(key, state, config=config, runtime_paths=runtime_paths)


def get_published_index(
    base_id: str,
    *,
    config: Config,
    runtime_paths: RuntimePaths,
    execution_identity: ToolExecutionIdentity | None = None,
) -> PublishedIndexResolution:
    """Return the currently published index, if one is usable."""
    key, binding = _resolve_published_index_key_and_binding(
        base_id,
        config=config,
        runtime_paths=runtime_paths,
        execution_identity=execution_identity,
        create=False,
    )
    metadata_path = published_index_metadata_path(key)
    state = load_published_index_state(metadata_path)
    availability = _published_index_availability(key=key, state=state, metadata_exists=metadata_path.exists())
    current_embedder_client_signature = embedder_client_signature(config, runtime_paths)

    index = _published_indexes.get(key)
    if index is not None:
        if (
            index.embedder_client_signature == current_embedder_client_signature
            and state is not None
            and _cached_index_matches_persisted_state(index, state)
            and _cached_index_still_queryable(index)
        ):
            if index.state != state:
                index = replace(index, state=state)
                _published_indexes[key] = index
            return PublishedIndexResolution(
                key=key,
                index=index,
                state=state,
                availability=availability,
                schedule_refresh_on_access=binding.incremental_sync_on_access,
            )
        _published_indexes.pop(key, None)

    if state is None:
        return PublishedIndexResolution(
            key=key,
            index=None,
            state=state,
            availability=availability,
            schedule_refresh_on_access=binding.incremental_sync_on_access,
        )

    try:
        knowledge = _load_queryable_index_from_state(key, state, config=config, runtime_paths=runtime_paths)
    except Exception:
        logger.warning(
            "Published knowledge index handle could not be opened",
            base_id=base_id,
            exc_info=True,
        )
        return PublishedIndexResolution(
            key=key,
            index=None,
            state=state,
            availability=KnowledgeAvailability.REFRESH_FAILED,
            schedule_refresh_on_access=binding.incremental_sync_on_access,
        )
    if knowledge is None:
        return PublishedIndexResolution(
            key=key,
            index=None,
            state=state,
            availability=availability
            if availability is not KnowledgeAvailability.READY
            else KnowledgeAvailability.REFRESH_FAILED,
            schedule_refresh_on_access=binding.incremental_sync_on_access,
        )

    index = _PublishedIndexHandle(
        key=key,
        knowledge=knowledge,
        state=state,
        metadata_path=published_index_metadata_path(key),
        embedder_client_signature=current_embedder_client_signature,
    )
    _cache_published_index(index)
    return PublishedIndexResolution(
        key=key,
        index=index,
        state=state,
        availability=availability,
        schedule_refresh_on_access=binding.incremental_sync_on_access,
    )


def _publish_knowledge_index(
    key: PublishedIndexKey,
    *,
    knowledge: Knowledge,
    state: PublishedIndexState,
    metadata_path: Path | None = None,
    embedder_client_signature: str | None = None,
) -> _PublishedIndexHandle:
    """Publish a read handle in this process."""
    _evict_published_indexes_for_refresh_target(refresh_target_for_published_index_key(key))
    index = _PublishedIndexHandle(
        key=key,
        knowledge=knowledge,
        state=state,
        metadata_path=metadata_path or published_index_metadata_path(key),
        embedder_client_signature=embedder_client_signature,
    )
    _cache_published_index(index)
    return index


def publish_knowledge_index_from_state(
    key: PublishedIndexKey,
    *,
    state: PublishedIndexState,
    config: Config,
    runtime_paths: RuntimePaths,
    metadata_path: Path | None = None,
) -> _PublishedIndexHandle | None:
    """Publish a read handle rebuilt from persisted metadata."""
    knowledge = _load_queryable_index_from_state(key, state, config=config, runtime_paths=runtime_paths)
    if knowledge is None:
        return None
    return _publish_knowledge_index(
        key,
        knowledge=knowledge,
        state=state,
        metadata_path=metadata_path,
        embedder_client_signature=embedder_client_signature(config, runtime_paths),
    )


def _same_physical_binding(key: PublishedIndexKey, refresh_target: KnowledgeRefreshTarget) -> bool:
    return (
        key.base_id == refresh_target.base_id
        and key.storage_root == refresh_target.storage_root
        and key.knowledge_path == refresh_target.knowledge_path
    )


def _same_physical_source(left: PublishedIndexKey, right: PublishedIndexKey) -> bool:
    return left.storage_root == right.storage_root and left.knowledge_path == right.knowledge_path


def _published_index_key_is_private(key: PublishedIndexKey) -> bool:
    return key.base_id.startswith(_PRIVATE_KNOWLEDGE_BASE_ID_PREFIX)


def prune_private_index_bookkeeping() -> None:
    """Bound PrivateAgentKnowledge in-process published index handles."""
    private_index_keys = [key for key in _published_indexes if _published_index_key_is_private(key)]
    for key in private_index_keys[:-_MAX_PRIVATE_PUBLISHED_INDEXES]:
        _published_indexes.pop(key, None)


def _cache_published_index(index: _PublishedIndexHandle) -> None:
    _published_indexes[index.key] = index
    prune_private_index_bookkeeping()


def _evict_published_indexes_for_refresh_target(refresh_target: KnowledgeRefreshTarget) -> None:
    for cached_key in tuple(_published_indexes):
        if _same_physical_binding(cached_key, refresh_target):
            _published_indexes.pop(cached_key, None)


def _published_index_keys_for_shared_source(
    base_id: str,
    *,
    config: Config,
    runtime_paths: RuntimePaths,
    execution_identity: ToolExecutionIdentity | None = None,
) -> tuple[PublishedIndexKey, ...]:
    base_mode = config.get_knowledge_base_config(base_id).mode
    key = resolve_published_index_key(
        base_id,
        config=config,
        runtime_paths=runtime_paths,
        execution_identity=execution_identity,
        create=False,
    )
    matching_keys = [key] if base_mode == "semantic" else []
    for candidate_base_id in config.knowledge_bases:
        if candidate_base_id == base_id:
            continue
        if config.get_knowledge_base_config(candidate_base_id).mode != "semantic":
            continue
        try:
            candidate_key = resolve_published_index_key(
                candidate_base_id,
                config=config,
                runtime_paths=runtime_paths,
                execution_identity=execution_identity,
                create=False,
            )
        except Exception:
            logger.warning(
                "Could not resolve related published knowledge index while marking source changed",
                base_id=base_id,
                related_base_id=candidate_base_id,
                exc_info=True,
            )
            continue
        if _same_physical_source(candidate_key, key):
            matching_keys.append(candidate_key)
    return tuple(matching_keys)


def _mark_published_index_key_stale_on_disk(matching_key: PublishedIndexKey, *, reason: str) -> bool:
    mark_published_index_stale_and_evict(matching_key, reason=reason)
    return True


def mark_published_index_stale_and_evict(matching_key: PublishedIndexKey, *, reason: str) -> bool:
    """Mark one published index stale and evict matching process-local handles."""
    mark_published_index_stale(matching_key, reason=reason)
    _evict_published_indexes_for_refresh_target(refresh_target_for_published_index_key(matching_key))
    return True


def _mark_knowledge_source_changed(
    base_id: str,
    *,
    config: Config,
    runtime_paths: RuntimePaths,
    execution_identity: ToolExecutionIdentity | None = None,
    reason: str = "source_mutated",
) -> tuple[str, ...]:
    """Mark same-source published indexes stale after a source mutation."""
    matching_keys = _published_index_keys_for_shared_source(
        base_id,
        config=config,
        runtime_paths=runtime_paths,
        execution_identity=execution_identity,
    )
    for matching_key in matching_keys:
        _mark_published_index_key_stale_on_disk(matching_key, reason=reason)
    return tuple(dict.fromkeys(key.base_id for key in matching_keys))


async def mark_knowledge_source_changed_async(
    base_id: str,
    *,
    config: Config,
    runtime_paths: RuntimePaths,
    execution_identity: ToolExecutionIdentity | None = None,
    reason: str = "source_mutated",
) -> tuple[str, ...]:
    """Async stale marker that keeps metadata I/O off the event loop."""
    return await _run_to_thread_to_completion_on_cancel(
        _mark_knowledge_source_changed,
        base_id,
        config=config,
        runtime_paths=runtime_paths,
        execution_identity=execution_identity,
        reason=reason,
    )
