--
-- Audit Log Partitioning Setup for LeoCRM
-- =========================================
--
-- This script creates a partitioned audit_log table with monthly partitions.
-- It includes functions for automatic partition creation and maintenance.
--
-- Usage:
--   psql -h localhost -U leocrm -d crm_db -f scripts/setup_audit_partitioning.sql
--
-- The script is idempotent — safe to run multiple times.
--

BEGIN;

-- ═══════════════════════════════════════════════════════════════
-- 1. Create the partitioned audit_log table
-- ═══════════════════════════════════════════════════════════════

-- Check if the partitioned table already exists
DO $$
BEGIN
    IF NOT EXISTS (
        SELECT 1 FROM pg_class WHERE relname = 'audit_log' AND relkind = 'p'
    ) THEN
        -- Check if there's an existing non-partitioned table
        IF EXISTS (SELECT 1 FROM pg_class WHERE relname = 'audit_log' AND relkind = 'r') THEN
            -- Rename existing table
            ALTER TABLE audit_log RENAME TO audit_log_old;
            RAISE NOTICE 'Renamed existing audit_log to audit_log_old';            
        END IF;

        -- Create the partitioned table
        CREATE TABLE audit_log (
            id UUID NOT NULL DEFAULT gen_random_uuid(),
            tenant_id UUID,
            user_id UUID,
            action VARCHAR(100) NOT NULL,
            entity_type VARCHAR(50),
            entity_id UUID,
            changes JSONB,
            ip_address VARCHAR(45),
            user_agent TEXT,
            created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
            PRIMARY KEY (id, created_at)
        ) PARTITION BY RANGE (created_at);

        RAISE NOTICE 'Created partitioned audit_log table';
    END IF;
END
$$;

-- ═══════════════════════════════════════════════════════════════
-- 2. Function: Create a single monthly partition
-- ═══════════════════════════════════════════════════════════════

CREATE OR REPLACE FUNCTION create_monthly_audit_partition(
    partition_date DATE DEFAULT CURRENT_DATE
)
RETURNS void
LANGUAGE plpgsql
AS $$
DECLARE
    year_str TEXT;
    month_str TEXT;
    partition_name TEXT;
    start_date TEXT;
    end_date TEXT;
BEGIN
    year_str := TO_CHAR(partition_date, 'YYYY');
    month_str := TO_CHAR(partition_date, 'MM');
    partition_name := 'audit_log_' || year_str || '_' || month_str;
    start_date := year_str || '-' || month_str || '-01';
    end_date := TO_CHAR(partition_date + INTERVAL '1 month', 'YYYY-MM-DD');

    -- Check if partition already exists
    IF NOT EXISTS (
        SELECT 1 FROM pg_class WHERE relname = partition_name
    ) THEN
        EXECUTE format('
            CREATE TABLE %I PARTITION OF audit_log
            FOR VALUES FROM (%L) TO (%L)',
            partition_name, start_date, end_date
        );

        -- Create indexes on the new partition
        EXECUTE format('
            CREATE INDEX %I ON %I (tenant_id)',
            'idx_' || partition_name || '_tenant', partition_name
        );
        EXECUTE format('
            CREATE INDEX %I ON %I (action)',
            'idx_' || partition_name || '_action', partition_name
        );
        EXECUTE format('
            CREATE INDEX %I ON %I (entity_type, entity_id)',
            'idx_' || partition_name || '_entity', partition_name
        );
        EXECUTE format('
            CREATE INDEX %I ON %I (created_at DESC)',
            'idx_' || partition_name || '_created', partition_name
        );

        RAISE NOTICE 'Created partition: %', partition_name;
    ELSE
        RAISE NOTICE 'Partition already exists: %', partition_name;
    END IF;
END;
$$;

-- ═══════════════════════════════════════════════════════════════
-- 3. Function: Create partitions for the next N months
-- ═══════════════════════════════════════════════════════════════

CREATE OR REPLACE FUNCTION create_future_audit_partitions(
    months_ahead INT DEFAULT 3
)
RETURNS void
LANGUAGE plpgsql
AS $$
DECLARE
    i INT;
    partition_date DATE;
BEGIN
    -- Create current month partition first
    partition_date := DATE_TRUNC('month', CURRENT_DATE)::DATE;
    PERFORM create_monthly_audit_partition(partition_date);

    -- Create future month partitions
    FOR i IN 1..months_ahead LOOP
        partition_date := (DATE_TRUNC('month', CURRENT_DATE) + (i || ' months')::INTERVAL)::DATE;
        PERFORM create_monthly_audit_partition(partition_date);
    END LOOP;

    RAISE NOTICE 'Created % future partitions (current + % months)', months_ahead + 1, months_ahead;
END;
$$;

-- ═══════════════════════════════════════════════════════════════
-- 4. Function: Drop old partitions beyond retention period
-- ═══════════════════════════════════════════════════════════════

CREATE OR REPLACE FUNCTION drop_old_audit_partitions(
    retention_months INT DEFAULT 12
)
RETURNS INT
LANGUAGE plpgsql
AS $$
DECLARE
    partition_record RECORD;
    drop_count INT := 0;
    cutoff_date DATE;
