-- Partitioning Infrastructure Migration 001: Foundation
-- Sprint: SPRINT_3422_0001_0001 - Time-Based Partitioning
-- Category: C (infrastructure setup, requires planned maintenance)
--
-- Purpose: Create partition management infrastructure including:
-- - Helper functions for partition creation and maintenance
-- - Utility functions for BRIN index optimization
-- - Partition maintenance scheduling support
--
-- This migration creates the foundation; table conversion is done in separate migrations.

BEGIN;

-- ============================================================================
-- Step 1: Create partition management schema
-- ============================================================================

CREATE SCHEMA IF NOT EXISTS partition_mgmt;

COMMENT ON SCHEMA partition_mgmt IS
    'Partition management utilities for time-series tables';

-- ============================================================================
-- Step 2: Managed table registration
-- ============================================================================

CREATE TABLE IF NOT EXISTS partition_mgmt.managed_tables (
    schema_name TEXT NOT NULL,
    table_name TEXT NOT NULL,
    partition_key TEXT NOT NULL,
    partition_type TEXT NOT NULL,
    retention_months INT NOT NULL DEFAULT 0,
    months_ahead INT NOT NULL DEFAULT 3,
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    PRIMARY KEY (schema_name, table_name)
);

COMMENT ON TABLE partition_mgmt.managed_tables IS
    'Tracks partitioned tables with retention and creation settings';

-- ============================================================================
-- Step 3: Partition creation function
-- ============================================================================

-- Creates a new partition for a given table and date range
CREATE OR REPLACE FUNCTION partition_mgmt.create_partition(
    p_schema_name TEXT,
    p_table_name TEXT,
    p_partition_column TEXT,
    p_start_date DATE,
    p_end_date DATE,
    p_partition_suffix TEXT DEFAULT NULL
)
RETURNS TEXT
LANGUAGE plpgsql
AS $$
DECLARE
    v_partition_name TEXT;
    v_parent_table TEXT;
    v_sql TEXT;
BEGIN
    v_parent_table := format('%I.%I', p_schema_name, p_table_name);

    -- Generate partition name: tablename_YYYY_MM or tablename_YYYY_Q#
    IF p_partition_suffix IS NOT NULL THEN
        v_partition_name := format('%s_%s', p_table_name, p_partition_suffix);
    ELSE
        v_partition_name := format('%s_%s', p_table_name, to_char(p_start_date, 'YYYY_MM'));
    END IF;

    -- Check if partition already exists
    IF EXISTS (
        SELECT 1 FROM pg_class c
        JOIN pg_namespace n ON c.relnamespace = n.oid
        WHERE n.nspname = p_schema_name AND c.relname = v_partition_name
    ) THEN
        RAISE NOTICE 'Partition % already exists, skipping', v_partition_name;
        RETURN v_partition_name;
    END IF;

    -- Create partition
    v_sql := format(
        'CREATE TABLE %I.%I PARTITION OF %s FOR VALUES FROM (%L) TO (%L)',
        p_schema_name,
        v_partition_name,
        v_parent_table,
        p_start_date,
        p_end_date
    );

    EXECUTE v_sql;

    RAISE NOTICE 'Created partition %.%', p_schema_name, v_partition_name;
    RETURN v_partition_name;
END;
$$;

-- ============================================================================
-- Step 4: Monthly partition creation helper
-- ============================================================================

CREATE OR REPLACE FUNCTION partition_mgmt.create_monthly_partitions(
    p_schema_name TEXT,
    p_table_name TEXT,
    p_partition_column TEXT,
    p_start_month DATE,
    p_months_ahead INT DEFAULT 3
)
RETURNS SETOF TEXT
LANGUAGE plpgsql
AS $$
DECLARE
    v_current_month DATE;
    v_end_month DATE;
    v_partition_name TEXT;
