core-hub/src/core_hub/migration.py

import hashlib
import json
from collections.abc import Iterable
from datetime import datetime, timezone
from typing import Any

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from .models import (
    ApiConsumer,
    ApiKey,
    Hub,
    HubCapabilityManifest,
    InteractionEvent,
    MigrationRun,
    Widget,
    uuid_str,
)

MIGRATION_SCHEMA_VERSION = "core-hub.migration.v1"
COLLECTIONS = (
    "hubs",
    "hubCapabilityManifests",
    "apiConsumers",
    "apiKeys",
    "widgets",
    "interactionEvents",
)
SECRET_FIELD_NAMES = {"fullKey", "rawKey", "apiKey", "secret", "token"}


def encode_body(payload: dict[str, Any]) -> str:
    return json.dumps(payload, sort_keys=True)


def bundle_digest(bundle: dict[str, Any]) -> str:
    digestable = {key: value for key, value in bundle.items() if key != "bundleSha256"}
    canonical = json.dumps(digestable, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(canonical.encode("utf-8")).hexdigest()


def _timestamp(value: datetime) -> str:
    if value.tzinfo is None:
        value = value.replace(tzinfo=timezone.utc)
    return value.astimezone(timezone.utc).isoformat().replace("+00:00", "Z")


def _decoded(raw: str | None) -> dict[str, Any]:
    value = json.loads(raw or "{}")
    return value if isinstance(value, dict) else {}


def _with_timestamps(
    body: str | None,
    *,
    created_at: datetime,
    updated_at: datetime | None = None,
) -> dict[str, Any]:
    value = _decoded(body)
    value["createdAt"] = _timestamp(created_at)
    if updated_at is not None:
        value["updatedAt"] = _timestamp(updated_at)
    return value


async def export_bundle(
    session: AsyncSession, *, source_revision: str | None = None
) -> dict[str, Any]:
    """Export all seven durable Core Hub tables without raw credential material."""
    hubs = list((await session.execute(select(Hub).order_by(Hub.id))).scalars())
    manifests = list(
        (await session.execute(select(HubCapabilityManifest).order_by(HubCapabilityManifest.id)))
        .scalars()
    )
    consumers = list(
        (await session.execute(select(ApiConsumer).order_by(ApiConsumer.id))).scalars()
    )
    keys = list((await session.execute(select(ApiKey).order_by(ApiKey.id))).scalars())
    widgets = list((await session.execute(select(Widget).order_by(Widget.id))).scalars())
    events = list(
        (await session.execute(select(InteractionEvent).order_by(InteractionEvent.id))).scalars()
    )
    runs = list(
        (await session.execute(select(MigrationRun).order_by(MigrationRun.id))).scalars()
    )

    records: dict[str, list[dict[str, Any]]] = {
        "hubs": [],
        "hubCapabilityManifests": [],
        "apiConsumers": [],
        "apiKeys": [],
        "widgets": [],
        "interactionEvents": [],
        "migrationRuns": [],
    }
    for row in hubs:
        value = _with_timestamps(
            row.body_json, created_at=row.created_at, updated_at=row.updated_at
        )
        value.update(
            id=row.id,
            slug=row.slug,
            name=row.name,
            domain=row.domain,
            hubKind=row.hub_kind,
            hubFamily=row.hub_family,
            vsmFunction=row.vsm_function,
            vsmSystem=row.vsm_system,
            status=row.status,
            description=row.description,
        )
        records["hubs"].append({key: child for key, child in value.items() if child is not None})
    for row in manifests:
        value = _with_timestamps(
            row.body_json, created_at=row.created_at, updated_at=row.updated_at
        )
        value.update(
            id=row.id,
            hubId=row.hub_id,
            hubSlug=row.hub_slug,
            manifestVersion=row.manifest_version,
            status=row.status,
        )
        records["hubCapabilityManifests"].append(
            {key: child for key, child in value.items() if child is not None}
        )
    for row in consumers:
        value = _with_timestamps(row.body_json, created_at=row.created_at)
        value.update(
            id=row.id,
            slug=row.slug,
            name=row.name,
            description=row.description,
            hubCapabilityManifestId=row.hub_capability_manifest_id,
            rateLimitPerMinute=row.rate_limit_per_minute,
            quotaPerDay=row.quota_per_day,
            keyPrefix=row.key_prefix,
            status=row.status,
        )
        records["apiConsumers"].append(
            {key: child for key, child in value.items() if child is not None}
        )
    records["apiKeys"] = [
        {
            "id": row.id,
            "apiConsumerId": row.api_consumer_id,
            "keyPrefix": row.key_prefix,
            "keyHash": row.key_hash,
            "scopes": row.scopes,
            "status": row.status,
            "createdAt": _timestamp(row.created_at),
        }
        for row in keys
    ]
    for row in widgets:
        value = _with_timestamps(row.body_json, created_at=row.created_at)
        value.update(
            id=row.id,
            hubId=row.hub_id,
            name=row.name,
            widgetType=row.widget_type,
            capabilityRef=row.capability_ref,
            viewContext=row.view_context,
            policyScope=row.policy_scope,
            status=row.status,
        )
        records["widgets"].append(
            {key: child for key, child in value.items() if child is not None}
        )
    for row in events:
        value = _with_timestamps(row.body_json, created_at=row.created_at)
        value.update(
            id=row.id,
            widgetId=row.widget_id,
            eventType=row.event_type,
            viewContext=row.view_context,
            metadata=_decoded(row.metadata_json),
        )
        records["interactionEvents"].append(
            {key: child for key, child in value.items() if child is not None}
        )
    records["migrationRuns"] = [
        {
            "id": row.id,
            "source": row.source,
            "schemaVersion": row.schema_version,
            "bundleSha256": row.bundle_sha256,
            "dryRun": row.dry_run,
            "status": row.status,
            "counts": _decoded(row.counts_json),
            "diagnostics": _decoded(row.diagnostics_json),
            "createdAt": _timestamp(row.created_at),
        }
        for row in runs
    ]
    high_water = max((_timestamp(row.created_at) for row in events), default=None)
    bundle: dict[str, Any] = {
        "schemaVersion": MIGRATION_SCHEMA_VERSION,
        "source": "core-hub",
        "exportedAt": _timestamp(datetime.now(timezone.utc)),
        "highWaterMark": high_water,
        "records": records,
    }
    if source_revision:
        bundle["sourceRevision"] = source_revision
    bundle["bundleSha256"] = bundle_digest(bundle)
    return bundle


def records_for(bundle: dict[str, Any], collection: str) -> list[dict[str, Any]]:
    records = bundle.get("records", {})
    if not isinstance(records, dict):
        return []
    values = records.get(collection, [])
    if not isinstance(values, list):
        return []
    return [record for record in values if isinstance(record, dict)]


def _input_counts(bundle: dict[str, Any]) -> dict[str, dict[str, int]]:
    return {
        collection: {"input": len(records_for(bundle, collection))} for collection in COLLECTIONS
    }


def _work_counts(bundle: dict[str, Any]) -> dict[str, dict[str, int]]:
    return {
        collection: {
            "input": len(records_for(bundle, collection)),
            "created": 0,
            "updated": 0,
            "skipped": 0,
        }
        for collection in COLLECTIONS
    }


def _require(record: dict[str, Any], fields: Iterable[str], path: str, errors: list[str]) -> None:
    for field in fields:
        if record.get(field) in (None, ""):
            errors.append(f"{path}.{field} is required")


def _note_duplicate(value: str | None, seen: set[str], path: str, errors: list[str]) -> None:
    if not value:
        return
    if value in seen:
        errors.append(f"duplicate {path}: {value}")
    seen.add(value)


def _contains_secret_field(value: Any) -> bool:
    if isinstance(value, dict):
        for key, child in value.items():
            if key in SECRET_FIELD_NAMES:
                return True
            if _contains_secret_field(child):
                return True
    if isinstance(value, list):
        return any(_contains_secret_field(child) for child in value)
    return False


def _relationship(status: str, name: str, message: str) -> dict[str, str]:
    return {"name": name, "status": status, "message": message}


def validate_bundle(bundle: dict[str, Any]) -> dict[str, Any]:
    errors: list[str] = []
    warnings: list[str] = []
    relationships: list[dict[str, str]] = []

    if bundle.get("schemaVersion") != MIGRATION_SCHEMA_VERSION:
        errors.append(
            f"schemaVersion must be {MIGRATION_SCHEMA_VERSION}, got {bundle.get('schemaVersion')!r}"
        )
    if not isinstance(bundle.get("records", {}), dict):
        errors.append("records must be an object")

    hub_ids: set[str] = set()
    hub_slugs: set[str] = set()
    manifest_ids: set[str] = set()
    consumer_ids: set[str] = set()
    consumer_slugs: set[str] = set()
    widget_ids: set[str] = set()

    for index, record in enumerate(records_for(bundle, "hubs")):
        path = f"records.hubs[{index}]"
        _require(record, ("slug", "name"), path, errors)
        _note_duplicate(record.get("id"), hub_ids, "hub id", errors)
        _note_duplicate(record.get("slug"), hub_slugs, "hub slug", errors)

    for index, record in enumerate(records_for(bundle, "hubCapabilityManifests")):
        path = f"records.hubCapabilityManifests[{index}]"
        _require(record, ("id",), path, errors)
        _note_duplicate(record.get("id"), manifest_ids, "manifest id", errors)
        if not record.get("hubId") and not record.get("hubSlug"):
            errors.append(f"{path} must include hubId or hubSlug")
        elif record.get("hubId") not in hub_ids and record.get("hubSlug") not in hub_slugs:
            warnings.append(f"{path} references a hub outside this bundle")

    for index, record in enumerate(records_for(bundle, "apiConsumers")):
        path = f"records.apiConsumers[{index}]"
        _require(record, ("name",), path, errors)
        if not record.get("id") and not record.get("slug"):
            errors.append(f"{path} must include id or slug for idempotent import")
        _note_duplicate(record.get("id"), consumer_ids, "api consumer id", errors)
        _note_duplicate(record.get("slug"), consumer_slugs, "api consumer slug", errors)
        manifest_id = record.get("hubCapabilityManifestId")
        if manifest_id and manifest_id not in manifest_ids:
            warnings.append(
                f"{path}.hubCapabilityManifestId references a manifest outside this bundle"
            )

    for index, record in enumerate(records_for(bundle, "apiKeys")):
        path = f"records.apiKeys[{index}]"
        _require(record, ("id", "keyPrefix", "keyHash"), path, errors)
        if _contains_secret_field(record):
            errors.append(
                f"{path} contains a secret-like key field; import only hashes and prefixes"
            )
        if not record.get("apiConsumerId") and not record.get("apiConsumerSlug"):
            errors.append(f"{path} must include apiConsumerId or apiConsumerSlug")
        elif (
            record.get("apiConsumerId") not in consumer_ids
            and record.get("apiConsumerSlug") not in consumer_slugs
        ):
            warnings.append(f"{path} references an API consumer outside this bundle")

    for index, record in enumerate(records_for(bundle, "widgets")):
        path = f"records.widgets[{index}]"
        _require(record, ("id", "name"), path, errors)
        _note_duplicate(record.get("id"), widget_ids, "widget id", errors)
        if not record.get("hubId") and not record.get("hubSlug"):
            errors.append(f"{path} must include hubId or hubSlug")
        elif record.get("hubId") not in hub_ids and record.get("hubSlug") not in hub_slugs:
            warnings.append(f"{path} references a hub outside this bundle")

    for index, record in enumerate(records_for(bundle, "interactionEvents")):
        path = f"records.interactionEvents[{index}]"
        _require(record, ("id", "eventType"), path, errors)
        if not record.get("widgetId"):
            errors.append(f"{path}.widgetId is required")
        elif record.get("widgetId") not in widget_ids:
            warnings.append(f"{path}.widgetId references a widget outside this bundle")

    if errors:
        relationships.append(_relationship("fail", "bundle_integrity", "validation errors present"))
    else:
        relationships.append(
            _relationship("pass", "bundle_integrity", "required fields are present")
        )
    if warnings:
        relationships.append(
            _relationship(
                "warn", "external_references", "some references must resolve in the target DB"
            )
        )
    else:
        relationships.append(
            _relationship("pass", "external_references", "all references resolve within the bundle")
        )

    return {
        "ok": not errors,
        "schemaVersion": bundle.get("schemaVersion"),
        "source": bundle.get("source", "unknown"),
        "bundleSha256": bundle_digest(bundle),
        "counts": _input_counts(bundle),
        "errors": errors,
        "warnings": warnings,
        "relationshipChecks": relationships,
    }


async def _one_or_none(session: AsyncSession, statement: Any) -> Any | None:
    result = await session.execute(statement)
    return result.scalars().first()


async def _hub_maps(session: AsyncSession) -> tuple[dict[str, Hub], dict[str, Hub]]:
    result = await session.execute(select(Hub))
    rows = list(result.scalars())
    return {row.id: row for row in rows}, {row.slug: row for row in rows}


async def _manifest_maps(session: AsyncSession) -> dict[str, HubCapabilityManifest]:
    result = await session.execute(select(HubCapabilityManifest))
    rows = list(result.scalars())
    return {row.id: row for row in rows}


async def _consumer_maps(
    session: AsyncSession,
) -> tuple[dict[str, ApiConsumer], dict[str, ApiConsumer]]:
    result = await session.execute(select(ApiConsumer))
    rows = list(result.scalars())
    return {row.id: row for row in rows}, {row.slug: row for row in rows if row.slug}


async def _widget_maps(session: AsyncSession) -> dict[str, Widget]:
    result = await session.execute(select(Widget))
    rows = list(result.scalars())
    return {row.id: row for row in rows}


def _resolved_hub_id(
    record: dict[str, Any], hubs_by_id: dict[str, Hub], hubs_by_slug: dict[str, Hub]
) -> str:
    if record.get("hubId"):
        return str(record["hubId"])
    return hubs_by_slug[str(record["hubSlug"])].id


def _resolved_consumer_id(record: dict[str, Any], consumers_by_slug: dict[str, ApiConsumer]) -> str:
    if record.get("apiConsumerId"):
        return str(record["apiConsumerId"])
    return consumers_by_slug[str(record["apiConsumerSlug"])].id


def _mark(counts: dict[str, dict[str, int]], collection: str, existing: Any | None) -> None:
    key = "updated" if existing else "created"
    counts[collection][key] += 1


async def import_bundle(
    session: AsyncSession, bundle: dict[str, Any], *, dry_run: bool = False
) -> dict[str, Any]:
    report = validate_bundle(bundle)
    report["dryRun"] = dry_run
    report["counts"] = _work_counts(bundle)
    if not report["ok"]:
        return report

    hubs_by_id, hubs_by_slug = await _hub_maps(session)
    for record in records_for(bundle, "hubs"):
        existing = hubs_by_id.get(record.get("id")) or hubs_by_slug.get(record.get("slug"))
        _mark(report["counts"], "hubs", existing)
        if dry_run:
            continue
        row = existing or Hub(
            id=record.get("id") or uuid_str(),
            slug=record["slug"],
            name=record["name"],
        )
        row.slug = record["slug"]
        row.name = record["name"]
        row.domain = record.get("domain")
        row.hub_kind = record.get("hubKind")
        row.hub_family = record.get("hubFamily")
        row.vsm_function = record.get("vsmFunction")
        row.vsm_system = record.get("vsmSystem")
        row.status = record.get("status", "active")
        row.description = record.get("description")
        row.body_json = encode_body(record)
        session.add(row)
    if not dry_run:
        await session.flush()
        hubs_by_id, hubs_by_slug = await _hub_maps(session)

    manifests_by_id = await _manifest_maps(session)
    for record in records_for(bundle, "hubCapabilityManifests"):
        existing = manifests_by_id.get(record.get("id"))
        _mark(report["counts"], "hubCapabilityManifests", existing)
        if dry_run:
            continue
        hub_id = _resolved_hub_id(record, hubs_by_id, hubs_by_slug)
        hub_slug = hubs_by_id[hub_id].slug if hub_id in hubs_by_id else record.get("hubSlug")
        row = existing or HubCapabilityManifest(id=record["id"], body_json="{}")
        row.hub_id = hub_id
        row.hub_slug = hub_slug
        row.manifest_version = record.get("manifestVersion", "0.1.0")
        row.status = record.get("status", "draft")
        row.body_json = encode_body(record)
        session.add(row)
    if not dry_run:
        await session.flush()

    consumers_by_id, consumers_by_slug = await _consumer_maps(session)
    for record in records_for(bundle, "apiConsumers"):
        existing = consumers_by_id.get(record.get("id")) or consumers_by_slug.get(
            record.get("slug")
        )
        _mark(report["counts"], "apiConsumers", existing)
        if dry_run:
            continue
        row = existing or ApiConsumer(
            id=record.get("id") or uuid_str(),
            slug=record.get("slug"),
            name=record["name"],
        )
        row.slug = record.get("slug")
        row.name = record["name"]
        row.description = record.get("description")
        row.hub_capability_manifest_id = record.get("hubCapabilityManifestId")
        row.rate_limit_per_minute = record.get("rateLimitPerMinute")
        row.quota_per_day = record.get("quotaPerDay")
        row.key_prefix = record.get("keyPrefix")
        row.key_hash = record.get("keyHash")
        row.status = record.get("status", "active")
        row.body_json = encode_body(record)
        session.add(row)
    if not dry_run:
        await session.flush()
        consumers_by_id, consumers_by_slug = await _consumer_maps(session)

    for record in records_for(bundle, "apiKeys"):
        existing = await session.get(ApiKey, record.get("id"))
        if existing is None:
            existing = await _one_or_none(
                session, select(ApiKey).where(ApiKey.key_hash == record["keyHash"])
            )
        _mark(report["counts"], "apiKeys", existing)
        if dry_run:
            continue
        row = existing or ApiKey(
            id=record["id"],
            api_consumer_id=_resolved_consumer_id(record, consumers_by_slug),
        )
        row.api_consumer_id = _resolved_consumer_id(record, consumers_by_slug)
        row.key_prefix = record["keyPrefix"]
        row.key_hash = record["keyHash"]
        row.scopes = record.get("scopes")
        row.status = record.get("status", "active")
        session.add(row)

    widgets_by_id = await _widget_maps(session)
    for record in records_for(bundle, "widgets"):
        existing = widgets_by_id.get(record.get("id"))
        _mark(report["counts"], "widgets", existing)
        if dry_run:
            continue
        row = existing or Widget(
            id=record["id"],
            hub_id=_resolved_hub_id(record, hubs_by_id, hubs_by_slug),
            name=record["name"],
        )
        row.hub_id = _resolved_hub_id(record, hubs_by_id, hubs_by_slug)
        row.name = record["name"]
        row.widget_type = record.get("widgetType")
        row.capability_ref = record.get("capabilityRef")
        row.view_context = record.get("viewContext")
        row.policy_scope = record.get("policyScope")
        row.status = record.get("status", "active")
        row.body_json = encode_body(record)
        session.add(row)
    if not dry_run:
        await session.flush()

    for record in records_for(bundle, "interactionEvents"):
        existing = await session.get(InteractionEvent, record.get("id"))
        _mark(report["counts"], "interactionEvents", existing)
        if dry_run:
            continue
        row = existing or InteractionEvent(
            id=record["id"],
            widget_id=record["widgetId"],
            event_type=record["eventType"],
        )
        row.widget_id = record["widgetId"]
        row.event_type = record["eventType"]
        row.view_context = record.get("viewContext")
        row.metadata_json = encode_body(record.get("metadata", {}))
        row.body_json = encode_body(record)
        session.add(row)

    if not dry_run:
        run = MigrationRun(
            source=report["source"],
            schema_version=report["schemaVersion"],
            bundle_sha256=report["bundleSha256"],
            dry_run=False,
            status="imported",
            counts_json=encode_body(report["counts"]),
            diagnostics_json=encode_body(
                {
                    "errors": report["errors"],
                    "warnings": report["warnings"],
                    "relationshipChecks": report["relationshipChecks"],
                }
            ),
        )
        session.add(run)
        await session.commit()
        await session.refresh(run)
        report["migrationRunId"] = run.id

    return report