-- ============================================================================
-- E2EE Session State Tracking
-- ============================================================================
-- Server-side session state tracking to detect and recover from session
-- mismatches between parties. The server cannot see session keys, only
-- metadata about whether sessions exist.
-- ============================================================================

-- ============================================================================
-- 1. Session State Table
-- ============================================================================

-- Track which users have established sessions with each other
-- This allows detection of asymmetric session states (one party has session, other doesn't)
CREATE TABLE IF NOT EXISTS e2ee_session_state (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),

    -- Session participants (always stored with user_id < peer_id for consistency)
    user_id UUID NOT NULL REFERENCES profiles(id) ON DELETE CASCADE,
    peer_id UUID NOT NULL REFERENCES profiles(id) ON DELETE CASCADE,

    -- Session state flags
    user_has_session BOOLEAN NOT NULL DEFAULT FALSE,
    peer_has_session BOOLEAN NOT NULL DEFAULT FALSE,

    -- Session metadata (no actual keys stored!)
    user_session_created_at TIMESTAMPTZ,
    peer_session_created_at TIMESTAMPTZ,

    -- Version tracking for conflict resolution
    user_session_version INTEGER NOT NULL DEFAULT 0,
    peer_session_version INTEGER NOT NULL DEFAULT 0,

    -- Timestamps
    created_at TIMESTAMPTZ DEFAULT NOW(),
    updated_at TIMESTAMPTZ DEFAULT NOW(),

    -- Ensure unique pair (user_id should always be < peer_id)
    CONSTRAINT e2ee_session_state_pair_unique UNIQUE (user_id, peer_id),
    CONSTRAINT e2ee_session_state_ordering CHECK (user_id < peer_id)
);

-- Indexes for efficient lookups
CREATE INDEX IF NOT EXISTS idx_session_state_user ON e2ee_session_state(user_id);
CREATE INDEX IF NOT EXISTS idx_session_state_peer ON e2ee_session_state(peer_id);
CREATE INDEX IF NOT EXISTS idx_session_state_mismatch ON e2ee_session_state(user_has_session, peer_has_session)
    WHERE user_has_session != peer_has_session;

-- ============================================================================
-- 2. RLS Policies
-- ============================================================================

ALTER TABLE e2ee_session_state ENABLE ROW LEVEL SECURITY;

-- Users can see session state for their own sessions
CREATE POLICY session_state_select_own ON e2ee_session_state
    FOR SELECT USING (auth.uid() = user_id OR auth.uid() = peer_id);

-- Users can only insert/update their own side of the session
-- (handled via function to enforce ordering)

-- ============================================================================
-- 3. Helper Functions
-- ============================================================================

-- Function to update session state (handles ordering automatically)
CREATE OR REPLACE FUNCTION update_e2ee_session_state(
    p_user_id UUID,
    p_peer_id UUID,
    p_has_session BOOLEAN
) RETURNS JSONB AS $$
DECLARE
    v_lower_id UUID;
    v_higher_id UUID;
    v_is_user_lower BOOLEAN;
    v_result JSONB;
    v_session_state RECORD;
BEGIN
    -- Determine ordering (user_id < peer_id constraint)
    IF p_user_id < p_peer_id THEN
        v_lower_id := p_user_id;
        v_higher_id := p_peer_id;
        v_is_user_lower := TRUE;
    ELSE
        v_lower_id := p_peer_id;
        v_higher_id := p_user_id;
        v_is_user_lower := FALSE;
    END IF;

    -- Insert or update
    INSERT INTO e2ee_session_state (user_id, peer_id, user_has_session, peer_has_session, user_session_created_at, peer_session_created_at, user_session_version, peer_session_version)
    VALUES (
        v_lower_id,
        v_higher_id,
        CASE WHEN v_is_user_lower THEN p_has_session ELSE FALSE END,
        CASE WHEN NOT v_is_user_lower THEN p_has_session ELSE FALSE END,
        CASE WHEN v_is_user_lower AND p_has_session THEN NOW() ELSE NULL END,
        CASE WHEN NOT v_is_user_lower AND p_has_session THEN NOW() ELSE NULL END,
        CASE WHEN v_is_user_lower THEN 1 ELSE 0 END,
        CASE WHEN NOT v_is_user_lower THEN 1 ELSE 0 END
    )
    ON CONFLICT (user_id, peer_id) DO UPDATE SET
        user_has_session = CASE
            WHEN v_is_user_lower THEN p_has_session
            ELSE e2ee_session_state.user_has_session
        END,
        peer_has_session = CASE
            WHEN NOT v_is_user_lower THEN p_has_session
            ELSE e2ee_session_state.peer_has_session
        END,
        user_session_created_at = CASE
            WHEN v_is_user_lower AND p_has_session AND e2ee_session_state.user_session_created_at IS NULL THEN NOW()
            WHEN v_is_user_lower AND NOT p_has_session THEN NULL
            ELSE e2ee_session_state.user_session_created_at
        END,
        peer_session_created_at = CASE
            WHEN NOT v_is_user_lower AND p_has_session AND e2ee_session_state.peer_session_created_at IS NULL THEN NOW()
            WHEN NOT v_is_user_lower AND NOT p_has_session THEN NULL
            ELSE e2ee_session_state.peer_session_created_at
        END,
        user_session_version = CASE
            WHEN v_is_user_lower THEN e2ee_session_state.user_session_version + 1
            ELSE e2ee_session_state.user_session_version
        END,
        peer_session_version = CASE
            WHEN NOT v_is_user_lower THEN e2ee_session_state.peer_session_version + 1
            ELSE e2ee_session_state.peer_session_version
        END,
        updated_at = NOW()
    RETURNING * INTO v_session_state;

    -- Build result with mismatch detection
    v_result := jsonb_build_object(
        'success', TRUE,
        'user_has_session', CASE WHEN v_is_user_lower THEN v_session_state.user_has_session ELSE v_session_state.peer_has_session END,
        'peer_has_session', CASE WHEN v_is_user_lower THEN v_session_state.peer_has_session ELSE v_session_state.user_has_session END,
        'session_mismatch', v_session_state.user_has_session != v_session_state.peer_has_session,
        'peer_session_version', CASE WHEN v_is_user_lower THEN v_session_state.peer_session_version ELSE v_session_state.user_session_version END
    );

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