BEGIN
    v_current_month := date_trunc('month', p_start_month)::DATE;
    v_end_month := date_trunc('month', NOW() + (p_months_ahead || ' months')::INTERVAL)::DATE;

    WHILE v_current_month <= v_end_month LOOP
        v_partition_name := partition_mgmt.create_partition(
            p_schema_name,
            p_table_name,
            p_partition_column,
            v_current_month,
            (v_current_month + INTERVAL '1 month')::DATE
        );
        RETURN NEXT v_partition_name;
        v_current_month := (v_current_month + INTERVAL '1 month')::DATE;
    END LOOP;
END;
$$;

-- ============================================================================
-- Step 5: Quarterly partition creation helper
-- ============================================================================

CREATE OR REPLACE FUNCTION partition_mgmt.create_quarterly_partitions(
    p_schema_name TEXT,
    p_table_name TEXT,
    p_partition_column TEXT,
    p_start_quarter DATE,
    p_quarters_ahead INT DEFAULT 2
)
RETURNS SETOF TEXT
LANGUAGE plpgsql
AS $$
DECLARE
    v_current_quarter DATE;
    v_end_quarter DATE;
    v_partition_name TEXT;
    v_suffix TEXT;
BEGIN
    v_current_quarter := date_trunc('quarter', p_start_quarter)::DATE;
    v_end_quarter := date_trunc('quarter', NOW() + (p_quarters_ahead * 3 || ' months')::INTERVAL)::DATE;

    WHILE v_current_quarter <= v_end_quarter LOOP
        -- Generate suffix like 2025_Q1, 2025_Q2, etc.
        v_suffix := to_char(v_current_quarter, 'YYYY') || '_Q' ||
                    EXTRACT(QUARTER FROM v_current_quarter)::TEXT;

        v_partition_name := partition_mgmt.create_partition(
            p_schema_name,
            p_table_name,
            p_partition_column,
            v_current_quarter,
            (v_current_quarter + INTERVAL '3 months')::DATE,
            v_suffix
        );
        RETURN NEXT v_partition_name;
        v_current_quarter := (v_current_quarter + INTERVAL '3 months')::DATE;
    END LOOP;
END;
$$;

-- ============================================================================
-- Step 6: Ensure future partitions exist
-- ============================================================================

CREATE OR REPLACE FUNCTION partition_mgmt.ensure_future_partitions(
    p_schema_name TEXT,
    p_table_name TEXT,
    p_months_ahead INT
)
RETURNS INT
LANGUAGE plpgsql
AS $$
DECLARE
    v_partition_key TEXT;
    v_partition_type TEXT;
    v_months_ahead INT;
    v_created INT := 0;
    v_current DATE;
    v_end DATE;
    v_suffix TEXT;
    v_partition_name TEXT;
