-- ============================================================================
-- Secure E2EE Chat System for Mutual Follows
-- ============================================================================
-- This migration creates the infrastructure for end-to-end encrypted messaging
-- using Signal Protocol concepts. Only mutual follows can exchange messages.
-- The server never sees plaintext - only encrypted blobs.
-- ============================================================================

-- ============================================================================
-- 1. Signal Protocol Key Storage
-- ============================================================================

-- Identity and pre-keys for Signal Protocol key exchange
CREATE TABLE IF NOT EXISTS signal_keys (
    user_id UUID PRIMARY KEY REFERENCES profiles(id) ON DELETE CASCADE,

    -- Identity Key (long-term, base64 encoded public key)
    identity_key_public TEXT NOT NULL,

    -- Signed Pre-Key (medium-term, rotated periodically)
    signed_prekey_public TEXT NOT NULL,
    signed_prekey_id INTEGER NOT NULL DEFAULT 1,
    signed_prekey_signature TEXT NOT NULL,

    -- One-Time Pre-Keys (for perfect forward secrecy, consumed on use)
    -- Stored as JSONB array: [{"id": 1, "key": "base64..."}, ...]
    one_time_prekeys JSONB DEFAULT '[]'::JSONB,

    -- Metadata
    created_at TIMESTAMPTZ DEFAULT NOW(),
    updated_at TIMESTAMPTZ DEFAULT NOW()
);

-- Index for fast key lookups
CREATE INDEX IF NOT EXISTS idx_signal_keys_user_id ON signal_keys(user_id);

-- ============================================================================
-- 2. Encrypted Conversations Metadata
-- ============================================================================

-- Conversation metadata (no content, just participants and state)
CREATE TABLE IF NOT EXISTS encrypted_conversations (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),

    -- Participants (always 2 for DM)
    participant_a UUID NOT NULL REFERENCES profiles(id) ON DELETE CASCADE,
    participant_b UUID NOT NULL REFERENCES profiles(id) ON DELETE CASCADE,

    -- Conversation state
    created_at TIMESTAMPTZ DEFAULT NOW(),
    last_message_at TIMESTAMPTZ DEFAULT NOW(),

    -- Ensure ordered participant storage (smaller UUID first)
    -- This prevents duplicate conversations
    CONSTRAINT ordered_participants CHECK (participant_a < participant_b),
    CONSTRAINT unique_conversation UNIQUE (participant_a, participant_b)
);

-- Indexes for conversation lookups
CREATE INDEX IF NOT EXISTS idx_conversations_participant_a ON encrypted_conversations(participant_a);
CREATE INDEX IF NOT EXISTS idx_conversations_participant_b ON encrypted_conversations(participant_b);
CREATE INDEX IF NOT EXISTS idx_conversations_last_message ON encrypted_conversations(last_message_at DESC);

-- ============================================================================
-- 3. Encrypted Messages
-- ============================================================================

-- Encrypted message storage - server sees ONLY ciphertext
CREATE TABLE IF NOT EXISTS encrypted_messages (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),

    -- Conversation reference
    conversation_id UUID NOT NULL REFERENCES encrypted_conversations(id) ON DELETE CASCADE,

    -- Sender (for routing, not content attribution)
    sender_id UUID NOT NULL REFERENCES profiles(id) ON DELETE CASCADE,

    -- Encrypted payload (what the server stores)
    -- This is the Signal Protocol message, completely opaque to server
    ciphertext BYTEA NOT NULL,

    -- Signal Protocol header (needed for decryption, but reveals nothing)
    -- Contains ephemeral key, previous chain length, message number
    message_header TEXT NOT NULL,

    -- Message type (for protocol handling)
    -- 1 = PreKeyWhisperMessage (initial message establishing session)
    -- 2 = WhisperMessage (subsequent messages in established session)
    message_type INTEGER NOT NULL DEFAULT 2,

    -- Delivery metadata
    created_at TIMESTAMPTZ DEFAULT NOW(),
    delivered_at TIMESTAMPTZ,
    read_at TIMESTAMPTZ,

    -- Expiration (optional ephemeral messaging)
    expires_at TIMESTAMPTZ
);

