kontextual-engine/src/kontextual_engine/adapters/sqlite/asset_registry.py

"""SQLite asset registry repository."""

from __future__ import annotations

import json
import sqlite3
from pathlib import Path
from typing import Any

from kontextual_engine.core import (
    Actor,
    AssetRepresentation,
    AssetVersion,
    AuditEvent,
    ContextEntity,
    CoreRelationship,
    IdempotencyRecord,
    IngestionJob,
    IngestionJobStatus,
    KnowledgeAsset,
    LifecycleState,
    MetadataRecord,
    RepresentationKind,
    RelationshipTargetKind,
)
from kontextual_engine.errors import NotFoundError, ValidationError


class SQLiteAssetRegistryRepository:
    def __init__(self, path: str | Path) -> None:
        self.path = Path(path)
        self.path.parent.mkdir(parents=True, exist_ok=True)
        self._initialize()

    def save_actor(self, actor: Actor) -> Actor:
        with self._connect() as conn:
            conn.execute(
                """
                insert into actors (id, actor_type, payload)
                values (?, ?, ?)
                on conflict(id) do update set
                  actor_type=excluded.actor_type,
                  payload=excluded.payload
                """,
                (actor.id, actor.actor_type.value, _json(actor.to_dict())),
            )
        return actor

    def get_actor(self, actor_id: str) -> Actor:
        row = self._one("select payload from actors where id = ?", (actor_id,))
        if row is None:
            raise NotFoundError("Actor not found", details={"actor_id": actor_id})
        return Actor.from_dict(_loads(row["payload"]))

    def save_asset(self, asset: KnowledgeAsset) -> KnowledgeAsset:
        with self._connect() as conn:
            conn.execute(
                """
                insert into assets (id, title, asset_type, lifecycle, payload)
                values (?, ?, ?, ?, ?)
                on conflict(id) do update set
                  title=excluded.title,
                  asset_type=excluded.asset_type,
                  lifecycle=excluded.lifecycle,
                  payload=excluded.payload
                """,
                (
                    asset.id,
                    asset.title,
                    asset.classification.asset_type,
                    asset.lifecycle.value,
                    _json(asset.to_dict()),
                ),
            )
        return asset

    def get_asset(self, asset_id: str) -> KnowledgeAsset:
        row = self._one("select payload from assets where id = ?", (asset_id,))
        if row is None:
            raise NotFoundError("Asset not found", details={"asset_id": asset_id})
        return KnowledgeAsset.from_dict(_loads(row["payload"]))

    def list_assets(
        self,
        *,
        lifecycle: LifecycleState | None = None,
        asset_type: str | None = None,
    ) -> list[KnowledgeAsset]:
        clauses = []
        params: list[Any] = []
        if lifecycle is not None:
            clauses.append("lifecycle = ?")
            params.append(lifecycle.value)
        if asset_type is not None:
            clauses.append("asset_type = ?")
            params.append(asset_type)
        where = f" where {' and '.join(clauses)}" if clauses else ""
        rows = self._all(f"select payload from assets{where} order by title, id", tuple(params))
        return [KnowledgeAsset.from_dict(_loads(row["payload"])) for row in rows]

    def save_representation(self, representation: AssetRepresentation) -> AssetRepresentation:
        try:
            with self._connect() as conn:
                conn.execute(
                    """
                    insert into representations (id, asset_id, kind, digest, payload)
                    values (?, ?, ?, ?, ?)
                    on conflict(id) do update set
                      asset_id=excluded.asset_id,
                      kind=excluded.kind,
                      digest=excluded.digest,
                      payload=excluded.payload
                    """,
                    (
                        representation.representation_id,
                        representation.asset_id,
                        representation.kind.value,
                        representation.digest,
                        _json(representation.to_dict()),
                    ),
                )
        except sqlite3.IntegrityError as exc:
            raise ValidationError(
                "Representation references an unknown asset",
                details={
                    "asset_id": representation.asset_id,
                    "representation_id": representation.representation_id,
                },
            ) from exc
        return representation

    def get_representation(self, representation_id: str) -> AssetRepresentation:
        row = self._one("select payload from representations where id = ?", (representation_id,))
        if row is None:
            raise NotFoundError(
                "Representation not found",
                details={"representation_id": representation_id},
            )
        return AssetRepresentation.from_dict(_loads(row["payload"]))

    def list_representations(
        self,
        *,
        asset_id: str | None = None,
        kind: RepresentationKind | None = None,
    ) -> list[AssetRepresentation]:
        clauses = []
        params: list[Any] = []
        if asset_id is not None:
            clauses.append("asset_id = ?")
            params.append(asset_id)
        if kind is not None:
            clauses.append("kind = ?")
            params.append(kind.value)
        where = f" where {' and '.join(clauses)}" if clauses else ""
        rows = self._all(
            f"select payload from representations{where} order by asset_id, kind, id",
            tuple(params),
        )
        return [AssetRepresentation.from_dict(_loads(row["payload"])) for row in rows]

    def save_metadata_record(self, asset_id: str, record: MetadataRecord) -> MetadataRecord:
        try:
            with self._connect() as conn:
                conn.execute(
                    """
                    insert into metadata_records (id, asset_id, key, payload)
                    values (?, ?, ?, ?)
                    on conflict(id) do update set
                      asset_id=excluded.asset_id,
                      key=excluded.key,
                      payload=excluded.payload
                    """,
                    (record.record_id, asset_id, record.key, _json(record.to_dict())),
                )
        except sqlite3.IntegrityError as exc:
            raise ValidationError(
                "Metadata record references an unknown asset",
                details={"asset_id": asset_id, "record_id": record.record_id},
            ) from exc
        return record

    def list_metadata_records(self, asset_id: str) -> list[MetadataRecord]:
        rows = self._all(
            "select payload from metadata_records where asset_id = ? order by key, id",
            (asset_id,),
        )
        if not rows:
            self.get_asset(asset_id)
        return [MetadataRecord.from_dict(_loads(row["payload"])) for row in rows]

    def save_context_entity(self, entity: ContextEntity) -> ContextEntity:
        with self._connect() as conn:
            conn.execute(
                """
                insert into context_entities (id, entity_type, name, payload)
                values (?, ?, ?, ?)
                on conflict(id) do update set
                  entity_type=excluded.entity_type,
                  name=excluded.name,
                  payload=excluded.payload
                """,
                (entity.entity_id, entity.entity_type.value, entity.name, _json(entity.to_dict())),
            )
        return entity

    def get_context_entity(self, entity_id: str) -> ContextEntity:
        row = self._one("select payload from context_entities where id = ?", (entity_id,))
        if row is None:
            raise NotFoundError("Context entity not found", details={"entity_id": entity_id})
        return ContextEntity.from_dict(_loads(row["payload"]))

    def list_context_entities(self) -> list[ContextEntity]:
        rows = self._all("select payload from context_entities order by entity_type, name, id", ())
        return [ContextEntity.from_dict(_loads(row["payload"])) for row in rows]

    def save_relationship(self, relationship: CoreRelationship) -> CoreRelationship:
        self.get_asset(relationship.source_id)
        if relationship.target_kind == RelationshipTargetKind.ASSET:
            self.get_asset(relationship.target_id)
        else:
            self.get_context_entity(relationship.target_id)
        with self._connect() as conn:
            conn.execute(
                """
                insert into core_relationships (id, source_id, target_id, target_kind, predicate, payload)
                values (?, ?, ?, ?, ?, ?)
                on conflict(id) do update set
                  source_id=excluded.source_id,
                  target_id=excluded.target_id,
                  target_kind=excluded.target_kind,
                  predicate=excluded.predicate,
                  payload=excluded.payload
                """,
                (
                    relationship.relationship_id,
                    relationship.source_id,
                    relationship.target_id,
                    relationship.target_kind.value,
                    relationship.predicate,
                    _json(relationship.to_dict()),
                ),
            )
        return relationship

    def get_relationship(self, relationship_id: str) -> CoreRelationship:
        row = self._one("select payload from core_relationships where id = ?", (relationship_id,))
        if row is None:
            raise NotFoundError(
                "Relationship not found",
                details={"relationship_id": relationship_id},
            )
        return CoreRelationship.from_dict(_loads(row["payload"]))

    def list_relationships(
        self,
        *,
        source_id: str | None = None,
        target_id: str | None = None,
    ) -> list[CoreRelationship]:
        clauses = []
        params: list[Any] = []
        if source_id is not None:
            clauses.append("source_id = ?")
            params.append(source_id)
        if target_id is not None:
            clauses.append("target_id = ?")
            params.append(target_id)
        where = f" where {' and '.join(clauses)}" if clauses else ""
        rows = self._all(
            f"select payload from core_relationships{where} order by source_id, target_id, predicate, id",
            tuple(params),
        )
        return [CoreRelationship.from_dict(_loads(row["payload"])) for row in rows]

    def save_version(self, version: AssetVersion) -> AssetVersion:
        try:
            with self._connect() as conn:
                conn.execute(
                    """
                    insert into asset_versions (id, asset_id, sequence, change_type, payload)
                    values (?, ?, ?, ?, ?)
                    """,
                    (
                        version.version_id,
                        version.asset_id,
                        version.sequence,
                        version.change_type.value,
                        _json(version.to_dict()),
                    ),
                )
        except sqlite3.IntegrityError as exc:
            raise ValidationError(
                "Version sequence already exists for asset",
                details={"asset_id": version.asset_id, "sequence": version.sequence},
            ) from exc
        return version

    def list_versions(self, asset_id: str) -> list[AssetVersion]:
        rows = self._all(
            "select payload from asset_versions where asset_id = ? order by sequence",
            (asset_id,),
        )
        if not rows:
            self.get_asset(asset_id)
        return [AssetVersion.from_dict(_loads(row["payload"])) for row in rows]

    def save_audit_event(self, event: AuditEvent) -> AuditEvent:
        with self._connect() as conn:
            conn.execute(
                """
                insert into audit_events (id, target, actor_id, correlation_id, outcome, occurred_at, payload)
                values (?, ?, ?, ?, ?, ?, ?)
                on conflict(id) do update set
                  target=excluded.target,
                  actor_id=excluded.actor_id,
                  correlation_id=excluded.correlation_id,
                  outcome=excluded.outcome,
                  occurred_at=excluded.occurred_at,
                  payload=excluded.payload
                """,
                (
                    event.event_id,
                    event.target,
                    event.actor_id,
                    event.correlation_id,
                    event.outcome.value,
                    event.occurred_at,
                    _json(event.to_dict()),
                ),
            )
        return event

    def get_audit_event(self, event_id: str) -> AuditEvent:
        row = self._one("select payload from audit_events where id = ?", (event_id,))
        if row is None:
            raise NotFoundError("Audit event not found", details={"event_id": event_id})
        return AuditEvent.from_dict(_loads(row["payload"]))

    def list_audit_events(
        self,
        *,
        target: str | None = None,
        correlation_id: str | None = None,
    ) -> list[AuditEvent]:
        clauses = []
        params: list[Any] = []
        if target is not None:
            clauses.append("target = ?")
            params.append(target)
        if correlation_id is not None:
            clauses.append("correlation_id = ?")
            params.append(correlation_id)
        where = f" where {' and '.join(clauses)}" if clauses else ""
        rows = self._all(f"select payload from audit_events{where} order by occurred_at, id", tuple(params))
        return [AuditEvent.from_dict(_loads(row["payload"])) for row in rows]

    def save_idempotency_record(self, record: IdempotencyRecord) -> IdempotencyRecord:
        with self._connect() as conn:
            conn.execute(
                """
                insert into idempotency_records (key, operation, request_hash, status, payload)
                values (?, ?, ?, ?, ?)
                on conflict(key) do update set
                  operation=excluded.operation,
                  request_hash=excluded.request_hash,
                  status=excluded.status,
                  payload=excluded.payload
                """,
                (
                    record.key,
                    record.operation,
                    record.request_hash,
                    record.status.value,
                    _json(record.to_dict()),
                ),
            )
        return record

    def get_idempotency_record(self, key: str) -> IdempotencyRecord | None:
        row = self._one("select payload from idempotency_records where key = ?", (key,))
        if row is None:
            return None
        return IdempotencyRecord.from_dict(_loads(row["payload"]))

    def save_ingestion_job(self, job: IngestionJob) -> IngestionJob:
        with self._connect() as conn:
            conn.execute(
                """
                insert into ingestion_jobs (id, status, actor_id, correlation_id, created_at, updated_at, payload)
                values (?, ?, ?, ?, ?, ?, ?)
                on conflict(id) do update set
                  status=excluded.status,
                  actor_id=excluded.actor_id,
                  correlation_id=excluded.correlation_id,
                  updated_at=excluded.updated_at,
                  payload=excluded.payload
                """,
                (
                    job.job_id,
                    job.status.value,
                    job.actor_id,
                    job.correlation_id,
                    job.created_at,
                    job.updated_at,
                    _json(job.to_dict()),
                ),
            )
        return job

    def get_ingestion_job(self, job_id: str) -> IngestionJob:
        row = self._one("select payload from ingestion_jobs where id = ?", (job_id,))
        if row is None:
            raise NotFoundError("Ingestion job not found", details={"job_id": job_id})
        return IngestionJob.from_dict(_loads(row["payload"]))

    def list_ingestion_jobs(
        self,
        *,
        status: IngestionJobStatus | None = None,
    ) -> list[IngestionJob]:
        if status is None:
            rows = self._all("select payload from ingestion_jobs order by created_at, id", ())
        else:
            rows = self._all(
                "select payload from ingestion_jobs where status = ? order by created_at, id",
                (status.value,),
            )
        return [IngestionJob.from_dict(_loads(row["payload"])) for row in rows]

    def _initialize(self) -> None:
        with self._connect() as conn:
            conn.executescript(
                """
                create table if not exists actors (
                  id text primary key,
                  actor_type text not null,
                  payload text not null
                );
                create table if not exists assets (
                  id text primary key,
                  title text not null,
                  asset_type text not null,
                  lifecycle text not null,
                  payload text not null
                );
                create table if not exists representations (
                  id text primary key,
                  asset_id text not null references assets(id) on delete cascade,
                  kind text not null,
                  digest text not null,
                  payload text not null
                );
                create table if not exists metadata_records (
                  id text primary key,
                  asset_id text not null references assets(id) on delete cascade,
                  key text not null,
                  payload text not null
                );
                create table if not exists context_entities (
                  id text primary key,
                  entity_type text not null,
                  name text not null,
                  payload text not null
                );
                create table if not exists core_relationships (
                  id text primary key,
                  source_id text not null references assets(id) on delete cascade,
                  target_id text not null,
                  target_kind text not null,
                  predicate text not null,
                  payload text not null
                );
                create table if not exists asset_versions (
                  id text primary key,
                  asset_id text not null references assets(id) on delete cascade,
                  sequence integer not null,
                  change_type text not null,
                  payload text not null,
                  unique(asset_id, sequence)
                );
                create table if not exists audit_events (
                  id text primary key,
                  target text not null,
                  actor_id text not null,
                  correlation_id text not null,
                  outcome text not null,
                  occurred_at text not null,
                  payload text not null,
                  foreign key(actor_id) references actors(id)
                );
                create table if not exists idempotency_records (
                  key text primary key,
                  operation text not null,
                  request_hash text not null,
                  status text not null,
                  payload text not null
                );
                create table if not exists ingestion_jobs (
                  id text primary key,
                  status text not null,
                  actor_id text not null,
                  correlation_id text not null,
                  created_at text not null,
                  updated_at text not null,
                  payload text not null
                );
                create index if not exists idx_assets_lifecycle on assets(lifecycle);
                create index if not exists idx_representations_asset on representations(asset_id);
                create index if not exists idx_metadata_asset on metadata_records(asset_id);
                create index if not exists idx_entities_type on context_entities(entity_type);
                create index if not exists idx_relationships_source on core_relationships(source_id);
                create index if not exists idx_relationships_target on core_relationships(target_id);
                create index if not exists idx_versions_asset on asset_versions(asset_id);
                create index if not exists idx_audit_target on audit_events(target);
                create index if not exists idx_audit_correlation on audit_events(correlation_id);
                create index if not exists idx_ingestion_jobs_status on ingestion_jobs(status);
                create index if not exists idx_ingestion_jobs_correlation on ingestion_jobs(correlation_id);
                """
            )

    def _connect(self) -> sqlite3.Connection:
        conn = sqlite3.connect(self.path)
        conn.row_factory = sqlite3.Row
        conn.execute("pragma foreign_keys = on")
        return conn

    def _one(self, query: str, params: tuple[Any, ...]) -> sqlite3.Row | None:
        with self._connect() as conn:
            return conn.execute(query, params).fetchone()

    def _all(self, query: str, params: tuple[Any, ...]) -> list[sqlite3.Row]:
        with self._connect() as conn:
            return list(conn.execute(query, params).fetchall())


def _json(value: dict[str, Any]) -> str:
    return json.dumps(value, sort_keys=True, separators=(",", ":"))


def _loads(value: str) -> dict[str, Any]:
    return json.loads(value)