BEGIN
    cutoff_date := DATE_TRUNC('month', CURRENT_DATE - (retention_months || ' months')::INTERVAL)::DATE;

    FOR partition_record IN
        SELECT
            inhrelid::regclass AS partition_name,
            pg_get_expr(relpartbound, relid) AS partition_bound
        FROM pg_catalog.pg_inherits
        JOIN pg_class ON pg_class.oid = inhrelid
        WHERE inhparent = 'audit_log'::regclass
    LOOP
        -- Extract the upper bound date from the partition bound expression
        -- Format: FOR VALUES FROM ('2026-01-01') TO ('2026-02-01')
        IF partition_record.partition_bound ~ 'TO \(''([0-9]{4}-[0-9]{2}-[0-9]{2})' THEN
            DECLARE
                upper_date DATE;
            BEGIN
                upper_date := SUBSTRING(
                    partition_record.partition_bound
                    FROM 'TO \(''([0-9]{4}-[0-9]{2}-[0-9]{2})'
                )::DATE;

                IF upper_date <= cutoff_date THEN
                    EXECUTE format('DROP TABLE IF EXISTS %I', partition_record.partition_name);
                    drop_count := drop_count + 1;
                    RAISE NOTICE 'Dropped old partition: %', partition_record.partition_name;
                END IF;
            END;
        END IF;
    END LOOP;

    RETURN drop_count;
END;
$$;

-- ═══════════════════════════════════════════════════════════════
-- 5. Create initial partitions (current + next 3 months)
-- ═══════════════════════════════════════════════════════════════

SELECT create_future_audit_partitions(3);

-- ═══════════════════════════════════════════════════════════════
-- 6. Migrate existing data (if any)
-- ═══════════════════════════════════════════════════════════════

DO $$
DECLARE
    old_count BIGINT;
    new_count BIGINT;
BEGIN
    -- Check if old table exists and has data
    IF EXISTS (SELECT 1 FROM pg_class WHERE relname = 'audit_log_old') THEN
        EXECUTE 'SELECT COUNT(*) FROM audit_log_old' INTO old_count;

        IF old_count > 0 THEN
            RAISE NOTICE 'Migrating % rows from audit_log_old...', old_count;

            -- Insert data into partitioned table
            EXECUTE '
                INSERT INTO audit_log (
                    id, tenant_id, user_id, action, entity_type,
                    entity_id, changes, ip_address, user_agent, created_at
                )
                SELECT
                    id, tenant_id, user_id, action, entity_type,
                    entity_id, changes, ip_address, user_agent, created_at
                FROM audit_log_old
            ';

            EXECUTE 'SELECT COUNT(*) FROM audit_log' INTO new_count;
            RAISE NOTICE 'Migration complete: % rows migrated', new_count;

            -- Verify data integrity
            IF old_count = new_count THEN
                RAISE NOTICE 'Data integrity verified: % rows match', old_count;
            ELSE
                RAISE WARNING 'Data mismatch: old=% rows, new=% rows', old_count, new_count;
            END IF;
        ELSE
            RAISE NOTICE 'No data to migrate in audit_log_old';
        END IF;
    END IF;
END;
$$;

-- ═══════════════════════════════════════════════════════════════
-- 7. Verify setup
-- ═══════════════════════════════════════════════════════════════

DO $$
DECLARE
    partition_count INT;
    table_type TEXT;
BEGIN
    -- Check table type
    SELECT relkind INTO table_type FROM pg_class WHERE relname = 'audit_log';
    
    IF table_type = 'p' THEN
        RAISE NOTICE 'audit_log is correctly set up as a partitioned table';
    ELSE
        RAISE WARNING 'audit_log is NOT a partitioned table (relkind=%)', table_type;
    END IF;

    -- Count partitions
    SELECT COUNT(*) INTO partition_count
    FROM pg_catalog.pg_inherits
    WHERE inhparent = 'audit_log'::regclass;

    RAISE NOTICE 'Number of partitions: %', partition_count;
END;
$$;

COMMIT;

-- ═══════════════════════════════════════════════════════════════
-- Usage Examples
-- ═══════════════════════════════════════════════════════════════
--
-- Create partitions for the next 6 months:
--   SELECT create_future_audit_partitions(6);
--
-- Drop partitions older than 12 months:
--   SELECT drop_old_audit_partitions(12);
--
-- Create a specific month partition:
--   SELECT create_monthly_audit_partition('2026-07-01'::DATE);
--
-- Check partition sizes:
--   SELECT
--       relname AS partition_name,
--       pg_size_pretty(pg_total_relation_size(relid)) AS total_size
--   FROM pg_catalog.pg_statio_user_tables
--   WHERE relname LIKE 'audit_log_%'
--   ORDER BY relname;
--
-- Check row counts per partition:
--   SELECT
--       relname AS partition_name,
--       n_live_tup AS row_count
--   FROM pg_catalog.pg_stat_user_tables
--   WHERE relname LIKE 'audit_log_%'
--   ORDER BY relname;
--
-- Cron job (run on 1st of each month at 2 AM):
--   0 2 1 * * /usr/bin/psql -h localhost -U leocrm -d crm_db -c "SELECT create_future_audit_partitions(3);"
--