-- Function to get session state between two users
CREATE OR REPLACE FUNCTION get_e2ee_session_state(
    p_user_id UUID,
    p_peer_id UUID
) RETURNS JSONB AS $$
DECLARE
    v_lower_id UUID;
    v_higher_id UUID;
    v_is_user_lower BOOLEAN;
    v_session_state RECORD;
BEGIN
    -- Determine ordering
    IF p_user_id < p_peer_id THEN
        v_lower_id := p_user_id;
        v_higher_id := p_peer_id;
        v_is_user_lower := TRUE;
    ELSE
        v_lower_id := p_peer_id;
        v_higher_id := p_user_id;
        v_is_user_lower := FALSE;
    END IF;

    SELECT * INTO v_session_state
    FROM e2ee_session_state
    WHERE user_id = v_lower_id AND peer_id = v_higher_id;

    IF NOT FOUND THEN
        RETURN jsonb_build_object(
            'exists', FALSE,
            'user_has_session', FALSE,
            'peer_has_session', FALSE,
            'session_mismatch', FALSE
        );
    END IF;

    RETURN jsonb_build_object(
        'exists', TRUE,
        'user_has_session', CASE WHEN v_is_user_lower THEN v_session_state.user_has_session ELSE v_session_state.peer_has_session END,
        'peer_has_session', CASE WHEN v_is_user_lower THEN v_session_state.peer_has_session ELSE v_session_state.user_has_session END,
        'session_mismatch', v_session_state.user_has_session != v_session_state.peer_has_session,
        'user_session_version', CASE WHEN v_is_user_lower THEN v_session_state.user_session_version ELSE v_session_state.peer_session_version END,
        'peer_session_version', CASE WHEN v_is_user_lower THEN v_session_state.peer_session_version ELSE v_session_state.user_session_version END
    );
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Function to clear session state (when resetting)
CREATE OR REPLACE FUNCTION clear_e2ee_session_state(
    p_user_id UUID,
    p_peer_id UUID
) RETURNS JSONB AS $$
DECLARE
    v_lower_id UUID;
    v_higher_id UUID;
    v_is_user_lower BOOLEAN;
BEGIN
    -- Determine ordering
    IF p_user_id < p_peer_id THEN
        v_lower_id := p_user_id;
        v_higher_id := p_peer_id;
        v_is_user_lower := TRUE;
    ELSE
        v_lower_id := p_peer_id;
        v_higher_id := p_user_id;
        v_is_user_lower := FALSE;
    END IF;

    -- Update only the caller's side of the session
    UPDATE e2ee_session_state SET
        user_has_session = CASE WHEN v_is_user_lower THEN FALSE ELSE user_has_session END,
        peer_has_session = CASE WHEN NOT v_is_user_lower THEN FALSE ELSE peer_has_session END,
        user_session_created_at = CASE WHEN v_is_user_lower THEN NULL ELSE user_session_created_at END,
        peer_session_created_at = CASE WHEN NOT v_is_user_lower THEN NULL ELSE peer_session_created_at END,
        user_session_version = CASE WHEN v_is_user_lower THEN user_session_version + 1 ELSE user_session_version END,
        peer_session_version = CASE WHEN NOT v_is_user_lower THEN peer_session_version + 1 ELSE peer_session_version END,
        updated_at = NOW()
    WHERE user_id = v_lower_id AND peer_id = v_higher_id;

    RETURN jsonb_build_object('success', TRUE);
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- ============================================================================
-- 4. Add decryption_failure event type if not exists
-- ============================================================================

-- Update the check constraint to include session_mismatch event type
ALTER TABLE e2ee_session_events
    DROP CONSTRAINT IF EXISTS e2ee_session_events_event_type_check;

ALTER TABLE e2ee_session_events
    ADD CONSTRAINT e2ee_session_events_event_type_check
    CHECK (event_type IN ('session_reset', 'conversation_cleanup', 'key_refresh', 'decryption_failure', 'session_mismatch', 'session_established'));

-- ============================================================================
-- 5. Realtime for session state changes
-- ============================================================================

ALTER PUBLICATION supabase_realtime ADD TABLE e2ee_session_state;

-- ============================================================================
-- 6. Comments
-- ============================================================================

COMMENT ON TABLE e2ee_session_state IS 'Server-side tracking of E2EE session existence between user pairs. Does NOT store actual keys.';
COMMENT ON COLUMN e2ee_session_state.user_has_session IS 'Whether the user with smaller UUID has an active session';
COMMENT ON COLUMN e2ee_session_state.peer_has_session IS 'Whether the user with larger UUID has an active session';
COMMENT ON COLUMN e2ee_session_state.user_session_version IS 'Incremented each time user updates their session state';
COMMENT ON FUNCTION update_e2ee_session_state IS 'Update session state for a user-peer pair. Handles UUID ordering automatically.';
COMMENT ON FUNCTION get_e2ee_session_state IS 'Get session state between two users. Returns mismatch detection.';
COMMENT ON FUNCTION clear_e2ee_session_state IS 'Clear session state for a user (used during session reset).';