BEGIN
    SELECT partition_key, partition_type, months_ahead
      INTO v_partition_key, v_partition_type, v_months_ahead
      FROM partition_mgmt.managed_tables
     WHERE schema_name = p_schema_name
       AND table_name = p_table_name;

    IF v_partition_key IS NULL THEN
        RETURN 0;
    END IF;

    IF p_months_ahead IS NOT NULL AND p_months_ahead > 0 THEN
        v_months_ahead := p_months_ahead;
    END IF;

    IF v_months_ahead IS NULL OR v_months_ahead <= 0 THEN
        RETURN 0;
    END IF;

    v_partition_type := lower(coalesce(v_partition_type, 'monthly'));

    IF v_partition_type = 'monthly' THEN
        v_current := date_trunc('month', NOW())::DATE;
        v_end := date_trunc('month', NOW() + (v_months_ahead || ' months')::INTERVAL)::DATE;

        WHILE v_current <= v_end LOOP
            v_partition_name := format('%s_%s', p_table_name, to_char(v_current, 'YYYY_MM'));
            IF NOT EXISTS (
                SELECT 1 FROM pg_class c
                JOIN pg_namespace n ON c.relnamespace = n.oid
                WHERE n.nspname = p_schema_name AND c.relname = v_partition_name
            ) THEN
                PERFORM partition_mgmt.create_partition(
                    p_schema_name,
                    p_table_name,
                    v_partition_key,
                    v_current,
                    (v_current + INTERVAL '1 month')::DATE
                );
                v_created := v_created + 1;
            END IF;

            v_current := (v_current + INTERVAL '1 month')::DATE;
        END LOOP;
    ELSIF v_partition_type = 'quarterly' THEN
        v_current := date_trunc('quarter', NOW())::DATE;
        v_end := date_trunc('quarter', NOW() + (v_months_ahead || ' months')::INTERVAL)::DATE;

        WHILE v_current <= v_end LOOP
            v_suffix := to_char(v_current, 'YYYY') || '_Q' ||
                        EXTRACT(QUARTER FROM v_current)::TEXT;
            v_partition_name := format('%s_%s', p_table_name, v_suffix);

            IF NOT EXISTS (
                SELECT 1 FROM pg_class c
                JOIN pg_namespace n ON c.relnamespace = n.oid
                WHERE n.nspname = p_schema_name AND c.relname = v_partition_name
            ) THEN
                PERFORM partition_mgmt.create_partition(
                    p_schema_name,
                    p_table_name,
                    v_partition_key,
                    v_current,
                    (v_current + INTERVAL '3 months')::DATE,
                    v_suffix
                );
                v_created := v_created + 1;
            END IF;

            v_current := (v_current + INTERVAL '3 months')::DATE;
        END LOOP;
    END IF;

    RETURN v_created;
END;
$$;

-- ============================================================================
-- Step 7: Retention enforcement function
-- ============================================================================

CREATE OR REPLACE FUNCTION partition_mgmt.enforce_retention(
    p_schema_name TEXT,
    p_table_name TEXT,
    p_retention_months INT
)
RETURNS INT
LANGUAGE plpgsql
AS $$
DECLARE
    v_retention_months INT;
    v_cutoff_date DATE;
    v_partition RECORD;
    v_dropped INT := 0;
BEGIN
    SELECT retention_months
      INTO v_retention_months
      FROM partition_mgmt.managed_tables
     WHERE schema_name = p_schema_name
       AND table_name = p_table_name;

    IF p_retention_months IS NOT NULL AND p_retention_months > 0 THEN
        v_retention_months := p_retention_months;
    END IF;

    IF v_retention_months IS NULL OR v_retention_months <= 0 THEN
        RETURN 0;
    END IF;

    v_cutoff_date := (NOW() - (v_retention_months || ' months')::INTERVAL)::DATE;

    FOR v_partition IN
        SELECT partition_name, partition_end
          FROM partition_mgmt.partition_stats
         WHERE schema_name = p_schema_name
           AND table_name = p_table_name
    LOOP
        IF v_partition.partition_end IS NOT NULL AND v_partition.partition_end < v_cutoff_date THEN
            EXECUTE format('DROP TABLE IF EXISTS %I.%I', p_schema_name, v_partition.partition_name);
            v_dropped := v_dropped + 1;
        END IF;
    END LOOP;

    RETURN v_dropped;
END;
$$;

-- ============================================================================
-- Step 8: Partition detach and archive function
-- ============================================================================

CREATE OR REPLACE FUNCTION partition_mgmt.detach_partition(
    p_schema_name TEXT,
    p_table_name TEXT,
    p_partition_name TEXT,
    p_archive_schema TEXT DEFAULT 'archive'
)
RETURNS BOOLEAN
LANGUAGE plpgsql
AS $$
DECLARE
    v_parent_table TEXT;
    v_partition_full TEXT;
    v_archive_table TEXT;