-- Indexes for message retrieval
CREATE INDEX IF NOT EXISTS idx_messages_conversation ON encrypted_messages(conversation_id, created_at DESC);
CREATE INDEX IF NOT EXISTS idx_messages_sender ON encrypted_messages(sender_id);
CREATE INDEX IF NOT EXISTS idx_messages_unread ON encrypted_messages(conversation_id, read_at) WHERE read_at IS NULL;
CREATE INDEX IF NOT EXISTS idx_messages_expiring ON encrypted_messages(expires_at) WHERE expires_at IS NOT NULL;

-- ============================================================================
-- 4. Helper Functions
-- ============================================================================

-- Check if two users have a mutual follow relationship
CREATE OR REPLACE FUNCTION is_mutual_follow(user_a UUID, user_b UUID)
RETURNS BOOLEAN AS $$
BEGIN
    RETURN EXISTS (
        SELECT 1 FROM follows f1
        WHERE f1.follower_id = user_a
        AND f1.following_id = user_b
    ) AND EXISTS (
        SELECT 1 FROM follows f2
        WHERE f2.follower_id = user_b
        AND f2.following_id = user_a
    );
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Get or create a conversation between two mutual follows
CREATE OR REPLACE FUNCTION get_or_create_conversation(user_a UUID, user_b UUID)
RETURNS UUID AS $$
DECLARE
    conv_id UUID;
    ordered_a UUID;
    ordered_b UUID;
BEGIN
    -- Verify mutual follow
    IF NOT is_mutual_follow(user_a, user_b) THEN
        RAISE EXCEPTION 'Users must have mutual follow to start conversation';
    END IF;

    -- Order participants for consistent storage
    IF user_a < user_b THEN
        ordered_a := user_a;
        ordered_b := user_b;
    ELSE
        ordered_a := user_b;
        ordered_b := user_a;
    END IF;

    -- Try to get existing conversation
    SELECT id INTO conv_id
    FROM encrypted_conversations
    WHERE participant_a = ordered_a AND participant_b = ordered_b;

    -- Create if doesn't exist
    IF conv_id IS NULL THEN
        INSERT INTO encrypted_conversations (participant_a, participant_b)
        VALUES (ordered_a, ordered_b)
        RETURNING id INTO conv_id;
    END IF;

    RETURN conv_id;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Consume a one-time pre-key (returns and removes it atomically)
CREATE OR REPLACE FUNCTION consume_one_time_prekey(target_user_id UUID)
RETURNS JSONB AS $$
DECLARE
    prekey JSONB;
    remaining JSONB;
BEGIN
    -- Get the first prekey
    SELECT one_time_prekeys->0 INTO prekey
    FROM signal_keys
    WHERE user_id = target_user_id
    AND jsonb_array_length(one_time_prekeys) > 0
    FOR UPDATE;

    IF prekey IS NULL THEN
        RETURN NULL;
    END IF;

    -- Remove it from the array
    UPDATE signal_keys
    SET one_time_prekeys = one_time_prekeys - 0,
        updated_at = NOW()
    WHERE user_id = target_user_id;

    RETURN prekey;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- ============================================================================
-- 5. Row Level Security Policies
-- ============================================================================

-- Enable RLS on all tables
ALTER TABLE signal_keys ENABLE ROW LEVEL SECURITY;
ALTER TABLE encrypted_conversations ENABLE ROW LEVEL SECURITY;
ALTER TABLE encrypted_messages ENABLE ROW LEVEL SECURITY;

-- Signal Keys: Users can only manage their own keys
CREATE POLICY signal_keys_select ON signal_keys
    FOR SELECT USING (true);  -- Anyone can read public keys

CREATE POLICY signal_keys_insert ON signal_keys
    FOR INSERT WITH CHECK (auth.uid() = user_id);

