"""Add governed model profiles, embedding cache, and invocation metadata.

Revision ID: 0002_model_profiles
Revises: 0001_initial_schema
Create Date: 2026-07-13
"""

from collections.abc import Sequence

from alembic import op

revision: str = "0002_model_profiles"
down_revision: str | None = "0001_initial_schema"
branch_labels: str | Sequence[str] | None = None
depends_on: str | Sequence[str] | None = None


def upgrade() -> None:
    op.execute(
        """
        CREATE TABLE rag.model_profiles (
          profile_hash char(64) PRIMARY KEY,
          alias text NOT NULL,
          kind text NOT NULL,
          provider text NOT NULL,
          model text NOT NULL,
          api_mode text NOT NULL,
          dimension smallint,
          endpoint_identity_hash char(64) NOT NULL,
          config_snapshot jsonb NOT NULL DEFAULT '{}'::jsonb,
          synthetic boolean NOT NULL DEFAULT false,
          enabled boolean NOT NULL DEFAULT true,
          created_at timestamptz NOT NULL DEFAULT now(),
          updated_at timestamptz NOT NULL DEFAULT now(),
          CONSTRAINT model_profiles_alias_key UNIQUE (alias),
          CONSTRAINT model_profiles_hash_kind_key UNIQUE (profile_hash, kind),
          CONSTRAINT model_profiles_hash_format
            CHECK (profile_hash ~ '^[0-9a-f]{64}$'),
          CONSTRAINT model_profiles_alias_nonempty
            CHECK (btrim(alias) <> ''),
          CONSTRAINT model_profiles_kind_valid
            CHECK (kind IN ('embedding', 'rerank', 'chat')),
          CONSTRAINT model_profiles_identity_nonempty
            CHECK (
              btrim(provider) <> ''
              AND btrim(model) <> ''
              AND btrim(api_mode) <> ''
            ),
          CONSTRAINT model_profiles_embedding_dimension
            CHECK (
              (kind = 'embedding' AND dimension = 1024)
              OR (kind IN ('rerank', 'chat') AND dimension IS NULL)
            ),
          CONSTRAINT model_profiles_endpoint_identity_hash_format
            CHECK (endpoint_identity_hash ~ '^[0-9a-f]{64}$'),
          CONSTRAINT model_profiles_config_snapshot_object
            CHECK (jsonb_typeof(config_snapshot) = 'object'),
          CONSTRAINT model_profiles_config_snapshot_has_no_credentials
            CHECK (
              config_snapshot::text !~*
                '"[^\"]*(api[_-]?key|secret|password|token|authorization|credential)[^\"]*"[[:space:]]*:'
            ),
          CONSTRAINT model_profiles_timestamps_valid
            CHECK (updated_at >= created_at)
        );
        """
    )

    op.execute(
        """
        ALTER TABLE rag.knowledge_bases
          ADD COLUMN active_embedding_profile_hash char(64),
          ADD COLUMN active_embedding_profile_kind text NOT NULL DEFAULT 'embedding',
          ADD CONSTRAINT knowledge_bases_active_embedding_profile_hash_format
            CHECK (
              active_embedding_profile_hash IS NULL
              OR active_embedding_profile_hash ~ '^[0-9a-f]{64}$'
            ),
          ADD CONSTRAINT knowledge_bases_active_embedding_profile_fk
            FOREIGN KEY (
              active_embedding_profile_hash,
              active_embedding_profile_kind
            )
            REFERENCES rag.model_profiles (profile_hash, kind)
            ON DELETE RESTRICT;
          ALTER TABLE rag.knowledge_bases
            ADD CONSTRAINT knowledge_bases_active_embedding_profile_kind
              CHECK (active_embedding_profile_kind = 'embedding');
        """
    )

    # The only profile that can be inferred safely from legacy rows is an explicitly
    # synthetic, searchable fake embedding profile with one unambiguous model name.
    # Live provider identity is never guessed from model names or endpoint values.
    op.execute(
        """
        INSERT INTO rag.model_profiles (
          profile_hash,
          alias,
          kind,
          provider,
          model,
          api_mode,
          dimension,
          endpoint_identity_hash,
          config_snapshot,
          synthetic,
          enabled
        )
        SELECT
          chunk.embedding_profile_hash,
          'fake-embedding-' || left(chunk.embedding_profile_hash, 12),
          'embedding',
          'local-synthetic',
          min(chunk.embedding_model),
          'deterministic-offline',
          1024,
          encode(sha256(convert_to('local-fake', 'UTF8')), 'hex'),
          jsonb_build_object(
            'migration_revision', '0002_model_profiles',
            'source', 'existing_searchable_fake_chunks'
          ),
          true,
          true
        FROM rag.chunks AS chunk
        WHERE chunk.searchable IS TRUE
          AND chunk.embedding_profile_hash ~ '^[0-9a-f]{64}$'
          AND chunk.embedding_dimension = 1024
          AND lower(chunk.embedding_model) LIKE 'fake-%'
        GROUP BY chunk.embedding_profile_hash
        HAVING count(DISTINCT chunk.embedding_model) = 1
        ON CONFLICT (profile_hash) DO NOTHING;
        """
    )

    # A knowledge base is activated only when its searchable legacy projection has
    # exactly one backfilled fake profile. Multiple profiles intentionally leave NULL.
    op.execute(
        """
        WITH unique_searchable_fake_profile AS (
          SELECT
            chunk.knowledge_base_id,
            min(chunk.embedding_profile_hash) AS profile_hash
          FROM rag.chunks AS chunk
          JOIN rag.model_profiles AS profile
            ON profile.profile_hash = chunk.embedding_profile_hash
           AND profile.kind = 'embedding'
           AND profile.synthetic IS TRUE
          WHERE chunk.searchable IS TRUE
            AND lower(chunk.embedding_model) LIKE 'fake-%'
          GROUP BY chunk.knowledge_base_id
          HAVING count(DISTINCT chunk.embedding_profile_hash) = 1
        )
        UPDATE rag.knowledge_bases AS knowledge_base
           SET active_embedding_profile_hash = candidate.profile_hash,
               updated_at = now()
          FROM unique_searchable_fake_profile AS candidate
         WHERE knowledge_base.id = candidate.knowledge_base_id
           AND knowledge_base.active_embedding_profile_hash IS NULL;
        """
    )

    op.execute(
        """
        ALTER TABLE rag.chunks
          ADD COLUMN citation_id uuid NOT NULL DEFAULT gen_random_uuid(),
          ADD CONSTRAINT chunks_citation_id_key UNIQUE (citation_id),
          ADD CONSTRAINT chunks_id_embedding_text_sha256_key
            UNIQUE (id, embedding_text_sha256);
        """
    )

    op.execute(
        """
        CREATE TABLE rag.embedding_cache (
          profile_hash char(64) NOT NULL,
          profile_kind text NOT NULL DEFAULT 'embedding',
          embedding_text_sha256 char(64) NOT NULL,
          embedding vector(1024) NOT NULL,
          resolved_model text NOT NULL,
          provider_request_id text,
          usage jsonb NOT NULL DEFAULT '{}'::jsonb,
          elapsed_ms integer NOT NULL,
          created_at timestamptz NOT NULL DEFAULT now(),
          updated_at timestamptz NOT NULL DEFAULT now(),
          CONSTRAINT embedding_cache_primary_key
            PRIMARY KEY (profile_hash, embedding_text_sha256),
          CONSTRAINT embedding_cache_profile_fk
            FOREIGN KEY (profile_hash, profile_kind)
            REFERENCES rag.model_profiles (profile_hash, kind)
            ON DELETE RESTRICT,
          CONSTRAINT embedding_cache_profile_kind
            CHECK (profile_kind = 'embedding'),
          CONSTRAINT embedding_cache_text_hash_format
            CHECK (embedding_text_sha256 ~ '^[0-9a-f]{64}$'),
          CONSTRAINT embedding_cache_vector_dimension
            CHECK (vector_dims(embedding) = 1024),
          CONSTRAINT embedding_cache_resolved_model_nonempty
            CHECK (btrim(resolved_model) <> ''),
          CONSTRAINT embedding_cache_request_id_valid
            CHECK (
              provider_request_id IS NULL
              OR (
                btrim(provider_request_id) <> ''
                AND length(provider_request_id) <= 512
              )
            ),
          CONSTRAINT embedding_cache_usage_object
            CHECK (jsonb_typeof(usage) = 'object'),
          CONSTRAINT embedding_cache_elapsed_valid
            CHECK (elapsed_ms >= 0),
          CONSTRAINT embedding_cache_timestamps_valid
            CHECK (updated_at >= created_at)
        );
        """
    )

    op.execute(
        """
        CREATE TABLE rag.chunk_embedding_assignments (
          chunk_id uuid NOT NULL,
          profile_hash char(64) NOT NULL,
          profile_kind text NOT NULL DEFAULT 'embedding',
          embedding_text_sha256 char(64) NOT NULL,
          cache_profile_hash char(64),
          cache_embedding_text_sha256 char(64),
          status text NOT NULL DEFAULT 'PENDING',
          error_code text,
          created_at timestamptz NOT NULL DEFAULT now(),
          updated_at timestamptz NOT NULL DEFAULT now(),
          completed_at timestamptz,
          CONSTRAINT chunk_embedding_assignments_primary_key
            PRIMARY KEY (chunk_id, profile_hash),
          CONSTRAINT chunk_embedding_assignments_chunk_text_fk
            FOREIGN KEY (chunk_id, embedding_text_sha256)
            REFERENCES rag.chunks (id, embedding_text_sha256)
            ON DELETE CASCADE,
          CONSTRAINT chunk_embedding_assignments_profile_fk
            FOREIGN KEY (profile_hash, profile_kind)
            REFERENCES rag.model_profiles (profile_hash, kind)
            ON DELETE RESTRICT,
          CONSTRAINT chunk_embedding_assignments_profile_kind
            CHECK (profile_kind = 'embedding'),
          CONSTRAINT chunk_embedding_assignments_cache_fk
            FOREIGN KEY (cache_profile_hash, cache_embedding_text_sha256)
            REFERENCES rag.embedding_cache (profile_hash, embedding_text_sha256)
            ON DELETE RESTRICT,
          CONSTRAINT chunk_embedding_assignments_text_hash_format
            CHECK (embedding_text_sha256 ~ '^[0-9a-f]{64}$'),
          CONSTRAINT chunk_embedding_assignments_status_valid
            CHECK (status IN ('PENDING', 'EMBEDDING', 'READY', 'FAILED', 'STALE')),
          CONSTRAINT chunk_embedding_assignments_cache_binding
            CHECK (
              (
                status = 'READY'
                AND cache_profile_hash = profile_hash
                AND cache_embedding_text_sha256 = embedding_text_sha256
              )
              OR (
                status <> 'READY'
                AND cache_profile_hash IS NULL
                AND cache_embedding_text_sha256 IS NULL
              )
            ),
          CONSTRAINT chunk_embedding_assignments_completion_consistent
            CHECK (
              (
                status IN ('READY', 'FAILED', 'STALE')
                AND completed_at IS NOT NULL
              )
              OR (
                status IN ('PENDING', 'EMBEDDING')
                AND completed_at IS NULL
              )
            ),
          CONSTRAINT chunk_embedding_assignments_error_code_valid
            CHECK (
              error_code IS NULL
              OR (btrim(error_code) <> '' AND length(error_code) <= 128)
            ),
          CONSTRAINT chunk_embedding_assignments_timestamps_valid
            CHECK (
              updated_at >= created_at
              AND (completed_at IS NULL OR completed_at >= created_at)
            )
        );
        """
    )

    op.execute(
        """
        CREATE INDEX chunk_embedding_assignments_work_queue
          ON rag.chunk_embedding_assignments (profile_hash, status, updated_at)
          WHERE status IN ('PENDING', 'EMBEDDING');
        """
    )

    op.execute(
        """
        CREATE TABLE rag.model_invocations (
          id uuid PRIMARY KEY DEFAULT gen_random_uuid(),
          trace_id uuid NOT NULL,
          caller text NOT NULL,
          operation text NOT NULL,
          profile_hash char(64) NOT NULL,
          model text NOT NULL,
          provider_request_id text,
          status text NOT NULL,
          item_count integer NOT NULL DEFAULT 0,
          prompt_tokens integer NOT NULL DEFAULT 0,
          completion_tokens integer NOT NULL DEFAULT 0,
          total_tokens integer NOT NULL DEFAULT 0,
          elapsed_ms integer,
          error_code text,
          started_at timestamptz NOT NULL DEFAULT now(),
          finished_at timestamptz,
          created_at timestamptz NOT NULL DEFAULT now(),
          CONSTRAINT model_invocations_profile_fk
            FOREIGN KEY (profile_hash, operation)
            REFERENCES rag.model_profiles (profile_hash, kind)
            ON DELETE RESTRICT,
          CONSTRAINT model_invocations_caller_nonempty
            CHECK (btrim(caller) <> ''),
          CONSTRAINT model_invocations_operation_valid
            CHECK (operation IN ('embedding', 'rerank', 'chat')),
          CONSTRAINT model_invocations_model_nonempty
            CHECK (btrim(model) <> ''),
          CONSTRAINT model_invocations_request_id_valid
            CHECK (
              provider_request_id IS NULL
              OR (
                btrim(provider_request_id) <> ''
                AND length(provider_request_id) <= 512
              )
            ),
          CONSTRAINT model_invocations_status_valid
            CHECK (status IN ('STARTED', 'SUCCEEDED', 'FAILED', 'UNKNOWN')),
          CONSTRAINT model_invocations_counts_valid
            CHECK (
              item_count >= 0
              AND prompt_tokens >= 0
              AND completion_tokens >= 0
              AND total_tokens >= 0
              AND total_tokens = prompt_tokens + completion_tokens
            ),
          CONSTRAINT model_invocations_elapsed_valid
            CHECK (
              (status = 'STARTED' AND elapsed_ms IS NULL)
              OR (status <> 'STARTED' AND elapsed_ms >= 0)
            ),
          CONSTRAINT model_invocations_error_code_valid
            CHECK (
              error_code IS NULL
              OR (btrim(error_code) <> '' AND length(error_code) <= 128)
            ),
          CONSTRAINT model_invocations_error_consistent
            CHECK (
              (status = 'SUCCEEDED' AND error_code IS NULL)
              OR (status = 'FAILED' AND error_code IS NOT NULL)
              OR (status = 'UNKNOWN' AND error_code IS NOT NULL)
              OR (status = 'STARTED' AND error_code IS NULL)
            ),
          CONSTRAINT model_invocations_timestamps_valid
            CHECK (
              created_at >= started_at
              AND (
                (status = 'STARTED' AND finished_at IS NULL)
                OR (status <> 'STARTED' AND finished_at >= started_at)
              )
            )
        );
        """
    )

    op.execute(
        """
        COMMENT ON TABLE rag.model_invocations IS
          'Metadata-only provider audit log. Provider inputs and outputs are forbidden.';
        """
    )

    op.execute(
        """
        CREATE INDEX model_invocations_trace_lookup
          ON rag.model_invocations (trace_id, started_at DESC);
        """
    )
    op.execute(
        """
        CREATE INDEX model_invocations_profile_status_lookup
          ON rag.model_invocations (profile_hash, operation, status, started_at DESC);
        """
    )
    op.execute(
        """
        CREATE INDEX chunks_active_embedding_profile_filter
          ON rag.chunks (
            knowledge_base_id,
            embedding_profile_hash,
            access_scope_id
          )
          WHERE searchable;
        """
    )


