-- Fix function security issues
-- Adds SECURITY DEFINER and proper search_path to all functions flagged by linter

-- Drop functions that have changed return types to avoid conflicts
DROP FUNCTION IF EXISTS get_user_subscription(UUID);
DROP FUNCTION IF EXISTS user_has_active_subscription(UUID);

-- 1. Fix cleanup_old_ai_operations function
CREATE OR REPLACE FUNCTION cleanup_old_ai_operations()
RETURNS INTEGER 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
DECLARE
    deleted_count INTEGER;
BEGIN
    -- Delete AI operations older than 30 days
    DELETE FROM ai_operations 
    WHERE created_at < NOW() - INTERVAL '30 days';
    
    GET DIAGNOSTICS deleted_count = ROW_COUNT;
    
    RETURN deleted_count;
END;
$$ LANGUAGE plpgsql;

-- 2. Fix set_billing_month function
CREATE OR REPLACE FUNCTION set_billing_month()
RETURNS TRIGGER 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
BEGIN
    NEW.billing_month := DATE_TRUNC('month', NEW.created_at)::DATE;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- 3. Fix get_monthly_ai_usage function
CREATE OR REPLACE FUNCTION get_monthly_ai_usage(p_user_id UUID, p_tier TEXT DEFAULT NULL)
RETURNS TABLE (
    tier TEXT,
    usage_count BIGINT,
    total_tokens BIGINT,
    total_cost DECIMAL
) 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
BEGIN
    RETURN QUERY
    SELECT 
        au.model_tier as tier,
        COUNT(*)::BIGINT as usage_count,
        COALESCE(SUM(au.total_tokens), 0)::BIGINT as total_tokens,
        COALESCE(SUM(au.cost_usd), 0)::DECIMAL as total_cost
    FROM ai_usage au
    WHERE au.user_id = p_user_id
        AND au.billing_month = DATE_TRUNC('month', NOW())
        AND au.success = true -- Only count successful requests
        AND (p_tier IS NULL OR au.model_tier = p_tier)
    GROUP BY au.model_tier;
END;
$$ LANGUAGE plpgsql;

-- 4. Fix get_user_usage_stats function
CREATE OR REPLACE FUNCTION get_user_usage_stats(p_user_id UUID)
RETURNS JSON 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
DECLARE
    result JSON;
BEGIN
    SELECT json_build_object(
        'current_month', json_build_object(
            'advanced', COALESCE((
                SELECT json_build_object(
                    'count', usage_count,
                    'tokens', total_tokens,
                    'cost', total_cost
                )
                FROM get_monthly_ai_usage(p_user_id, 'advanced')
                WHERE tier = 'advanced'
            ), json_build_object('count', 0, 'tokens', 0, 'cost', 0)),
            'auto', COALESCE((
                SELECT json_build_object(
                    'count', usage_count,
                    'tokens', total_tokens,
                    'cost', total_cost
                )
                FROM get_monthly_ai_usage(p_user_id, 'auto')
                WHERE tier = 'auto'
            ), json_build_object('count', 0, 'tokens', 0, 'cost', 0))
        ),
        'total_books', (
            SELECT COUNT(*) FROM books WHERE user_id = p_user_id
        ),
        'total_words', (
            SELECT COALESCE(SUM(word_count), 0) FROM books WHERE user_id = p_user_id
        )
    ) INTO result;
    
    RETURN result;
END;
$$ LANGUAGE plpgsql;

-- 5. Fix can_user_make_request function
CREATE OR REPLACE FUNCTION can_user_make_request(
    p_user_id UUID, 
    p_tier TEXT, 
    p_limit INTEGER
) RETURNS JSON 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
DECLARE
    current_usage BIGINT;
    result JSON;
BEGIN
    -- Get current month usage for the tier
    SELECT COALESCE(usage_count, 0) INTO current_usage
    FROM get_monthly_ai_usage(p_user_id, p_tier)
    WHERE tier = p_tier;
    
    -- Return usage info and whether request can proceed
    SELECT json_build_object(
        'can_proceed', (current_usage < p_limit OR p_limit = 0),
        'current_usage', current_usage,
        'limit', p_limit,
        'remaining', GREATEST(0, p_limit - current_usage)
    ) INTO result;
    
    RETURN result;
END;
$$ LANGUAGE plpgsql;