CREATE POLICY signal_keys_update ON signal_keys
    FOR UPDATE USING (auth.uid() = user_id);

CREATE POLICY signal_keys_delete ON signal_keys
    FOR DELETE USING (auth.uid() = user_id);

-- Conversations: Only participants can see their conversations
CREATE POLICY conversations_select ON encrypted_conversations
    FOR SELECT USING (
        auth.uid() = participant_a OR auth.uid() = participant_b
    );

-- Conversations are created via the get_or_create_conversation function
-- which enforces mutual follow, so we allow insert if user is a participant
CREATE POLICY conversations_insert ON encrypted_conversations
    FOR INSERT WITH CHECK (
        (auth.uid() = participant_a OR auth.uid() = participant_b)
        AND is_mutual_follow(participant_a, participant_b)
    );

-- Messages: Only conversation participants can see/send messages
CREATE POLICY messages_select ON encrypted_messages
    FOR SELECT USING (
        EXISTS (
            SELECT 1 FROM encrypted_conversations c
            WHERE c.id = conversation_id
            AND (c.participant_a = auth.uid() OR c.participant_b = auth.uid())
        )
    );

-- Critical: Messages can only be inserted by sender who is in a mutual follow
CREATE POLICY messages_insert ON encrypted_messages
    FOR INSERT WITH CHECK (
        auth.uid() = sender_id
        AND EXISTS (
            SELECT 1 FROM encrypted_conversations c
            WHERE c.id = conversation_id
            AND (c.participant_a = auth.uid() OR c.participant_b = auth.uid())
            AND is_mutual_follow(c.participant_a, c.participant_b)
        )
    );

-- Users can update their own sent messages (for read receipts on received messages)
CREATE POLICY messages_update ON encrypted_messages
    FOR UPDATE USING (
        EXISTS (
            SELECT 1 FROM encrypted_conversations c
            WHERE c.id = conversation_id
            AND (c.participant_a = auth.uid() OR c.participant_b = auth.uid())
        )
    );

-- ============================================================================
-- 6. Triggers for Metadata Updates
-- ============================================================================

-- Update conversation last_message_at when new message is inserted
CREATE OR REPLACE FUNCTION update_conversation_timestamp()
RETURNS TRIGGER AS $$
BEGIN
    UPDATE encrypted_conversations
    SET last_message_at = NEW.created_at
    WHERE id = NEW.conversation_id;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER trg_update_conversation_timestamp
    AFTER INSERT ON encrypted_messages
    FOR EACH ROW
    EXECUTE FUNCTION update_conversation_timestamp();

-- Auto-expire old messages (for ephemeral messaging)
CREATE OR REPLACE FUNCTION cleanup_expired_messages()
RETURNS void AS $$
BEGIN
    DELETE FROM encrypted_messages
    WHERE expires_at IS NOT NULL AND expires_at < NOW();
END;
$$ LANGUAGE plpgsql;

-- ============================================================================
-- 7. Realtime Subscriptions
-- ============================================================================

-- Enable realtime for messages (participants will subscribe to their conversations)
ALTER PUBLICATION supabase_realtime ADD TABLE encrypted_messages;

-- ============================================================================
-- 8. Comments for Documentation
-- ============================================================================

COMMENT ON TABLE signal_keys IS 'Public cryptographic keys for Signal Protocol key exchange. Private keys stored only on device.';
COMMENT ON TABLE encrypted_conversations IS 'Metadata for E2EE conversations. No message content stored here.';
COMMENT ON TABLE encrypted_messages IS 'Encrypted message blobs. Server cannot decrypt - only route.';
COMMENT ON FUNCTION is_mutual_follow IS 'Returns true if both users follow each other.';
COMMENT ON FUNCTION get_or_create_conversation IS 'Creates or retrieves a conversation, enforcing mutual follow requirement.';
COMMENT ON FUNCTION consume_one_time_prekey IS 'Atomically retrieves and removes a one-time pre-key for forward secrecy.';