-- Migration: Create vector database tables for pgvector integration
-- Created: 2025-01-28

-- Create document_chunks table for storing text chunks with embeddings
CREATE TABLE IF NOT EXISTS document_chunks (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    document_id UUID NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
    content TEXT NOT NULL,
    metadata JSONB NOT NULL DEFAULT '{}',
    embedding vector(1536), -- OpenAI text-embedding-3-small dimension
    chunk_index INTEGER NOT NULL,
    section VARCHAR(100),
    page_number INTEGER,
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- Create indexes for better query performance
CREATE INDEX IF NOT EXISTS idx_document_chunks_document_id ON document_chunks(document_id);
CREATE INDEX IF NOT EXISTS idx_document_chunks_section ON document_chunks(section);
CREATE INDEX IF NOT EXISTS idx_document_chunks_chunk_index ON document_chunks(chunk_index);
CREATE INDEX IF NOT EXISTS idx_document_chunks_created_at ON document_chunks(created_at);

-- Create vector similarity search index with optimized parameters for 1536 dimensions
CREATE INDEX IF NOT EXISTS idx_document_chunks_embedding ON document_chunks USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);

-- Create composite indexes for common queries
CREATE INDEX IF NOT EXISTS idx_document_chunks_document_section ON document_chunks(document_id, section);
CREATE INDEX IF NOT EXISTS idx_document_chunks_document_chunk ON document_chunks(document_id, chunk_index);

-- Create metadata indexes for filtering
CREATE INDEX IF NOT EXISTS idx_document_chunks_metadata_gin ON document_chunks USING GIN (metadata);