def downgrade() -> None:
    op.execute("DROP INDEX IF EXISTS rag.chunks_active_embedding_profile_filter;")
    op.execute("DROP TABLE IF EXISTS rag.model_invocations;")
    op.execute("DROP TABLE IF EXISTS rag.chunk_embedding_assignments;")
    op.execute("DROP TABLE IF EXISTS rag.embedding_cache;")
    op.execute(
        """
        ALTER TABLE rag.chunks
          DROP CONSTRAINT IF EXISTS chunks_id_embedding_text_sha256_key,
          DROP CONSTRAINT IF EXISTS chunks_citation_id_key,
          DROP COLUMN IF EXISTS citation_id;
        """
    )
    op.execute(
        """
        ALTER TABLE rag.knowledge_bases
          DROP CONSTRAINT IF EXISTS knowledge_bases_active_embedding_profile_fk,
          DROP CONSTRAINT IF EXISTS knowledge_bases_active_embedding_profile_hash_format,
          DROP CONSTRAINT IF EXISTS knowledge_bases_active_embedding_profile_kind,
          DROP COLUMN IF EXISTS active_embedding_profile_kind,
          DROP COLUMN IF EXISTS active_embedding_profile_hash;
        """
    )
    op.execute("DROP TABLE IF EXISTS rag.model_profiles;")