BEGIN
    v_parent_table := format('%I.%I', p_schema_name, p_table_name);
    v_partition_full := format('%I.%I', p_schema_name, p_partition_name);
    v_archive_table := format('%I.%I', p_archive_schema, p_partition_name);

    -- Create archive schema if not exists
    EXECUTE format('CREATE SCHEMA IF NOT EXISTS %I', p_archive_schema);

    -- Detach partition
    EXECUTE format(
        'ALTER TABLE %s DETACH PARTITION %s',
        v_parent_table,
        v_partition_full
    );

    -- Move to archive schema
    EXECUTE format(
        'ALTER TABLE %s SET SCHEMA %I',
        v_partition_full,
        p_archive_schema
    );

    RAISE NOTICE 'Detached and archived partition % to %', p_partition_name, v_archive_table;
    RETURN TRUE;
EXCEPTION
    WHEN OTHERS THEN
        RAISE WARNING 'Failed to detach partition %: %', p_partition_name, SQLERRM;
        RETURN FALSE;
END;
$$;

-- ============================================================================
-- Step 9: Partition retention cleanup function
-- ============================================================================

CREATE OR REPLACE FUNCTION partition_mgmt.cleanup_old_partitions(
    p_schema_name TEXT,
    p_table_name TEXT,
    p_retention_months INT,
    p_archive_schema TEXT DEFAULT 'archive',
    p_dry_run BOOLEAN DEFAULT TRUE
)
RETURNS TABLE(partition_name TEXT, action TEXT)
LANGUAGE plpgsql
AS $$
DECLARE
    v_cutoff_date DATE;
    v_partition RECORD;
    v_partition_end DATE;
BEGIN
    v_cutoff_date := (NOW() - (p_retention_months || ' months')::INTERVAL)::DATE;

    FOR v_partition IN
        SELECT c.relname as name,
               pg_get_expr(c.relpartbound, c.oid) as bound_expr
        FROM pg_class c
        JOIN pg_namespace n ON c.relnamespace = n.oid
        JOIN pg_inherits i ON c.oid = i.inhrelid
        JOIN pg_class parent ON i.inhparent = parent.oid
        WHERE n.nspname = p_schema_name
          AND parent.relname = p_table_name
          AND c.relkind = 'r'
    LOOP
        -- Parse the partition bound to get end date
        -- Format: FOR VALUES FROM ('2024-01-01') TO ('2024-02-01')
        v_partition_end := (regexp_match(v_partition.bound_expr,
            'TO \(''([^'']+)''\)'))[1]::DATE;

        IF v_partition_end IS NOT NULL AND v_partition_end < v_cutoff_date THEN
            partition_name := v_partition.name;

            IF p_dry_run THEN
                action := 'WOULD_ARCHIVE';
            ELSE
                IF partition_mgmt.detach_partition(
                    p_schema_name, p_table_name, v_partition.name, p_archive_schema
                ) THEN
                    action := 'ARCHIVED';
                ELSE
                    action := 'FAILED';
                END IF;
            END IF;

            RETURN NEXT;
        END IF;
    END LOOP;
END;
$$;

-- ============================================================================
-- Step 10: Partition statistics view
-- ============================================================================

CREATE OR REPLACE VIEW partition_mgmt.partition_stats AS
SELECT
    n.nspname AS schema_name,
    parent.relname AS table_name,
    c.relname AS partition_name,
    pg_get_expr(c.relpartbound, c.oid) AS partition_range,
    (regexp_match(pg_get_expr(c.relpartbound, c.oid), 'FROM \(''([^'']+)''\)'))[1]::DATE AS partition_start,
    (regexp_match(pg_get_expr(c.relpartbound, c.oid), 'TO \(''([^'']+)''\)'))[1]::DATE AS partition_end,
    pg_size_pretty(pg_relation_size(c.oid)) AS size,
    pg_relation_size(c.oid) AS size_bytes,
    COALESCE(s.n_live_tup, 0) AS estimated_rows,
    s.last_vacuum,
    s.last_autovacuum,
    s.last_analyze,
    s.last_autoanalyze
