DataOps-platform/app/core/knowledge/repository.py

from __future__ import annotations

import json
from collections.abc import Mapping, Sequence
from typing import Any

from sqlalchemy import text
from sqlalchemy.orm import Session

from app.core.common.identifiers import new_governance_uid
from app.core.knowledge.chunking import embedding_input_hash
from app.core.knowledge.contracts import (
    KnowledgeDependencyDraft,
    KnowledgePointDraft,
    KnowledgeSnapshot,
    PointChange,
    PreparedPublication,
)


def _json(value: object) -> str:
    return json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))


def _mapping(value: Any) -> Mapping[str, Any]:
    if isinstance(value, Mapping):
        return value
    if isinstance(value, str):
        decoded = json.loads(value)
        if isinstance(decoded, Mapping):
            return decoded
    return {}


def _vector(value: Sequence[float]) -> str:
    return "[" + ",".join(str(float(item)) for item in value) + "]"


class SqlKnowledgeRepository:
    """Persist a canonical knowledge publication in the caller's transaction."""

    def __init__(
        self,
        session: Session,
        *,
        embedding_profile_id: str,
        embedding_profile_key: str,
        generation: int,
        workspace: str,
    ) -> None:
        self._session = session
        self._embedding_profile_id = embedding_profile_id
        self._embedding_profile_key = embedding_profile_key
        self._generation = generation
        self._workspace = workspace

    def load_active_snapshot(
        self, source_type: str, source_uid: str
    ) -> KnowledgeSnapshot | None:
        document = (
            self._session.execute(
                text(
                    """
                SELECT object_version, content_hash, point_set_hash,
                       permission_scope, source_updated_at
                FROM public.governance_documents
                WHERE object_uid = CAST(:source_uid AS uuid)
                  AND object_type = :source_type
                  AND status = 'active'
                """
                ),
                {"source_type": source_type, "source_uid": source_uid},
            )
            .mappings()
            .one_or_none()
        )
        if document is None:
            return None

        point_rows = self._session.execute(
            text(
                """
                SELECT point_key, semantic_path, content, content_hash,
                       metadata, metadata_hash, permission_scope, permission_hash
                FROM public.knowledge_points
                WHERE source_type = :source_type
                  AND source_uid = CAST(:source_uid AS uuid)
                  AND status = 'active'
                ORDER BY point_key
                """
            ),
            {"source_type": source_type, "source_uid": source_uid},
        ).mappings()
        points = tuple(
            KnowledgePointDraft(
                point_key=row["point_key"],
                semantic_path=row["semantic_path"],
                content=row["content"],
                content_hash=row["content_hash"],
                metadata=_mapping(row["metadata"]),
                metadata_hash=row["metadata_hash"],
                permission_scope=_mapping(row["permission_scope"]),
                permission_hash=row["permission_hash"],
            )
            for row in point_rows
        )
        prefix = f"{source_type}/{source_uid}/%"
        dependency_rows = self._session.execute(
            text(
                """
                SELECT from_point_key, to_point_key, relation_type, source
                FROM public.knowledge_point_dependencies
                WHERE from_point_key LIKE :prefix AND status = 'active'
                ORDER BY from_point_key, to_point_key, relation_type, source
                """
            ),
            {"prefix": prefix},
        ).mappings()
        dependencies = tuple(
            KnowledgeDependencyDraft(
                from_point_key=row["from_point_key"],
                to_point_key=row["to_point_key"],
                relation_type=row["relation_type"],
                source=row["source"],
            )
            for row in dependency_rows
        )
        updated_at = document["source_updated_at"]
        return KnowledgeSnapshot(
            source_type=source_type,
            source_uid=source_uid,
            source_revision=int(document["object_version"]),
            source_snapshot_hash=document["content_hash"],
            point_set_hash=document["point_set_hash"],
            permission_scope=_mapping(document["permission_scope"]),
            points=points,
            dependencies=dependencies,
            source_updated_at=updated_at.isoformat() if updated_at else None,
        )

    def load_embedding(self, embedding_hash: str) -> Sequence[float] | None:
        found = self._session.execute(
            text(
                """
                SELECT 1
                FROM public.knowledge_chunk_embeddings
                WHERE profile_id = CAST(:profile_id AS uuid)
                  AND embedding_hash = :embedding_hash
                LIMIT 1
                """
            ),
            {
                "profile_id": self._embedding_profile_id,
                "embedding_hash": embedding_hash,
            },
        ).scalar_one_or_none()
        return (0.0,) if found is not None else None

    def activate(self, publication: PreparedPublication) -> None:
        snapshot = publication.snapshot
        self._lock_source(snapshot.source_type, snapshot.source_uid)
        self._assert_revision_is_publishable(snapshot)

        change_set_id = new_governance_uid()
        correlation_id = new_governance_uid()
        document_id = new_governance_uid()
        change_groups = (
            publication.diff.added,
            publication.diff.modified,
            publication.diff.deleted,
        )
        change_type = "create" if not self._active_document_id(snapshot) else "update"
        self._insert_change_set(
            change_set_id=change_set_id,
            correlation_id=correlation_id,
            publication=publication,
            change_type=change_type,
        )

        old_revisions = self._load_active_point_revisions(snapshot)
        for group in change_groups:
            for change in group:
                self._retire_old_point(change)

        self._retire_active_document(snapshot)
        self._insert_document(document_id, change_set_id, snapshot)

        new_revisions: dict[str, int] = {}
        for change in (*publication.diff.added, *publication.diff.modified):
            point_revision = self._next_point_revision(change.point_key)
            new_revisions[change.point_key] = point_revision
            self._insert_point(change.new, snapshot, point_revision)

        self._replace_dependencies(snapshot)
        self._insert_chunks_and_embeddings(
            document_id=document_id,
            change_set_id=change_set_id,
            publication=publication,
        )
        self._insert_change_items(
            change_set_id=change_set_id,
            publication=publication,
            old_revisions=old_revisions,
            new_revisions=new_revisions,
        )
        self._invalidate_cache_dependencies(publication)
        self._insert_projections(document_id, change_set_id, snapshot)
        self._session.execute(
            text(
                """
                UPDATE public.knowledge_change_sets
                SET status = 'canonical_active',
                    updated_at = CURRENT_TIMESTAMP,
                    activated_at = CURRENT_TIMESTAMP
                WHERE id = CAST(:id AS uuid)
                """
            ),
            {"id": change_set_id},
        )

    def _lock_source(self, source_type: str, source_uid: str) -> None:
        self._session.execute(
            text("SELECT pg_advisory_xact_lock(hashtextextended(:key, 0))"),
            {"key": f"knowledge:{source_type}:{source_uid}"},
        )

    def _active_document_id(self, snapshot: KnowledgeSnapshot) -> str | None:
        return self._session.execute(
            text(
                """
                SELECT CAST(id AS text)
                FROM public.governance_documents
                WHERE object_uid = CAST(:source_uid AS uuid) AND status = 'active'
                """
            ),
            {"source_uid": snapshot.source_uid},
        ).scalar_one_or_none()

    def _assert_revision_is_publishable(self, snapshot: KnowledgeSnapshot) -> None:
        active_revision = self._session.execute(
            text(
                """
                SELECT object_version
                FROM public.governance_documents
                WHERE object_uid = CAST(:source_uid AS uuid) AND status = 'active'
                FOR UPDATE
                """
            ),
            {"source_uid": snapshot.source_uid},
        ).scalar_one_or_none()
        if active_revision is not None and snapshot.source_revision <= active_revision:
            raise ValueError("source revision is no longer publishable")

    def _insert_change_set(
        self,
        *,
        change_set_id: str,
        correlation_id: str,
        publication: PreparedPublication,
        change_type: str,
    ) -> None:
        snapshot = publication.snapshot
        self._session.execute(
            text(
                """
                INSERT INTO public.knowledge_change_sets (
                    id, correlation_id, source_type, source_uid, source_revision,
                    change_type, source_snapshot_hash, added_count, modified_count,
                    deleted_count, status, target_generation
                ) VALUES (
                    CAST(:id AS uuid), CAST(:correlation_id AS uuid), :source_type,
                    CAST(:source_uid AS uuid), :source_revision, :change_type,
                    :snapshot_hash, :added_count, :modified_count, :deleted_count,
                    'validating', :generation
                )
                """
            ),
            {
                "id": change_set_id,
                "correlation_id": correlation_id,
                "source_type": snapshot.source_type,
                "source_uid": snapshot.source_uid,
                "source_revision": snapshot.source_revision,
                "change_type": change_type,
                "snapshot_hash": snapshot.source_snapshot_hash,
                "added_count": len(publication.diff.added),
                "modified_count": len(publication.diff.modified),
                "deleted_count": len(publication.diff.deleted),
                "generation": self._generation,
            },
        )

    def _load_active_point_revisions(
        self, snapshot: KnowledgeSnapshot
    ) -> dict[str, int]:
        rows = self._session.execute(
            text(
                """
                SELECT point_key, point_revision
                FROM public.knowledge_points
                WHERE source_type = :source_type
                  AND source_uid = CAST(:source_uid AS uuid)
                  AND status = 'active'
                """
            ),
            {"source_type": snapshot.source_type, "source_uid": snapshot.source_uid},
        )
        return {row[0]: int(row[1]) for row in rows}

    def _retire_old_point(self, change: PointChange) -> None:
        if change.old is None:
            return
        target_status = "deleted" if change.change_kind == "deleted" else "superseded"
        self._session.execute(
            text(
                """
                UPDATE public.knowledge_points
                SET status = :status, valid_to = CURRENT_TIMESTAMP
                WHERE point_key = :point_key AND status = 'active'
                """
            ),
            {"status": target_status, "point_key": change.point_key},
        )

    def _retire_active_document(self, snapshot: KnowledgeSnapshot) -> None:
        self._session.execute(
            text(
                """
                UPDATE public.governance_documents
                SET status = 'superseded'
                WHERE object_uid = CAST(:source_uid AS uuid) AND status = 'active'
                """
            ),
            {"source_uid": snapshot.source_uid},
        )

    def _insert_document(
        self, document_id: str, change_set_id: str, snapshot: KnowledgeSnapshot
    ) -> None:
        point_by_path = {
            point.semantic_path: point.content for point in snapshot.points
        }
        object_name = (
            point_by_path.get("name") or f"{snapshot.source_type}:{snapshot.source_uid}"
        )
        domains = snapshot.permission_scope.get("business_domains", [])
        business_domain_uid = str(domains[0]) if domains else None
        content = _json(
            {
                "source_type": snapshot.source_type,
                "source_uid": snapshot.source_uid,
                "source_revision": snapshot.source_revision,
                "points": [
                    {"point_key": point.point_key, "content": point.content}
                    for point in snapshot.points
                ],
            }
        )
        self._session.execute(
            text(
                """
                INSERT INTO public.governance_documents (
                    id, object_uid, object_type, object_version, object_name,
                    business_domain_uid, permission_scope, content, content_hash,
                    status, source_updated_at, point_set_hash, active_generation,
                    change_set_id
                ) VALUES (
                    CAST(:id AS uuid), CAST(:source_uid AS uuid), :source_type,
                    :source_revision, :object_name, CAST(:business_domain_uid AS uuid),
                    CAST(:permission_scope AS jsonb), :content, :snapshot_hash,
                    'active', CAST(:source_updated_at AS timestamptz), :point_set_hash,
                    :generation, CAST(:change_set_id AS uuid)
                )
                """
            ),
            {
                "id": document_id,
                "source_uid": snapshot.source_uid,
                "source_type": snapshot.source_type,
                "source_revision": snapshot.source_revision,
                "object_name": object_name,
                "business_domain_uid": business_domain_uid,
                "permission_scope": _json(snapshot.permission_scope),
                "content": content,
                "snapshot_hash": snapshot.source_snapshot_hash,
                "source_updated_at": snapshot.source_updated_at,
                "point_set_hash": snapshot.point_set_hash,
                "generation": self._generation,
                "change_set_id": change_set_id,
            },
        )

    def _next_point_revision(self, point_key: str) -> int:
        return int(
            self._session.execute(
                text(
                    """
                    SELECT COALESCE(MAX(point_revision), 0) + 1
                    FROM public.knowledge_points
                    WHERE point_key = :point_key
                    """
                ),
                {"point_key": point_key},
            ).scalar_one()
        )

    def _insert_point(
        self,
        point: KnowledgePointDraft | None,
        snapshot: KnowledgeSnapshot,
        point_revision: int,
    ) -> None:
        if point is None:
            raise ValueError("added or modified point is missing its new value")
        self._session.execute(
            text(
                """
                INSERT INTO public.knowledge_points (
                    id, point_key, point_revision, source_type, source_uid,
                    source_revision, semantic_path, content, content_hash,
                    metadata, metadata_hash, permission_scope, permission_hash,
                    status, valid_from, activated_at
                ) VALUES (
                    CAST(:id AS uuid), :point_key, :point_revision, :source_type,
                    CAST(:source_uid AS uuid), :source_revision, :semantic_path,
                    :content, :content_hash, CAST(:metadata AS jsonb), :metadata_hash,
                    CAST(:permission_scope AS jsonb), :permission_hash, 'active',
                    CURRENT_TIMESTAMP, CURRENT_TIMESTAMP
                )
                """
            ),
            {
                "id": new_governance_uid(),
                "point_key": point.point_key,
                "point_revision": point_revision,
                "source_type": snapshot.source_type,
                "source_uid": snapshot.source_uid,
                "source_revision": snapshot.source_revision,
                "semantic_path": point.semantic_path,
                "content": point.content,
                "content_hash": point.content_hash,
                "metadata": _json(point.metadata),
                "metadata_hash": point.metadata_hash,
                "permission_scope": _json(point.permission_scope),
                "permission_hash": point.permission_hash,
            },
        )

    def _replace_dependencies(self, snapshot: KnowledgeSnapshot) -> None:
        prefix = f"{snapshot.source_type}/{snapshot.source_uid}/%"
        self._session.execute(
            text(
                """
                UPDATE public.knowledge_point_dependencies
                SET status = 'superseded'
                WHERE from_point_key LIKE :prefix AND status = 'active'
                """
            ),
            {"prefix": prefix},
        )
        statement = text(
            """
            INSERT INTO public.knowledge_point_dependencies (
                from_point_key, to_point_key, relation_type, source, generation, status
            ) VALUES (
                :from_point_key, :to_point_key, :relation_type, :source,
                :generation, 'active'
            )
            ON CONFLICT (
                from_point_key, to_point_key, relation_type, source, generation
            ) DO UPDATE SET status = 'active'
            """
        )
        for dependency in snapshot.dependencies:
            self._session.execute(
                statement,
                {
                    "from_point_key": dependency.from_point_key,
                    "to_point_key": dependency.to_point_key,
                    "relation_type": dependency.relation_type,
                    "source": dependency.source,
                    "generation": self._generation,
                },
            )

    def _insert_chunks_and_embeddings(
        self,
        *,
        document_id: str,
        change_set_id: str,
        publication: PreparedPublication,
    ) -> None:
        chunk_statement = text(
            """
            INSERT INTO public.governance_chunks (
                id, document_id, chunk_no, content, content_hash, chunk_kind,
                section_path, metadata, lexical_text, search_vector,
                primary_point_key, point_keys, point_set_hash, change_set_id
            ) VALUES (
                CAST(:id AS uuid), CAST(:document_id AS uuid), :chunk_no, :content,
                :content_hash, :chunk_kind, :section_path, CAST(:metadata AS jsonb),
                :lexical_text, to_tsvector('simple', COALESCE(:lexical_text, '')),
                :primary_point_key, CAST(:point_keys AS jsonb), :point_set_hash,
                CAST(:change_set_id AS uuid)
            )
            """
        )
        for chunk_no, chunk in enumerate(publication.chunks):
            chunk_id = new_governance_uid()
            self._session.execute(
                chunk_statement,
                {
                    "id": chunk_id,
                    "document_id": document_id,
                    "chunk_no": chunk_no,
                    "content": chunk.content,
                    "content_hash": chunk.content_hash,
                    "chunk_kind": chunk.chunk_kind,
                    "section_path": chunk.section_path,
                    "metadata": _json(chunk.metadata),
                    "lexical_text": chunk.lexical_text,
                    "primary_point_key": chunk.primary_point_key,
                    "point_keys": _json(chunk.point_keys),
                    "point_set_hash": chunk.point_set_hash,
                    "change_set_id": change_set_id,
                },
            )
            embedding_hash = embedding_input_hash(chunk, self._embedding_profile_key)
            vector = publication.embeddings.get(embedding_hash)
            if vector is not None:
                self._session.execute(
                    text(
                        """
                        INSERT INTO public.knowledge_chunk_embeddings (
                            id, chunk_id, profile_id, embedding, embedding_hash
                        ) VALUES (
                            CAST(:id AS uuid), CAST(:chunk_id AS uuid),
                            CAST(:profile_id AS uuid), CAST(:embedding AS vector),
                            :embedding_hash
                        )
                        """
                    ),
                    {
                        "id": new_governance_uid(),
                        "chunk_id": chunk_id,
                        "profile_id": self._embedding_profile_id,
                        "embedding": _vector(vector),
                        "embedding_hash": embedding_hash,
                    },
                )
                continue
            inserted = self._session.execute(
                text(
                    """
                    INSERT INTO public.knowledge_chunk_embeddings (
                        id, chunk_id, profile_id, embedding, embedding_hash
                    )
                    SELECT CAST(:id AS uuid), CAST(:chunk_id AS uuid), profile_id,
                           embedding, embedding_hash
                    FROM public.knowledge_chunk_embeddings
                    WHERE profile_id = CAST(:profile_id AS uuid)
                      AND embedding_hash = :embedding_hash
                    LIMIT 1
                    """
                ),
                {
                    "id": new_governance_uid(),
                    "chunk_id": chunk_id,
                    "profile_id": self._embedding_profile_id,
                    "embedding_hash": embedding_hash,
                },
            ).rowcount
            if inserted != 1:
                raise RuntimeError("reusable embedding disappeared during publication")

    def _insert_change_items(
        self,
        *,
        change_set_id: str,
        publication: PreparedPublication,
        old_revisions: Mapping[str, int],
        new_revisions: Mapping[str, int],
    ) -> None:
        statement = text(
            """
            INSERT INTO public.knowledge_change_items (
                change_set_id, point_key, change_kind, old_point_revision,
                new_point_revision, old_content_hash, new_content_hash,
                canonical_status, embedding_status, cache_status, lightrag_status
            ) VALUES (
                CAST(:change_set_id AS uuid), :point_key, :change_kind,
                :old_revision, :new_revision, :old_hash, :new_hash,
                'complete', 'complete', 'complete', 'pending'
            )
            """
        )
        for change in (
            *publication.diff.added,
            *publication.diff.modified,
            *publication.diff.deleted,
        ):
            self._session.execute(
                statement,
                {
                    "change_set_id": change_set_id,
                    "point_key": change.point_key,
                    "change_kind": change.change_kind,
                    "old_revision": old_revisions.get(change.point_key),
                    "new_revision": new_revisions.get(change.point_key),
                    "old_hash": change.old.content_hash if change.old else None,
                    "new_hash": change.new.content_hash if change.new else None,
                },
            )

    def _invalidate_cache_dependencies(self, publication: PreparedPublication) -> None:
        changed_keys = sorted(publication.diff.changed_point_keys)
        if not changed_keys:
            return
        self._session.execute(
            text(
                """
                DELETE FROM public.knowledge_cache_dependencies
                WHERE point_key = ANY(CAST(:point_keys AS text[]))
                """
            ),
            {"point_keys": changed_keys},
        )

    def _insert_projections(
        self,
        document_id: str,
        change_set_id: str,
        snapshot: KnowledgeSnapshot,
    ) -> None:
        statement = text(
            """
            INSERT INTO public.knowledge_index_projections (
                id, document_id, change_set_id, engine, generation, workspace,
                external_document_id, content_hash, status
            ) VALUES (
                CAST(:id AS uuid), CAST(:document_id AS uuid),
                CAST(:change_set_id AS uuid), :engine, :generation, :workspace,
                :external_document_id, :content_hash, :status
            )
            """
        )
        for engine, status in (("canonical_vector", "ready"), ("lightrag", "pending")):
            self._session.execute(
                statement,
                {
                    "id": new_governance_uid(),
                    "document_id": document_id,
                    "change_set_id": change_set_id,
                    "engine": engine,
                    "generation": self._generation,
                    "workspace": self._workspace,
                    "external_document_id": (
                        f"{snapshot.source_type}:{snapshot.source_uid}:"
                        f"{snapshot.source_revision}"
                    ),
                    "content_hash": snapshot.source_snapshot_hash,
                    "status": status,
                },
            )