-- Create vector_similarity_searches table for tracking search queries
CREATE TABLE IF NOT EXISTS vector_similarity_searches (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
    query_text TEXT NOT NULL,
    query_embedding vector(1536),
    search_results JSONB NOT NULL DEFAULT '[]',
    filters JSONB DEFAULT '{}',
    limit_count INTEGER DEFAULT 10,
    similarity_threshold DECIMAL(3,2) DEFAULT 0.7,
    processing_time_ms INTEGER,
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- Create indexes for search tracking
CREATE INDEX IF NOT EXISTS idx_vector_searches_user_id ON vector_similarity_searches(user_id);
CREATE INDEX IF NOT EXISTS idx_vector_searches_created_at ON vector_similarity_searches(created_at);
CREATE INDEX IF NOT EXISTS idx_vector_searches_query_embedding ON vector_similarity_searches USING ivfflat (query_embedding vector_cosine_ops) WITH (lists = 50);

-- Create document_similarities table for tracking document-to-document similarities
CREATE TABLE IF NOT EXISTS document_similarities (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    source_document_id UUID NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
    target_document_id UUID NOT NULL REFERENCES documents(id) ON DELETE CASCADE,
    similarity_score DECIMAL(5,4) NOT NULL,
    similarity_type VARCHAR(50) NOT NULL DEFAULT 'content', -- 'content', 'financial', 'industry', etc.
    metadata JSONB DEFAULT '{}',
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    
    -- Ensure unique combinations
    UNIQUE(source_document_id, target_document_id, similarity_type)
);

-- Create indexes for document similarities
CREATE INDEX IF NOT EXISTS idx_document_similarities_source ON document_similarities(source_document_id);
CREATE INDEX IF NOT EXISTS idx_document_similarities_target ON document_similarities(target_document_id);
CREATE INDEX IF NOT EXISTS idx_document_similarities_score ON document_similarities(similarity_score DESC);
CREATE INDEX IF NOT EXISTS idx_document_similarities_type ON document_similarities(similarity_type);

-- Create industry_embeddings table for industry-specific analysis
CREATE TABLE IF NOT EXISTS industry_embeddings (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    industry_name VARCHAR(100) NOT NULL,
    industry_description TEXT,
    embedding vector(1536),
    document_count INTEGER DEFAULT 0,
    average_similarity DECIMAL(5,4),
    created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    
    UNIQUE(industry_name)
);

-- Create indexes for industry embeddings
CREATE INDEX IF NOT EXISTS idx_industry_embeddings_name ON industry_embeddings(industry_name);
CREATE INDEX IF NOT EXISTS idx_industry_embeddings_embedding ON industry_embeddings USING ivfflat (embedding vector_cosine_ops) WITH (lists = 50);

-- Create functions for vector operations

-- Function to calculate cosine similarity between two vectors
CREATE OR REPLACE FUNCTION cosine_similarity(a vector, b vector)
RETURNS DECIMAL
LANGUAGE plpgsql
AS $$
BEGIN
    RETURN 1 - (a <=> b);
END;
$$;

-- Function to find similar documents with 3072-dimensional vectors
CREATE OR REPLACE FUNCTION find_similar_documents(
    query_embedding vector(3072),
    similarity_threshold DECIMAL DEFAULT 0.7,
    max_results INTEGER DEFAULT 10,
    document_filter UUID DEFAULT NULL
)
RETURNS TABLE (
    document_id UUID,
    similarity_score DECIMAL,
    chunk_content TEXT,
    metadata JSONB
)
LANGUAGE plpgsql
AS $$
BEGIN
    RETURN QUERY
    SELECT 
        dc.document_id,
        1 - (dc.embedding <=> query_embedding) as similarity_score,
        dc.content as chunk_content,
        dc.metadata
    FROM document_chunks dc
    WHERE dc.embedding IS NOT NULL
        AND (document_filter IS NULL OR dc.document_id = document_filter)
        AND 1 - (dc.embedding <=> query_embedding) >= similarity_threshold
    ORDER BY dc.embedding <=> query_embedding
    LIMIT max_results;
END;
$$;

-- Function to update document similarities
CREATE OR REPLACE FUNCTION update_document_similarities()
RETURNS void
LANGUAGE plpgsql
AS $$
DECLARE
    doc1 RECORD;
    doc2 RECORD;
    similarity DECIMAL;
BEGIN
    -- Clear existing similarities
    DELETE FROM document_similarities;
    
    -- Calculate similarities between all document pairs
    FOR doc1 IN SELECT DISTINCT document_id FROM document_chunks LOOP
        FOR doc2 IN SELECT DISTINCT document_id FROM document_chunks WHERE document_id > doc1.document_id LOOP
            -- Calculate average similarity between chunks
            SELECT AVG(1 - (c1.embedding <=> c2.embedding)) INTO similarity
            FROM document_chunks c1
            CROSS JOIN document_chunks c2
            WHERE c1.document_id = doc1.document_id
            AND c2.document_id = doc2.document_id
            AND c1.embedding IS NOT NULL
            AND c2.embedding IS NOT NULL;
            
            -- Insert if similarity is above threshold
            IF similarity > 0.5 THEN
                INSERT INTO document_similarities (
                    source_document_id, target_document_id, similarity_score, similarity_type
                ) VALUES (
                    doc1.document_id, doc2.document_id, similarity, 'content'
                );
            END IF;
        END LOOP;
    END LOOP;
END;
$$;

-- Function to update document_chunks updated_at timestamp
CREATE OR REPLACE FUNCTION update_document_chunks_updated_at()
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at = CURRENT_TIMESTAMP;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER trigger_update_document_chunks_updated_at
    BEFORE UPDATE ON document_chunks
    FOR EACH ROW
    EXECUTE FUNCTION update_document_chunks_updated_at();

-- Function to update industry_embeddings updated_at timestamp
CREATE OR REPLACE FUNCTION update_industry_embeddings_updated_at()
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at = CURRENT_TIMESTAMP;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER trigger_update_industry_embeddings_updated_at
    BEFORE UPDATE ON industry_embeddings
    FOR EACH ROW
    EXECUTE FUNCTION update_industry_embeddings_updated_at();

-- Add comments for documentation
COMMENT ON TABLE document_chunks IS 'Stores document text chunks with 3072-dimensional embeddings for semantic search';
COMMENT ON COLUMN document_chunks.embedding IS 'OpenAI text-embedding-3-large vector (3072 dimensions)';
COMMENT ON TABLE vector_similarity_searches IS 'Tracks search queries and results for analytics';
COMMENT ON TABLE document_similarities IS 'Stores document-to-document similarity scores';
COMMENT ON TABLE industry_embeddings IS 'Stores industry-specific embeddings for sector analysis'; 