-- Enable pgvector extension
CREATE EXTENSION IF NOT EXISTS vector;

-- Chunks table (matching Drizzle schema from Part 8)
CREATE TABLE chunks (
  id           UUID PRIMARY KEY DEFAULT gen_random_uuid(),
  document_id  UUID NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
  org_id       UUID NOT NULL REFERENCES organisations(id),
  content      TEXT NOT NULL,
  heading      TEXT NOT NULL,
  chunk_index  INTEGER NOT NULL,
  sha256       TEXT NOT NULL,
  token_count  INTEGER NOT NULL,
  embedding    vector(1536) NOT NULL,
  created_at   TIMESTAMPTZ DEFAULT now()
);

-- CRITICAL: Set maintenance_work_mem before building HNSW.
-- Without this, pgvector falls back to disk-based building (10–20× slower).
SET maintenance_work_mem = '2GB';  -- Use ≥4 GB for >500k vectors

-- HNSW index using inner product (requires L2-normalised vectors — Part 8)
-- vector_ip_ops = inner product = cosine similarity for unit vectors (fastest)
-- ef_construction=64 is minimum; use 128–256 for production-grade recall (>98%)
CREATE INDEX CONCURRENTLY chunks_embedding_hnsw_idx
  ON chunks
  USING hnsw (embedding vector_ip_ops)
  WITH (m = 16, ef_construction = 128);  -- 128 is the production baseline

-- Composite index for tenant-scoped lookups
CREATE INDEX chunks_org_doc_idx ON chunks (org_id, document_id);

-- Deduplication index
CREATE UNIQUE INDEX chunks_sha256_doc_idx ON chunks (sha256, document_id);

// ✅ Tune ef_search per query — higher recall at cost of latency
// file: packages/db/src/retrieval/VectorSearch.ts

import { db } from '../drizzle'
import { sql } from 'drizzle-orm'

export async function vectorSearch(
  orgId: string,
  queryEmbedding: number[],
  limit = 20,
  efSearch = 40,  // Default; increase to 80 for higher recall at cost of latency
): Promise<SearchResult[]> {
  // Set ef_search for this session/transaction
  await db.execute(sql`SET LOCAL hnsw.ef_search = ${efSearch}`)

  const vectorLiteral = `[${queryEmbedding.join(',')}]`

  // Inner product query — works correctly only with L2-normalised vectors
  // <#> = negative inner product (pgvector uses this for similarity ordering)
  return db.execute(sql`
    SELECT
      id,
      content,
      heading,
      document_id,
      token_count,
      1 - (embedding <#> ${vectorLiteral}::vector) AS similarity
    FROM chunks
    WHERE org_id = ${orgId}::uuid    -- RLS also enforces this, but explicit is faster
    ORDER BY embedding <#> ${vectorLiteral}::vector
    LIMIT ${limit}
  `)
}

-- Enable RLS on chunks table
ALTER TABLE chunks ENABLE ROW LEVEL SECURITY;
ALTER TABLE chunks FORCE ROW LEVEL SECURITY;  -- Applies to table owners too

-- Create a PostgreSQL role for the application
CREATE ROLE workspace_app LOGIN PASSWORD 'strong_password_here';

-- Grant DML — but NOT bypass RLS
GRANT SELECT, INSERT, UPDATE, DELETE ON chunks TO workspace_app;

-- Policy: app can only see/modify chunks in its current org context
-- The app sets current_setting('app.current_org_id') at session start
CREATE POLICY chunks_org_isolation ON chunks
  USING (org_id = current_setting('app.current_org_id')::uuid)
  WITH CHECK (org_id = current_setting('app.current_org_id')::uuid);

// ✅ Session-level org context — set before every query in the Workspace
// file: packages/db/src/middleware/orgContext.ts

import { db } from '../drizzle'
import { sql } from 'drizzle-orm'

/**
 * Sets the PostgreSQL session variable that RLS policies read.
 * MUST be called at the start of every request handler.
 */
export async function setOrgContext(orgId: string): Promise<void> {
  // Validate orgId is a valid UUID before passing to database
  if (!/^[0-9a-f-]{36}$/.test(orgId)) {
    throw new Error(`Invalid org ID format: ${orgId}`)
  }
  await db.execute(sql`SET LOCAL app.current_org_id = ${orgId}`)
}

// Fastify middleware usage:
// fastify.addHook('preHandler', async (req) => {
//   const orgId = req.user.orgId  // from JWT
//   await setOrgContext(orgId)
// })

-- Add shadow column for new embedding model — old column untouched
ALTER TABLE chunks ADD COLUMN embedding_v2 vector(3072);

-- New HNSW index on shadow column — built CONCURRENTLY (non-blocking)
CREATE INDEX CONCURRENTLY chunks_embedding_v2_hnsw_idx
  ON chunks
  USING hnsw (embedding_v2 vector_ip_ops)
  WITH (m = 16, ef_construction = 64);

// ✅ Dual-write during migration window
// file: packages/db/src/ingestion/DualWriteIngester.ts

export async function ingestWithDualWrite(chunk: DocumentChunk): Promise<void> {
  // Embed in both old and new model simultaneously
  const [embeddingV1, embeddingV2] = await Promise.all([
    embedWithModel('text-embedding-3-large', chunk.content),   // 1536d
    embedWithModel('text-embedding-future', chunk.content),    // 3072d
  ])

  await db.update(chunks)
    .set({
      embedding: embeddingV1,      // Existing column
      embeddingV2: embeddingV2,    // Shadow column
    })
    .where(eq(chunks.id, chunk.id))
}

// After all chunks are backfilled to shadow column:
// 1. Switch query routing to embedding_v2 column
// 2. DROP COLUMN embedding (old) — non-blocking
// 3. ALTER TABLE chunks RENAME COLUMN embedding_v2 TO embedding