-- Advanced RLS Configuration with User Authentication
-- This provides more granular control over data access

-- First, create a users table if it doesn't exist
CREATE TABLE IF NOT EXISTS public.users (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    email TEXT UNIQUE NOT NULL,
    role TEXT NOT NULL DEFAULT 'viewer', -- viewer, editor, admin
    created_at TIMESTAMPTZ DEFAULT NOW(),
    updated_at TIMESTAMPTZ DEFAULT NOW()
);

-- Add user_id columns to track ownership (if not exists)
DO $$
BEGIN
    -- Add user_id to datasources if not exists
    IF NOT EXISTS (SELECT 1 FROM information_schema.columns
                   WHERE table_schema = 'public'
                   AND table_name = 'datasources'
                   AND column_name = 'user_id') THEN
        ALTER TABLE public.datasources ADD COLUMN user_id UUID REFERENCES public.users(id);
    END IF;

    -- Add user_id to documents if not exists
    IF NOT EXISTS (SELECT 1 FROM information_schema.columns
                   WHERE table_schema = 'public'
                   AND table_name = 'documents'
                   AND column_name = 'user_id') THEN
        ALTER TABLE public.documents ADD COLUMN user_id UUID REFERENCES public.users(id);
    END IF;

    -- Add user_id to collections if not exists
    IF NOT EXISTS (SELECT 1 FROM information_schema.columns
                   WHERE table_schema = 'public'
                   AND table_name = 'collections'
                   AND column_name = 'user_id') THEN
        ALTER TABLE public.collections ADD COLUMN user_id UUID REFERENCES public.users(id);
    END IF;

    -- Add user_id to chat_logs if not exists
    IF NOT EXISTS (SELECT 1 FROM information_schema.columns
                   WHERE table_schema = 'public'
                   AND table_name = 'chat_logs'
                   AND column_name = 'user_id') THEN
        ALTER TABLE public.chat_logs ADD COLUMN user_id UUID REFERENCES public.users(id);
    END IF;
END $$;

-- Enable RLS on users table
ALTER TABLE public.users ENABLE ROW LEVEL SECURITY;

-- Drop existing policies if they exist (to avoid conflicts)
DROP POLICY IF EXISTS "Service role can manage agents" ON public.agents;
DROP POLICY IF EXISTS "Service role can manage agent versions" ON public.agent_versions;
DROP POLICY IF EXISTS "Service role can manage datasources" ON public.datasources;
DROP POLICY IF EXISTS "Service role can manage documents" ON public.documents;
DROP POLICY IF EXISTS "Service role can manage chunks" ON public.chunks;
DROP POLICY IF EXISTS "Service role can manage embeddings" ON public.embeddings;
DROP POLICY IF EXISTS "Service role can manage collections" ON public.collections;
DROP POLICY IF EXISTS "Service role can manage collection items" ON public.collection_items;
DROP POLICY IF EXISTS "Service role can manage pipeline runs" ON public.pipeline_runs;
DROP POLICY IF EXISTS "Service role can manage chat logs" ON public.chat_logs;

-- Users table policies
CREATE POLICY "Users can read their own profile" ON public.users
    FOR SELECT
    USING (auth.uid() = id OR auth.role() = 'service_role');

CREATE POLICY "Service role can manage all users" ON public.users
    FOR ALL
    USING (auth.role() = 'service_role')
    WITH CHECK (auth.role() = 'service_role');

-- Datasources policies (user-owned)
CREATE POLICY "Users can read their own datasources" ON public.datasources
    FOR SELECT
    USING (
        auth.role() = 'service_role' OR
        user_id = auth.uid() OR
        user_id IS NULL -- For backward compatibility
    );

CREATE POLICY "Users can create their own datasources" ON public.datasources
    FOR INSERT
    WITH CHECK (
        auth.role() = 'service_role' OR
        user_id = auth.uid()
    );

CREATE POLICY "Users can update their own datasources" ON public.datasources
    FOR UPDATE
    USING (
        auth.role() = 'service_role' OR
        user_id = auth.uid()
    );

CREATE POLICY "Users can delete their own datasources" ON public.datasources
    FOR DELETE
    USING (
        auth.role() = 'service_role' OR
        user_id = auth.uid()
    );

-- Documents policies (user-owned through datasource)
CREATE POLICY "Users can read documents from their datasources" ON public.documents
    FOR SELECT
    USING (
        auth.role() = 'service_role' OR
        user_id = auth.uid() OR
        user_id IS NULL OR
        EXISTS (
            SELECT 1 FROM public.datasources
            WHERE datasources.id = documents.datasource_id
            AND (datasources.user_id = auth.uid() OR datasources.user_id IS NULL)
        )
    );

