-- ACR PostgreSQL Schema V2
-- Purpose:
--   1. Support canonical_song/work/recording/asset/window hierarchy
--   2. Support encoder-first evolution via model_registry + feature_set_registry
--   3. Support pgvector-backed hot reference sets without binding the entire system to one vector table

CREATE EXTENSION IF NOT EXISTS vector;

-- =========================================================
-- 1. Canonical business entities
-- =========================================================

CREATE TABLE IF NOT EXISTS canonical_song (
    canonical_song_id          BIGSERIAL PRIMARY KEY,
    biz_song_code              TEXT UNIQUE,
    title                      TEXT NOT NULL,
    title_norm                 TEXT,
    primary_artist             TEXT,
    primary_artist_norm        TEXT,
    language_code              TEXT,
    rights_status              TEXT,
    metadata_json              JSONB NOT NULL DEFAULT '{}'::jsonb,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS work (
    work_id                    BIGSERIAL PRIMARY KEY,
    canonical_song_id          BIGINT NOT NULL REFERENCES canonical_song(canonical_song_id),
    work_code                  TEXT UNIQUE,
    work_title                 TEXT NOT NULL,
    work_title_norm            TEXT,
    composer                   TEXT,
    lyricist                   TEXT,
    publisher                  TEXT,
    iswc                       TEXT,
    metadata_json              JSONB NOT NULL DEFAULT '{}'::jsonb,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS recording (
    recording_id               BIGSERIAL PRIMARY KEY,
    work_id                    BIGINT NOT NULL REFERENCES work(work_id),
    canonical_song_id          BIGINT NOT NULL REFERENCES canonical_song(canonical_song_id),
    recording_code             TEXT UNIQUE,
    recording_title            TEXT,
    artist_name                TEXT,
    album_name                 TEXT,
    version_type               TEXT NOT NULL,
    is_reference               BOOLEAN NOT NULL DEFAULT FALSE,
    reference_priority         INTEGER NOT NULL DEFAULT 100,
    release_date               DATE,
    isrc                       TEXT,
    duration_sec               NUMERIC(10,3),
    metadata_json              JSONB NOT NULL DEFAULT '{}'::jsonb,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

-- =========================================================
-- 2. Assets and windows
-- =========================================================

CREATE TABLE IF NOT EXISTS recording_asset (
    asset_id                   BIGSERIAL PRIMARY KEY,
    recording_id               BIGINT NOT NULL REFERENCES recording(recording_id),
    asset_role                 TEXT NOT NULL,
    storage_uri                TEXT NOT NULL,
    storage_scheme             TEXT NOT NULL,
    file_ext                   TEXT,
    mime_type                  TEXT,
    file_size_bytes            BIGINT,
    audio_sha256               TEXT,
    sample_rate                INTEGER,
    channels                   INTEGER,
    bit_rate_kbps              INTEGER,
    codec_name                 TEXT,
    duration_sec               NUMERIC(10,3),
    loudness_lufs              NUMERIC(8,3),
    normalized_storage_uri     TEXT,
    ingest_status              TEXT NOT NULL DEFAULT 'ready',
    metadata_json              JSONB NOT NULL DEFAULT '{}'::jsonb,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE UNIQUE INDEX IF NOT EXISTS uq_recording_asset_sha256
ON recording_asset(audio_sha256)
WHERE audio_sha256 IS NOT NULL;

CREATE TABLE IF NOT EXISTS audio_window (
    window_id                  BIGSERIAL PRIMARY KEY,
    asset_id                   BIGINT NOT NULL REFERENCES recording_asset(asset_id),
    recording_id               BIGINT NOT NULL REFERENCES recording(recording_id),
    work_id                    BIGINT NOT NULL REFERENCES work(work_id),
    canonical_song_id          BIGINT NOT NULL REFERENCES canonical_song(canonical_song_id),
    window_index               INTEGER NOT NULL,
    start_sec                  NUMERIC(10,3) NOT NULL,
    end_sec                    NUMERIC(10,3) NOT NULL,
    duration_sec               NUMERIC(10,3) NOT NULL,
    segment_role               TEXT NOT NULL DEFAULT 'reference',
    segment_type               TEXT,
    quality_score              NUMERIC(8,5),
    active_for_index           BOOLEAN NOT NULL DEFAULT TRUE,
    metadata_json              JSONB NOT NULL DEFAULT '{}'::jsonb,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE UNIQUE INDEX IF NOT EXISTS uq_audio_window_asset_idx
ON audio_window(asset_id, window_index);

-- =========================================================
-- 3. Model and feature registries
-- =========================================================

CREATE TABLE IF NOT EXISTS model_registry (
    model_id                   BIGSERIAL PRIMARY KEY,
    model_name                 TEXT NOT NULL,
    model_family               TEXT NOT NULL,
    model_version              TEXT NOT NULL,
    model_source               TEXT,
    model_uri                  TEXT,
    license_name               TEXT,
    input_modality             TEXT NOT NULL DEFAULT 'audio',
    input_sample_rate          INTEGER,
    input_channel_mode         TEXT DEFAULT 'mono',
    default_window_sec         NUMERIC(10,3),
    default_hop_sec            NUMERIC(10,3),
    output_embedding_dim       INTEGER,
    pooling_supported          TEXT[],
    layer_selection_supported  BOOLEAN NOT NULL DEFAULT FALSE,
    is_trainable               BOOLEAN NOT NULL DEFAULT FALSE,
    is_active                  BOOLEAN NOT NULL DEFAULT TRUE,
    metadata_json              JSONB NOT NULL DEFAULT '{}'::jsonb,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    UNIQUE(model_name, model_version)
);

CREATE TABLE IF NOT EXISTS feature_set_registry (
    feature_set_id             BIGSERIAL PRIMARY KEY,
    model_id                   BIGINT NOT NULL REFERENCES model_registry(model_id),
    feature_name               TEXT NOT NULL,
    feature_level              TEXT NOT NULL,
    extraction_granularity     TEXT NOT NULL,
    window_sec                 NUMERIC(10,3),
    hop_sec                    NUMERIC(10,3),
    embedding_dim              INTEGER,
    pooling_strategy           TEXT,
    layer_selection            TEXT,
    normalize_l2              BOOLEAN NOT NULL DEFAULT TRUE,
    distance_metric            TEXT NOT NULL,
    quantization_type          TEXT,
    feature_schema_version     TEXT NOT NULL,
    config_json                JSONB NOT NULL DEFAULT '{}'::jsonb,
    status                     TEXT NOT NULL DEFAULT 'active',
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS feature_extraction_job (
    extraction_job_id          BIGSERIAL PRIMARY KEY,
    feature_set_id             BIGINT NOT NULL REFERENCES feature_set_registry(feature_set_id),
    target_scope               TEXT NOT NULL,
    job_status                 TEXT NOT NULL DEFAULT 'pending',
    shard_key                  TEXT,
    input_count                BIGINT,
    output_count               BIGINT,
    started_at                 TIMESTAMPTZ,
    finished_at                TIMESTAMPTZ,
    log_uri                    TEXT,
    metadata_json              JSONB NOT NULL DEFAULT '{}'::jsonb,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

-- =========================================================
-- 4. Feature facts
-- =========================================================

CREATE TABLE IF NOT EXISTS audio_embedding (
    embedding_id               BIGSERIAL PRIMARY KEY,
    feature_set_id             BIGINT NOT NULL REFERENCES feature_set_registry(feature_set_id),
    extraction_job_id          BIGINT REFERENCES feature_extraction_job(extraction_job_id),
    asset_id                   BIGINT REFERENCES recording_asset(asset_id),
    window_id                  BIGINT REFERENCES audio_window(window_id),
    recording_id               BIGINT NOT NULL REFERENCES recording(recording_id),
    work_id                    BIGINT NOT NULL REFERENCES work(work_id),
    canonical_song_id          BIGINT NOT NULL REFERENCES canonical_song(canonical_song_id),
    embedding_storage_mode     TEXT NOT NULL,
    embedding_uri              TEXT,
    vector_norm                NUMERIC(12,6),
    checksum                   TEXT,
    is_indexed                 BOOLEAN NOT NULL DEFAULT FALSE,
    metadata_json              JSONB NOT NULL DEFAULT '{}'::jsonb,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    CONSTRAINT ck_audio_embedding_scope CHECK (asset_id IS NOT NULL OR window_id IS NOT NULL)
);

CREATE TABLE IF NOT EXISTS audio_embedding_vector_192 (
    embedding_id               BIGINT PRIMARY KEY REFERENCES audio_embedding(embedding_id) ON DELETE CASCADE,
    embedding                  VECTOR(192) NOT NULL
);

CREATE TABLE IF NOT EXISTS audio_embedding_vector_768 (
    embedding_id               BIGINT PRIMARY KEY REFERENCES audio_embedding(embedding_id) ON DELETE CASCADE,
    embedding                  VECTOR(768) NOT NULL
);

CREATE TABLE IF NOT EXISTS audio_fingerprint (
    fingerprint_id             BIGSERIAL PRIMARY KEY,
    feature_set_id             BIGINT NOT NULL REFERENCES feature_set_registry(feature_set_id),
    asset_id                   BIGINT REFERENCES recording_asset(asset_id),
    window_id                  BIGINT REFERENCES audio_window(window_id),
    recording_id               BIGINT NOT NULL REFERENCES recording(recording_id),
    work_id                    BIGINT NOT NULL REFERENCES work(work_id),
    canonical_song_id          BIGINT NOT NULL REFERENCES canonical_song(canonical_song_id),
    fingerprint_uri            TEXT,
    hash_count                 INTEGER,
    is_indexed                 BOOLEAN NOT NULL DEFAULT FALSE,
    metadata_json              JSONB NOT NULL DEFAULT '{}'::jsonb,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE UNIQUE INDEX IF NOT EXISTS uq_audio_fingerprint_feature_asset
ON audio_fingerprint(feature_set_id, asset_id);

CREATE TABLE IF NOT EXISTS reference_set_registry (
    reference_set_id           BIGSERIAL PRIMARY KEY,
    set_name                   TEXT NOT NULL UNIQUE,
    description                TEXT,
    encoder_scope              TEXT,
    status                     TEXT NOT NULL DEFAULT 'active',
    metadata_json              JSONB NOT NULL DEFAULT '{}'::jsonb,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS reference_set_member (
    reference_set_id           BIGINT NOT NULL REFERENCES reference_set_registry(reference_set_id) ON DELETE CASCADE,
    recording_id               BIGINT NOT NULL REFERENCES recording(recording_id) ON DELETE CASCADE,
    member_role                TEXT NOT NULL DEFAULT 'hot_reference',
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    PRIMARY KEY(reference_set_id, recording_id)
);

-- =========================================================
-- 4.5 Lineage invariants (recommended production hardening)
-- =========================================================

CREATE OR REPLACE FUNCTION check_recording_lineage()
RETURNS trigger AS $$
DECLARE
    work_song_id BIGINT;
BEGIN
    SELECT canonical_song_id INTO work_song_id
    FROM work
    WHERE work_id = NEW.work_id;

    IF work_song_id IS NULL THEN
        RAISE EXCEPTION 'Invalid work_id=% for recording', NEW.work_id;
    END IF;

    IF NEW.canonical_song_id <> work_song_id THEN
        RAISE EXCEPTION 'recording.canonical_song_id % mismatches work.canonical_song_id %', NEW.canonical_song_id, work_song_id;
    END IF;

    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION check_audio_window_lineage()
RETURNS trigger AS $$
DECLARE
    asset_recording_id BIGINT;
    rec_work_id BIGINT;
    rec_song_id BIGINT;
BEGIN
    SELECT recording_id INTO asset_recording_id
    FROM recording_asset
    WHERE asset_id = NEW.asset_id;

    SELECT work_id, canonical_song_id INTO rec_work_id, rec_song_id
    FROM recording
    WHERE recording_id = NEW.recording_id;

    IF asset_recording_id IS NULL OR rec_work_id IS NULL THEN
        RAISE EXCEPTION 'Invalid asset_id=% or recording_id=% for audio_window', NEW.asset_id, NEW.recording_id;
    END IF;

    IF NEW.recording_id <> asset_recording_id THEN
        RAISE EXCEPTION 'audio_window.recording_id % mismatches recording_asset.recording_id %', NEW.recording_id, asset_recording_id;
    END IF;

    IF NEW.work_id <> rec_work_id OR NEW.canonical_song_id <> rec_song_id THEN
        RAISE EXCEPTION 'audio_window lineage mismatch for recording_id=%', NEW.recording_id;
    END IF;

    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION check_audio_embedding_lineage()
RETURNS trigger AS $$
DECLARE
    parent_recording_id BIGINT;
    parent_work_id BIGINT;
    parent_song_id BIGINT;
BEGIN
    IF NEW.window_id IS NOT NULL THEN
        SELECT recording_id, work_id, canonical_song_id
          INTO parent_recording_id, parent_work_id, parent_song_id
        FROM audio_window
        WHERE window_id = NEW.window_id;
    ELSE
        SELECT r.recording_id, r.work_id, r.canonical_song_id
          INTO parent_recording_id, parent_work_id, parent_song_id
        FROM recording_asset ra
        JOIN recording r ON r.recording_id = ra.recording_id
        WHERE ra.asset_id = NEW.asset_id;
    END IF;

    IF parent_recording_id IS NULL THEN
        RAISE EXCEPTION 'Invalid parent reference for audio_embedding';
    END IF;

    IF NEW.recording_id <> parent_recording_id OR NEW.work_id <> parent_work_id OR NEW.canonical_song_id <> parent_song_id THEN
        RAISE EXCEPTION 'audio_embedding lineage mismatch';
    END IF;

    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- =========================================================
-- 5. Retrieval index and results
-- =========================================================

CREATE TABLE IF NOT EXISTS retrieval_index_registry (
    retrieval_index_id         BIGSERIAL PRIMARY KEY,
    feature_set_id             BIGINT NOT NULL REFERENCES feature_set_registry(feature_set_id),
    index_name                 TEXT NOT NULL,
    index_backend              TEXT NOT NULL,
    index_type                 TEXT NOT NULL,
    storage_uri                TEXT,
    shard_no                   INTEGER,
    row_count                  BIGINT,
    index_status               TEXT NOT NULL DEFAULT 'active',
    config_json                JSONB NOT NULL DEFAULT '{}'::jsonb,
    built_at                   TIMESTAMPTZ,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS retrieval_candidate (
    retrieval_candidate_id     BIGSERIAL PRIMARY KEY,
    query_id                   TEXT NOT NULL,
    retrieval_index_id         BIGINT REFERENCES retrieval_index_registry(retrieval_index_id),
    feature_set_id             BIGINT REFERENCES feature_set_registry(feature_set_id),
    source_lane                TEXT NOT NULL CHECK (source_lane IN ('fingerprint', 'semantic', 'cover', 'melody', 'fusion')),
    candidate_level            TEXT NOT NULL CHECK (candidate_level IN ('window', 'recording', 'work', 'canonical_song')),
    candidate_id               BIGINT NOT NULL,
    evidence_window_id         BIGINT REFERENCES audio_window(window_id),
    raw_score                  NUMERIC(14,8),
    normalized_score           NUMERIC(14,8),
    rank_no                    INTEGER,
    metadata_json              JSONB NOT NULL DEFAULT '{}'::jsonb,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS match_decision (
    match_decision_id          BIGSERIAL PRIMARY KEY,
    query_id                   TEXT NOT NULL,
    canonical_song_id          BIGINT REFERENCES canonical_song(canonical_song_id),
    work_id                    BIGINT REFERENCES work(work_id),
    recording_id               BIGINT REFERENCES recording(recording_id),
    decision_status            TEXT NOT NULL,
    decision_score             NUMERIC(14,8),
    decision_reason            TEXT,
    metadata_json              JSONB NOT NULL DEFAULT '{}'::jsonb,
    created_at                 TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TRIGGER trg_recording_lineage
BEFORE INSERT OR UPDATE ON recording
FOR EACH ROW EXECUTE FUNCTION check_recording_lineage();

CREATE TRIGGER trg_audio_window_lineage
BEFORE INSERT OR UPDATE ON audio_window
FOR EACH ROW EXECUTE FUNCTION check_audio_window_lineage();

CREATE TRIGGER trg_audio_embedding_lineage
BEFORE INSERT OR UPDATE ON audio_embedding
FOR EACH ROW EXECUTE FUNCTION check_audio_embedding_lineage();

-- =========================================================
-- 6. Recommended indexes
-- =========================================================

CREATE INDEX IF NOT EXISTS idx_work_canonical_song_id
ON work(canonical_song_id);

CREATE INDEX IF NOT EXISTS idx_recording_work_id
ON recording(work_id);

CREATE INDEX IF NOT EXISTS idx_recording_canonical_song_id
ON recording(canonical_song_id);

CREATE INDEX IF NOT EXISTS idx_recording_reference
ON recording(is_reference, reference_priority);

CREATE INDEX IF NOT EXISTS idx_recording_asset_recording_id
ON recording_asset(recording_id);

CREATE INDEX IF NOT EXISTS idx_audio_window_asset_id
ON audio_window(asset_id);

CREATE INDEX IF NOT EXISTS idx_audio_window_recording_id
ON audio_window(recording_id);

CREATE INDEX IF NOT EXISTS idx_audio_window_canonical_song_id
ON audio_window(canonical_song_id);

CREATE INDEX IF NOT EXISTS idx_audio_window_active_for_index
ON audio_window(active_for_index);

CREATE UNIQUE INDEX IF NOT EXISTS uq_audio_embedding_feature_window
ON audio_embedding(feature_set_id, window_id)
WHERE window_id IS NOT NULL;

CREATE UNIQUE INDEX IF NOT EXISTS uq_audio_embedding_feature_asset
ON audio_embedding(feature_set_id, asset_id)
WHERE window_id IS NULL AND asset_id IS NOT NULL;

CREATE INDEX IF NOT EXISTS idx_audio_embedding_feature_set_id
ON audio_embedding(feature_set_id);

CREATE INDEX IF NOT EXISTS idx_audio_embedding_window_id
ON audio_embedding(window_id);

CREATE INDEX IF NOT EXISTS idx_audio_embedding_recording_id
ON audio_embedding(recording_id);

CREATE INDEX IF NOT EXISTS idx_reference_set_member_recording_id
ON reference_set_member(recording_id);

CREATE INDEX IF NOT EXISTS idx_retrieval_candidate_query_id
ON retrieval_candidate(query_id);

CREATE INDEX IF NOT EXISTS idx_match_decision_query_id
ON match_decision(query_id);

-- Optional hot-set HNSW indexes for pgvector-backed online retrieval
CREATE INDEX IF NOT EXISTS idx_audio_embedding_vector_192_cos_hnsw
ON audio_embedding_vector_192 USING hnsw (embedding vector_cosine_ops);

CREATE INDEX IF NOT EXISTS idx_audio_embedding_vector_768_cos_hnsw
ON audio_embedding_vector_768 USING hnsw (embedding vector_cosine_ops);