FROM pg_class c
JOIN pg_namespace n ON c.relnamespace = n.oid
JOIN pg_inherits i ON c.oid = i.inhrelid
JOIN pg_class parent ON i.inhparent = parent.oid
LEFT JOIN pg_stat_user_tables s ON c.oid = s.relid
WHERE c.relkind = 'r'
  AND parent.relkind = 'p'
ORDER BY n.nspname, parent.relname, c.relname;

COMMENT ON VIEW partition_mgmt.partition_stats IS
    'Statistics for all partitioned tables in the database';

-- ============================================================================
-- Step 11: BRIN index optimization helper
-- ============================================================================

CREATE OR REPLACE FUNCTION partition_mgmt.create_brin_index_if_not_exists(
    p_schema_name TEXT,
    p_table_name TEXT,
    p_column_name TEXT,
    p_pages_per_range INT DEFAULT 128
)
RETURNS BOOLEAN
LANGUAGE plpgsql
AS $$
DECLARE
    v_index_name TEXT;
    v_sql TEXT;
BEGIN
    v_index_name := format('brin_%s_%s', p_table_name, p_column_name);

    -- Check if index exists
    IF EXISTS (
        SELECT 1 FROM pg_indexes
        WHERE schemaname = p_schema_name AND indexname = v_index_name
    ) THEN
        RAISE NOTICE 'BRIN index % already exists', v_index_name;
        RETURN FALSE;
    END IF;

    v_sql := format(
        'CREATE INDEX %I ON %I.%I USING brin (%I) WITH (pages_per_range = %s)',
        v_index_name,
        p_schema_name,
        p_table_name,
        p_column_name,
        p_pages_per_range
    );

    EXECUTE v_sql;

    RAISE NOTICE 'Created BRIN index % on %.%(%)',
        v_index_name, p_schema_name, p_table_name, p_column_name;
    RETURN TRUE;
END;
$$;

-- ============================================================================
-- Step 12: Maintenance job tracking table
-- ============================================================================

CREATE TABLE IF NOT EXISTS partition_mgmt.maintenance_log (
    id BIGSERIAL PRIMARY KEY,
    operation TEXT NOT NULL,
    schema_name TEXT NOT NULL,
    table_name TEXT NOT NULL,
    partition_name TEXT,
    status TEXT NOT NULL DEFAULT 'started',
    details JSONB NOT NULL DEFAULT '{}',
    started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    completed_at TIMESTAMPTZ,
    error_message TEXT
);

CREATE INDEX idx_maintenance_log_table ON partition_mgmt.maintenance_log(schema_name, table_name);
CREATE INDEX idx_maintenance_log_status ON partition_mgmt.maintenance_log(status, started_at);

-- ============================================================================
-- Step 13: Archive schema for detached partitions
-- ============================================================================

CREATE SCHEMA IF NOT EXISTS archive;

COMMENT ON SCHEMA archive IS
    'Storage for detached/archived partitions awaiting deletion or offload';

COMMIT;

-- ============================================================================
-- Usage Examples (commented out)
-- ============================================================================

/*
-- Create monthly partitions for audit table, 3 months ahead
SELECT partition_mgmt.create_monthly_partitions(
    'scheduler', 'audit', 'created_at', '2024-01-01'::DATE, 3
);

-- Preview old partitions that would be archived (dry run)
SELECT * FROM partition_mgmt.cleanup_old_partitions(
    'scheduler', 'audit', 12, 'archive', TRUE
);

-- Actually archive old partitions
SELECT * FROM partition_mgmt.cleanup_old_partitions(
    'scheduler', 'audit', 12, 'archive', FALSE
);

-- View partition statistics
SELECT * FROM partition_mgmt.partition_stats
WHERE schema_name = 'scheduler'
ORDER BY table_name, partition_name;
*/