-- 6. Fix handle_new_user_email_preferences function
CREATE OR REPLACE FUNCTION handle_new_user_email_preferences()
RETURNS TRIGGER 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
BEGIN
    INSERT INTO email_preferences (user_id)
    VALUES (NEW.id);
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- 6a. Fix get_user_subscription function (recreated with proper security)
CREATE OR REPLACE FUNCTION get_user_subscription(user_uuid UUID)
RETURNS TABLE (
    id UUID,
    status subscription_status,
    price_id TEXT,
    current_period_end TIMESTAMPTZ,
    cancel_at_period_end BOOLEAN
)
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
BEGIN
    RETURN QUERY
    SELECT 
        s.id,
        s.status::subscription_status,
        s.price_id,
        s.current_period_end,
        s.cancel_at_period_end
    FROM subscriptions s
    WHERE s.user_id = user_uuid
    AND s.status IN ('active', 'trialing', 'past_due')
    ORDER BY s.created_at DESC
    LIMIT 1;
END;
$$ LANGUAGE plpgsql;

-- 6b. Fix user_has_active_subscription function (recreated with proper security)
CREATE OR REPLACE FUNCTION user_has_active_subscription(user_uuid UUID)
RETURNS BOOLEAN
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
BEGIN
    RETURN EXISTS(
        SELECT 1 
        FROM subscriptions 
        WHERE user_id = user_uuid 
        AND status IN ('active', 'trialing')
        AND (current_period_end IS NULL OR current_period_end > NOW())
    );
END;
$$ LANGUAGE plpgsql;

-- 7. Fix get_current_usage_period function
CREATE OR REPLACE FUNCTION get_current_usage_period(user_uuid UUID)
RETURNS TABLE (
    period_start TIMESTAMPTZ,
    period_end TIMESTAMPTZ
)
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
DECLARE
    subscription_record RECORD;
    current_time TIMESTAMPTZ := NOW();
BEGIN
    -- Get the user's active subscription
    SELECT * INTO subscription_record
    FROM subscriptions s
    WHERE s.user_id = user_uuid 
    AND s.status = 'active'
    AND s.current_period_end > current_time
    ORDER BY s.created_at DESC
    LIMIT 1;
    
    IF subscription_record IS NOT NULL THEN
        -- Return subscription period
        RETURN QUERY 
        SELECT 
            subscription_record.current_period_start,
            subscription_record.current_period_end;
    ELSE
        -- Return current month for free tier
        RETURN QUERY
        SELECT 
            DATE_TRUNC('month', current_time),
            (DATE_TRUNC('month', current_time) + INTERVAL '1 month' - INTERVAL '1 second');
    END IF;
END;
$$ LANGUAGE plpgsql;

-- 8. Fix get_or_create_current_usage function
CREATE OR REPLACE FUNCTION get_or_create_current_usage(user_uuid UUID)
RETURNS user_usage 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
DECLARE
    usage_period RECORD;
    usage_record user_usage;
BEGIN
    -- Get current usage period
    SELECT * INTO usage_period FROM get_current_usage_period(user_uuid);
    
    -- Try to find existing usage record
    SELECT * INTO usage_record
    FROM user_usage uu
    WHERE uu.user_id = user_uuid
    AND uu.period_start = usage_period.period_start
    AND uu.period_end = usage_period.period_end;
    
    -- Create if doesn't exist
    IF usage_record IS NULL THEN
        INSERT INTO user_usage (
            user_id,
            period_start,
            period_end
        ) VALUES (
            user_uuid,
            usage_period.period_start,
            usage_period.period_end
        ) RETURNING * INTO usage_record;
    END IF;
    
    RETURN usage_record;
END;
$$ LANGUAGE plpgsql;

-- 9. Fix increment_usage function
CREATE OR REPLACE FUNCTION increment_usage(
    user_uuid UUID,
    feature_name TEXT,
    increment_amount DECIMAL DEFAULT 1
)
RETURNS user_usage 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
DECLARE
    usage_record user_usage;
