-- Install pgvector
CREATE EXTENSION vector;

-- Create a table with a vector column
CREATE TABLE documents (
    id BIGSERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding vector(1536),  -- OpenAI text-embedding-3-small dimensions
    metadata JSONB,
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- Create an HNSW index (better for high-recall queries)
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops)
WITH (m = 16, ef_construction = 64);

-- pgvectorscale requires pgvector
CREATE EXTENSION vectorscale CASCADE;

-- StreamingDiskANN index
CREATE INDEX ON documents USING diskann (embedding vector_cosine_ops);

-- That's it! Much simpler configuration than HNSW tuning

-- Main documents table
CREATE TABLE documents (
    id BIGSERIAL PRIMARY KEY,
    source_url TEXT,
    title TEXT,
    content TEXT NOT NULL,
    content_hash TEXT UNIQUE NOT NULL,  -- dedup
    metadata JSONB DEFAULT '{}',
    created_at TIMESTAMPTZ DEFAULT NOW(),
    updated_at TIMESTAMPTZ DEFAULT NOW()
);

-- Chunks table (for chunked documents)
CREATE TABLE document_chunks (
    id BIGSERIAL PRIMARY KEY,
    document_id BIGINT REFERENCES documents(id) ON DELETE CASCADE,
    chunk_index INT NOT NULL,
    content TEXT NOT NULL,
    embedding vector(1536),
    token_count INT,
    metadata JSONB DEFAULT '{}',
    created_at TIMESTAMPTZ DEFAULT NOW(),
    UNIQUE(document_id, chunk_index)
);

-- Indexes
CREATE INDEX ON document_chunks USING diskann (embedding vector_cosine_ops);
CREATE INDEX ON document_chunks (document_id);
CREATE INDEX ON document_chunks USING gin (metadata);

-- Full text search index for hybrid search
CREATE INDEX ON document_chunks USING gin (to_tsvector('english', content));

-- Hybrid search with RRF (Reciprocal Rank Fusion)
CREATE OR REPLACE FUNCTION hybrid_search(
    query_text TEXT,
    query_embedding vector(1536),
    match_count INT DEFAULT 10,
    full_text_weight FLOAT DEFAULT 1.0,
    vector_weight FLOAT DEFAULT 1.0,
    rrf_k INT DEFAULT 60
)
RETURNS TABLE (
    id BIGINT,
    document_id BIGINT,
    content TEXT,
    metadata JSONB,
    rrf_score FLOAT
)
LANGUAGE SQL
AS $$
WITH 
-- Full text search results
full_text AS (
    SELECT
        dc.id,
        dc.document_id,
        dc.content,
        dc.metadata,
        ROW_NUMBER() OVER (ORDER BY ts_rank_cd(
            to_tsvector('english', dc.content),
            plainto_tsquery('english', query_text)
        ) DESC) AS rank
    FROM document_chunks dc
    WHERE to_tsvector('english', dc.content) @@ plainto_tsquery('english', query_text)
    LIMIT match_count * 2
),
-- Vector similarity search results  
vector AS (
    SELECT
        dc.id,
        dc.document_id,
        dc.content,
        dc.metadata,
        ROW_NUMBER() OVER (ORDER BY dc.embedding <=> query_embedding) AS rank
    FROM document_chunks dc
    ORDER BY dc.embedding <=> query_embedding
    LIMIT match_count * 2
)
-- Reciprocal Rank Fusion
SELECT
    COALESCE(ft.id, v.id) AS id,
    COALESCE(ft.document_id, v.document_id) AS document_id,
    COALESCE(ft.content, v.content) AS content,
    COALESCE(ft.metadata, v.metadata) AS metadata,
    (
        COALESCE(full_text_weight * 1.0 / (rrf_k + ft.rank), 0) +
        COALESCE(vector_weight * 1.0 / (rrf_k + v.rank), 0)
    ) AS rrf_score
FROM full_text ft
FULL OUTER JOIN vector v ON ft.id = v.id
ORDER BY rrf_score DESC
LIMIT match_count;
$$;

-- Tune for high recall
CREATE INDEX ON document_chunks USING hnsw (embedding vector_cosine_ops)
WITH (
    m = 16,           -- connections per node (higher = better recall, more memory)
    ef_construction = 128  -- candidate list size during build (higher = better, slower build)
);

-- For searches, tune ef_search per query
SET hnsw.ef_search = 100;  -- higher = better recall, slower query

CREATE TABLE document_chunks (
    id BIGSERIAL,
    document_id BIGINT NOT NULL,
    chunk_index INT NOT NULL,
    content TEXT NOT NULL,
    embedding vector(1536),
    created_at TIMESTAMPTZ DEFAULT NOW()
) PARTITION BY RANGE (created_at);

CREATE TABLE document_chunks_2026_q1
PARTITION OF document_chunks
FOR VALUES FROM ('2026-01-01') TO ('2026-04-01');

CREATE TABLE document_chunks_2026_q2
PARTITION OF document_chunks
FOR VALUES FROM ('2026-04-01') TO ('2026-07-01');

-- Each partition gets its own index
CREATE INDEX ON document_chunks_2026_q2 
USING diskann (embedding vector_cosine_ops);

-- Check index usage
SELECT schemaname, tablename, indexname, idx_scan, idx_tup_read
FROM pg_stat_user_indexes
WHERE tablename = 'document_chunks'
ORDER BY idx_scan DESC;

-- Check slow vector queries
SELECT query, calls, mean_exec_time, max_exec_time
FROM pg_stat_statements
WHERE query LIKE '%<=>%'
ORDER BY mean_exec_time DESC
LIMIT 20;