CREATE POLICY "Service role can manage all documents" ON public.documents
    FOR ALL
    USING (auth.role() = 'service_role')
    WITH CHECK (auth.role() = 'service_role');

-- Chunks policies (read-only for users, managed by service)
CREATE POLICY "Users can read chunks from their documents" ON public.chunks
    FOR SELECT
    USING (
        auth.role() = 'service_role' OR
        EXISTS (
            SELECT 1 FROM public.documents
            WHERE documents.id = chunks.doc_id
            AND (
                documents.user_id = auth.uid() OR
                documents.user_id IS NULL OR
                EXISTS (
                    SELECT 1 FROM public.datasources
                    WHERE datasources.id = documents.datasource_id
                    AND (datasources.user_id = auth.uid() OR datasources.user_id IS NULL)
                )
            )
        )
    );

CREATE POLICY "Service role can manage all chunks" ON public.chunks
    FOR ALL
    USING (auth.role() = 'service_role')
    WITH CHECK (auth.role() = 'service_role');

-- Embeddings policies (read-only for users, managed by service)
CREATE POLICY "Users can read embeddings from their chunks" ON public.embeddings
    FOR SELECT
    USING (
        auth.role() = 'service_role' OR
        EXISTS (
            SELECT 1 FROM public.chunks
            JOIN public.documents ON chunks.doc_id = documents.id
            WHERE chunks.id = embeddings.chunk_id
            AND (
                documents.user_id = auth.uid() OR
                documents.user_id IS NULL OR
                EXISTS (
                    SELECT 1 FROM public.datasources
                    WHERE datasources.id = documents.datasource_id
                    AND (datasources.user_id = auth.uid() OR datasources.user_id IS NULL)
                )
            )
        )
    );

CREATE POLICY "Service role can manage all embeddings" ON public.embeddings
    FOR ALL
    USING (auth.role() = 'service_role')
    WITH CHECK (auth.role() = 'service_role');

-- Collections policies (user-owned)
CREATE POLICY "Users can manage their own collections" ON public.collections
    FOR ALL
    USING (
        auth.role() = 'service_role' OR
        user_id = auth.uid() OR
        user_id IS NULL
    )
    WITH CHECK (
        auth.role() = 'service_role' OR
        user_id = auth.uid()
    );

-- Collection items policies
CREATE POLICY "Users can manage items in their collections" ON public.collection_items
    FOR ALL
    USING (
        auth.role() = 'service_role' OR
        EXISTS (
            SELECT 1 FROM public.collections
            WHERE collections.id = collection_items.collection_id
            AND (collections.user_id = auth.uid() OR collections.user_id IS NULL)
        )
    )
    WITH CHECK (
        auth.role() = 'service_role' OR
        EXISTS (
            SELECT 1 FROM public.collections
            WHERE collections.id = collection_items.collection_id
            AND (collections.user_id = auth.uid() OR collections.user_id IS NULL)
        )
    );

-- Chat logs policies (user-owned)
CREATE POLICY "Users can manage their own chat logs" ON public.chat_logs
    FOR ALL
    USING (
        auth.role() = 'service_role' OR
        user_id = auth.uid() OR
        user_id IS NULL
    )
    WITH CHECK (
        auth.role() = 'service_role' OR
        user_id = auth.uid()
    );

-- Agents policies (read-only for all authenticated users)
CREATE POLICY "All authenticated users can read agents" ON public.agents
    FOR SELECT
    USING (auth.role() IN ('authenticated', 'service_role'));

CREATE POLICY "Service role can manage agents" ON public.agents
    FOR ALL
    USING (auth.role() = 'service_role')
    WITH CHECK (auth.role() = 'service_role');

-- Agent versions policies (read-only for all authenticated users)
CREATE POLICY "All authenticated users can read agent versions" ON public.agent_versions
    FOR SELECT
    USING (auth.role() IN ('authenticated', 'service_role'));

CREATE POLICY "Service role can manage agent versions" ON public.agent_versions
    FOR ALL
    USING (auth.role() = 'service_role')
    WITH CHECK (auth.role() = 'service_role');

-- Pipeline runs policies (service role only)
CREATE POLICY "Service role can manage pipeline runs" ON public.pipeline_runs
    FOR ALL
    USING (auth.role() = 'service_role')
    WITH CHECK (auth.role() = 'service_role');

-- Create indexes for better performance
CREATE INDEX IF NOT EXISTS idx_datasources_user_id ON public.datasources(user_id);
CREATE INDEX IF NOT EXISTS idx_documents_user_id ON public.documents(user_id);
CREATE INDEX IF NOT EXISTS idx_collections_user_id ON public.collections(user_id);
CREATE INDEX IF NOT EXISTS idx_chat_logs_user_id ON public.chat_logs(user_id);