BEGIN
    -- Get or create current usage record
    SELECT * INTO usage_record FROM get_or_create_current_usage(user_uuid);
    
    -- Update the appropriate counter
    UPDATE user_usage 
    SET 
        advanced_prompts_used = CASE 
            WHEN feature_name = 'advanced_prompts' THEN advanced_prompts_used + increment_amount
            ELSE advanced_prompts_used 
        END,
        auto_prompts_used = CASE 
            WHEN feature_name = 'auto_prompts' THEN auto_prompts_used + increment_amount
            ELSE auto_prompts_used 
        END,
        storage_used_gb = CASE 
            WHEN feature_name = 'storage_gb' THEN storage_used_gb + increment_amount
            ELSE storage_used_gb 
        END,
        books_created = CASE 
            WHEN feature_name = 'books' THEN books_created + increment_amount
            ELSE books_created 
        END,
        custom_templates_used = CASE 
            WHEN feature_name = 'custom_templates' THEN custom_templates_used + increment_amount
            ELSE custom_templates_used 
        END,
        updated_at = NOW()
    WHERE id = usage_record.id
    RETURNING * INTO usage_record;
    
    -- Log the usage increment
    INSERT INTO usage_logs (user_id, feature_type, action, amount, metadata)
    VALUES (user_uuid, feature_name, 'increment', increment_amount, '{}'::jsonb);
    
    RETURN usage_record;
END;
$$ LANGUAGE plpgsql;

-- 10. Fix handle_new_user function
CREATE OR REPLACE FUNCTION public.handle_new_user()
RETURNS TRIGGER 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
BEGIN
    INSERT INTO public.profiles (
        id,
        full_name,
        avatar_url,
        email,
        created_at,
        updated_at
    ) VALUES (
        NEW.id,
        COALESCE(NEW.raw_user_meta_data->>'full_name', NEW.raw_user_meta_data->>'name'),
        NEW.raw_user_meta_data->>'avatar_url',
        NEW.email,
        NOW(),
        NOW()
    );
    
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- 11. Fix handle_user_update function
CREATE OR REPLACE FUNCTION public.handle_user_update()
RETURNS TRIGGER 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
BEGIN
    -- Only update if the profile exists
    IF EXISTS (SELECT 1 FROM public.profiles WHERE id = NEW.id) THEN
        UPDATE public.profiles
        SET
            full_name = COALESCE(
                NEW.raw_user_meta_data->>'full_name',
                NEW.raw_user_meta_data->>'name',
                full_name
            ),
            avatar_url = COALESCE(
                NEW.raw_user_meta_data->>'avatar_url',
                avatar_url
            ),
            email = COALESCE(NEW.email, email),
            updated_at = NOW()
        WHERE id = NEW.id;
    ELSE
        -- If profile doesn't exist, create it
        INSERT INTO public.profiles (
            id,
            full_name,
            avatar_url,
            email,
            created_at,
            updated_at
        ) VALUES (
            NEW.id,
            COALESCE(NEW.raw_user_meta_data->>'full_name', NEW.raw_user_meta_data->>'name'),
            NEW.raw_user_meta_data->>'avatar_url',
            NEW.email,
            NOW(),
            NOW()
        );
    END IF;
    
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- 12. Fix update_updated_at_column function
CREATE OR REPLACE FUNCTION update_updated_at_column()
RETURNS TRIGGER 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
BEGIN
    NEW.updated_at = NOW();
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- 13. Fix calculate_book_word_count function
CREATE OR REPLACE FUNCTION calculate_book_word_count(book_uuid UUID)
RETURNS INTEGER 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
DECLARE
    total_words INTEGER := 0;
    item_record RECORD;
BEGIN
    -- Sum word counts from all text files in the book
    FOR item_record IN 
        SELECT content 
        FROM file_system_items 
        WHERE book_id = book_uuid 
        AND item_type = 'file'
        AND file_extension IN ('txt', 'md', 'markdown')
        AND content IS NOT NULL
    LOOP
        -- Simple word count: split by whitespace and count non-empty elements
        total_words := total_words + (
            SELECT array_length(
                string_to_array(
                    regexp_replace(item_record.content, '\s+', ' ', 'g'), 
                    ' '
                ), 
                1
            )
        );
    END LOOP;
    
    RETURN COALESCE(total_words, 0);
END;
$$ LANGUAGE plpgsql;

-- 14. Fix update_book_word_count function
CREATE OR REPLACE FUNCTION update_book_word_count()
RETURNS TRIGGER 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
DECLARE
    book_uuid UUID;
    new_word_count INTEGER;
