-- Initialize pgvector extension and research tables
CREATE EXTENSION IF NOT EXISTS vector;
CREATE EXTENSION IF NOT EXISTS pg_trgm;

-- Document storage
CREATE TABLE IF NOT EXISTS documents (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    filename TEXT NOT NULL,
    doc_id TEXT,
    mime_type TEXT,
    file_path TEXT,
    status TEXT DEFAULT 'pending',
    page_count INTEGER DEFAULT 0,
    full_text TEXT,
    metadata JSONB DEFAULT '{}',
    created_at TIMESTAMPTZ DEFAULT NOW(),
    updated_at TIMESTAMPTZ DEFAULT NOW()
);

-- Vector embeddings for similarity search
CREATE TABLE IF NOT EXISTS chunks (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    doc_id UUID REFERENCES documents(id) ON DELETE CASCADE,
    content TEXT NOT NULL,
    vector VECTOR(4096),
    page_num INTEGER,
    block_index INTEGER,
    polygon JSONB,
    chunk_type TEXT DEFAULT 'text',
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- Research memories/knowledge store
CREATE TABLE IF NOT EXISTS memories (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    session_id UUID REFERENCES documents(id),
    content TEXT NOT NULL,
    vector VECTOR(4096),
    memory_type TEXT DEFAULT 'fact',
    importance INTEGER DEFAULT 3,
    source_doc_id UUID REFERENCES documents(id),
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- Research findings
CREATE TABLE IF NOT EXISTS findings (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    session_id UUID REFERENCES documents(id),
    question TEXT,
    answer TEXT,
    summary TEXT,
    vector VECTOR(4096),
    relevant_chunks JSONB,
    agent_name TEXT,
    confidence FLOAT,
    created_at TIMESTAMPTZ DEFAULT NOW()
);

-- Research sessions
CREATE TABLE IF NOT EXISTS research_sessions (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    query TEXT NOT NULL,
    status TEXT DEFAULT 'running',
    documents JSONB DEFAULT '[]',
    findings JSONB DEFAULT '[]',
    created_at TIMESTAMPTZ DEFAULT NOW(),
    completed_at TIMESTAMPTZ
);

-- Indexes
CREATE INDEX IF NOT EXISTS idx_chunks_vector ON chunks USING ivfflat (vector vector_ip_opclass) WITH (lists = 100);
CREATE INDEX IF NOT EXISTS idx_memories_vector ON memories USING ivfflat (vector vector_ip_opclass) WITH (lists = 100);
CREATE INDEX IF NOT EXISTS idx_findings_vector ON findings USING ivfflat (vector vector_ip_opclass) WITH (lists = 100);
CREATE INDEX IF NOT EXISTS idx_chunks_trgm ON chunks USING gin (content gin_trgm_ops);
CREATE INDEX IF NOT EXISTS idx_chunks_doc ON chunks(doc_id);
CREATE INDEX IF NOT EXISTS idx_chunks_doc ON chunks(doc_id);
CREATE INDEX IF NOT EXISTS idx_memories_session ON memories(session_id);
CREATE INDEX IF NOT EXISTS idx_findings_session ON findings(session_id);
CREATE INDEX IF NOT EXISTS idx_documents_status ON documents(status);

-- Pipeline intermediate stage storage
CREATE TABLE IF NOT EXISTS pipeline_stages (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    session_id UUID REFERENCES research_sessions(id) ON DELETE CASCADE,
    stage TEXT NOT NULL CHECK (stage IN ('triage','evidence','synthesis')),
    output TEXT,
    state JSONB DEFAULT '{}',
    created_at TIMESTAMPTZ DEFAULT NOW()
);
CREATE INDEX IF NOT EXISTS idx_pipeline_session ON pipeline_stages(session_id);
CREATE INDEX IF NOT EXISTS idx_pipeline_stage ON pipeline_stages(stage, session_id);

-- Create function for cosine similarity
CREATE OR REPLACE FUNCTION cosine_similarity(v1 VECTOR, v2 VECTOR)
RETURNS FLOAT AS $$
    SELECT 1 - (v1 <-> v2) / 2;
$$ LANGUAGE SQL;

-- Create function for vector similarity search
CREATE OR REPLACE FUNCTION match_vectors(query_embedding VECTOR(4096), match_count INTEGER DEFAULT 10, min_score FLOAT DEFAULT 0)
RETURNS TABLE(id UUID, chunk_id UUID, content TEXT, doc_id UUID, page_num INT, similarity FLOAT)
LANGUAGE plpgsql
AS $$
BEGIN
    RETURN QUERY
    SELECT
        m.id::TEXT,
        m.id::TEXT,
        m.content::TEXT,
        m.doc_id::TEXT,
        m.page_num::INT,
        cosine_similarity(m.vector, query_embedding)::FLOAT
    FROM chunks m
    WHERE cosine_similarity(m.vector, query_embedding) >= min_score
    ORDER BY m.vector <-> query_embedding
    LIMIT match_count;
END;
$$;