BEGIN
    -- Get the book_id for the file that was modified
    IF TG_OP = 'DELETE' THEN
        book_uuid := OLD.book_id;
    ELSE
        book_uuid := NEW.book_id;
    END IF;
    
    -- Only process if this is a text file
    IF (TG_OP = 'DELETE' AND OLD.item_type = 'file' AND OLD.file_extension IN ('txt', 'md', 'markdown')) OR
       (TG_OP IN ('INSERT', 'UPDATE') AND NEW.item_type = 'file' AND NEW.file_extension IN ('txt', 'md', 'markdown')) THEN
        
        -- Calculate new word count for the book
        new_word_count := calculate_book_word_count(book_uuid);
        
        -- Update the book's word count
        UPDATE books 
        SET word_count = new_word_count, 
            updated_at = NOW()
        WHERE id = book_uuid;
    END IF;
    
    -- Return appropriate record based on operation
    IF TG_OP = 'DELETE' THEN
        RETURN OLD;
    ELSE
        RETURN NEW;
    END IF;
END;
$$ LANGUAGE plpgsql;

-- 15. Fix update_chat_metadata function
CREATE OR REPLACE FUNCTION update_chat_metadata()
RETURNS TRIGGER 
SECURITY DEFINER
SET search_path = public, pg_temp
AS $$
BEGIN
    -- Update metadata when a message is added/updated/deleted
    IF TG_OP = 'INSERT' OR TG_OP = 'UPDATE' THEN
        UPDATE chats SET
            last_message_at = NOW(),
            message_count = (
                SELECT COUNT(*) 
                FROM chat_messages 
                WHERE chat_id = NEW.chat_id
            ),
            updated_at = NOW()
        WHERE id = NEW.chat_id;
        
        RETURN NEW;
    ELSIF TG_OP = 'DELETE' THEN
        UPDATE chats SET
            message_count = (
                SELECT COUNT(*) 
                FROM chat_messages 
                WHERE chat_id = OLD.chat_id
            ),
            updated_at = NOW()
        WHERE id = OLD.chat_id;
        
        RETURN OLD;
    END IF;
    
    RETURN NULL;
END;
$$ LANGUAGE plpgsql;

-- Add comments explaining the security changes
COMMENT ON FUNCTION cleanup_old_ai_operations() IS 'Cleans up old AI operations - secured with SECURITY DEFINER';
COMMENT ON FUNCTION set_billing_month() IS 'Sets billing month for AI usage records - secured with SECURITY DEFINER';
COMMENT ON FUNCTION get_monthly_ai_usage(UUID, TEXT) IS 'Gets AI usage stats for current month by tier - secured with SECURITY DEFINER';
COMMENT ON FUNCTION get_user_usage_stats(UUID) IS 'Gets comprehensive usage stats for dashboard - secured with SECURITY DEFINER';
COMMENT ON FUNCTION can_user_make_request(UUID, TEXT, INTEGER) IS 'Checks if user can make AI request within limits - secured with SECURITY DEFINER';
COMMENT ON FUNCTION handle_new_user_email_preferences() IS 'Creates email preferences for new users - secured with SECURITY DEFINER';
COMMENT ON FUNCTION get_current_usage_period(UUID) IS 'Gets current usage period for user - secured with SECURITY DEFINER';
COMMENT ON FUNCTION get_or_create_current_usage(UUID) IS 'Gets or creates current usage record - secured with SECURITY DEFINER';
COMMENT ON FUNCTION increment_usage(UUID, TEXT, DECIMAL) IS 'Increments usage counter for user - secured with SECURITY DEFINER';
COMMENT ON FUNCTION handle_new_user() IS 'Creates profile for new users - secured with SECURITY DEFINER';
COMMENT ON FUNCTION handle_user_update() IS 'Updates user profile when auth user changes - secured with SECURITY DEFINER';
COMMENT ON FUNCTION update_updated_at_column() IS 'Updates updated_at timestamp - secured with SECURITY DEFINER';
COMMENT ON FUNCTION calculate_book_word_count(UUID) IS 'Calculates word count for book - secured with SECURITY DEFINER';
COMMENT ON FUNCTION update_book_word_count() IS 'Updates book word count when files change - secured with SECURITY DEFINER';
COMMENT ON FUNCTION update_chat_metadata() IS 'Updates chat metadata when messages change - secured with SECURITY DEFINER';
COMMENT ON FUNCTION get_user_subscription(UUID) IS 'Gets user subscription details - secured with SECURITY DEFINER';
COMMENT ON FUNCTION user_has_active_subscription(UUID) IS 'Checks if user has active subscription - secured with SECURITY DEFINER';