-- Migration: Audit Log System (Enterprise-Ready v4)
-- Version: 014

-- Enable required extensions
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
CREATE EXTENSION IF NOT EXISTS "pgcrypto";

-- Create audit schema for RBAC, backup, and compliance
CREATE SCHEMA IF NOT EXISTS audit;

-- Create dedicated role for audit maintenance
DO $$
BEGIN
    IF NOT EXISTS (
        SELECT 1 FROM pg_roles WHERE rolname = 'audit_maintainer'
    ) THEN
        CREATE ROLE audit_maintainer;
    END IF;
END $$;

-- Severity enum
DO $$ BEGIN
    IF NOT EXISTS (SELECT 1 FROM pg_type WHERE typname = 'audit_severity') THEN
        CREATE TYPE audit_severity AS ENUM ('LOW', 'MEDIUM', 'HIGH', 'CRITICAL');
    END IF;
END $$;

-- Outbox pattern for transactional audit (true polling pattern)
CREATE TABLE IF NOT EXISTS audit.audit_outbox (
    id             UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    request_id     UUID NOT NULL,
    trace_id       UUID NOT NULL,
    event_type     VARCHAR(50) NOT NULL,
    payload        JSONB NOT NULL,
    payload_size   INTEGER CHECK (payload_size <= 65536),
    status         VARCHAR(20) DEFAULT 'PENDING',
    processed_at   TIMESTAMP WITH TIME ZONE,
    error_message  TEXT,
    retry_count    INTEGER DEFAULT 0,
    max_retries    INTEGER DEFAULT 3,
    failed_at      TIMESTAMP WITH TIME ZONE,
    created_at     TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    updated_at     TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    
    CONSTRAINT audit_outbox_status_check 
        CHECK (status IN ('PENDING', 'PROCESSING', 'COMPLETED', 'FAILED', 'RETRYING', 'FAILED_PERMANENT'))
);

-- Audit outbox indexes
CREATE INDEX idx_audit_outbox_status ON audit.audit_outbox(status);
CREATE INDEX idx_audit_outbox_status_created ON audit.audit_outbox(status, created_at);
CREATE INDEX idx_audit_outbox_retry ON audit.audit_outbox(status, retry_count) WHERE status = 'FAILED';
CREATE INDEX idx_audit_outbox_processing ON audit.audit_outbox(status, updated_at) WHERE status = 'PROCESSING';

-- Hot-Path Index (Critical Performance) - Partial index for poller's PENDING query
CREATE INDEX idx_audit_outbox_pending_created
ON audit.audit_outbox(status, created_at)
WHERE status = 'PENDING';

-- Main audit log table for all system operations (business + security)
CREATE TABLE IF NOT EXISTS audit.system_audit_logs (
    id                  UUID DEFAULT uuid_generate_v4(),
    outbox_event_id     UUID NOT NULL,
    request_id          UUID NOT NULL,
    trace_id            UUID NOT NULL,
    session_id          UUID,
    tenant_id           UUID,
    actor_id            UUID REFERENCES public.users(id) ON DELETE SET NULL,
    actor_name          VARCHAR(200),
    actor_email         VARCHAR(255),
    actor_role          VARCHAR(50),
    action              VARCHAR(50) NOT NULL,
    module              VARCHAR(50) NOT NULL,
    entity_id           UUID,
    entity_type         VARCHAR(50),
    batch_id            UUID,
    description         TEXT,
    severity            audit_severity DEFAULT 'LOW',
    ip_address          INET,
    user_agent          TEXT,
    payload             JSONB,
    payload_size        INTEGER,
    success             BOOLEAN DEFAULT TRUE,
    failure_reason      TEXT,
    event_version       INTEGER DEFAULT 1,
    created_at          TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    PRIMARY KEY (id, created_at),
    UNIQUE (outbox_event_id, created_at)
) PARTITION BY RANGE (created_at);

-- Create partitions for current and next 3 months (dynamic)
DO $$
DECLARE
    i INTEGER;
    start_date DATE;
    end_date DATE;
    partition_name TEXT;
BEGIN
    FOR i IN 0..3 LOOP
        start_date := date_trunc('month', CURRENT_DATE + (i || ' months')::interval);
        end_date := start_date + INTERVAL '1 month';
        partition_name := 'system_audit_logs_' || to_char(start_date, 'YYYY_MM');

        EXECUTE format(
            'CREATE TABLE IF NOT EXISTS audit.%I PARTITION OF audit.system_audit_logs
             FOR VALUES FROM (%L) TO (%L)',
            partition_name,
            start_date,
            end_date
        );
        RAISE NOTICE 'Created partition: audit.%', partition_name;
    END LOOP;
END $$;

-- Default partition for safety (catches data for months without explicit partitions)
CREATE TABLE audit.system_audit_logs_default PARTITION OF audit.system_audit_logs DEFAULT;

-- System audit logs indexes
CREATE INDEX idx_system_audit_request_id ON audit.system_audit_logs(request_id);
CREATE INDEX idx_system_audit_trace_id ON audit.system_audit_logs(trace_id);
CREATE INDEX idx_system_audit_actor_id ON audit.system_audit_logs(actor_id);
CREATE INDEX idx_system_audit_action ON audit.system_audit_logs(action);
CREATE INDEX idx_system_audit_module ON audit.system_audit_logs(module);
CREATE INDEX idx_system_audit_entity ON audit.system_audit_logs(entity_type, entity_id);
CREATE INDEX idx_system_audit_created_at ON audit.system_audit_logs(created_at DESC);
CREATE INDEX idx_system_audit_severity ON audit.system_audit_logs(severity);
CREATE INDEX idx_system_audit_batch_id ON audit.system_audit_logs(batch_id);
CREATE INDEX idx_system_audit_tenant_id ON audit.system_audit_logs(tenant_id) WHERE tenant_id IS NOT NULL;

-- Track field-level changes using snapshot diff (not row-per-field)
CREATE TABLE IF NOT EXISTS audit.entity_change_logs (
    id                  UUID DEFAULT uuid_generate_v4(),
    outbox_event_id     UUID NOT NULL,
    request_id          UUID NOT NULL,
    trace_id            UUID NOT NULL,
    entity_type         VARCHAR(50) NOT NULL,
    entity_id           UUID NOT NULL,
    actor_id            UUID REFERENCES public.users(id) ON DELETE SET NULL,
    actor_name          VARCHAR(200),
    actor_email         VARCHAR(255),
    change_type         VARCHAR(50) NOT NULL,
    payload             JSONB NOT NULL,
    payload_size        INTEGER,
    reason              TEXT,
    event_version       INTEGER DEFAULT 1,
    created_at          TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    PRIMARY KEY (id, created_at),
    UNIQUE (outbox_event_id, created_at)
) PARTITION BY RANGE (created_at);

-- Entity change logs indexes
CREATE INDEX idx_entity_change_entity ON audit.entity_change_logs(entity_type, entity_id);
CREATE INDEX idx_entity_change_request ON audit.entity_change_logs(request_id);
CREATE INDEX idx_entity_change_created ON audit.entity_change_logs(created_at DESC);
CREATE INDEX idx_entity_change_type ON audit.entity_change_logs(change_type);

-- Create partitions for entity_change_logs (dynamic)
DO $$
DECLARE
    i INTEGER;
    start_date DATE;
    end_date DATE;
    partition_name TEXT;
BEGIN
    FOR i IN 0..3 LOOP
        start_date := date_trunc('month', CURRENT_DATE + (i || ' months')::interval);
        end_date := start_date + INTERVAL '1 month';
        partition_name := 'entity_change_logs_' || to_char(start_date, 'YYYY_MM');

        EXECUTE format(
            'CREATE TABLE IF NOT EXISTS audit.%I PARTITION OF audit.entity_change_logs
             FOR VALUES FROM (%L) TO (%L)',
            partition_name,
            start_date,
            end_date
        );
        RAISE NOTICE 'Created partition: audit.%', partition_name;
    END LOOP;
END $$;

-- Default partition for safety
CREATE TABLE audit.entity_change_logs_default PARTITION OF audit.entity_change_logs DEFAULT;

-- Dead Letter Queue for failed audit events
CREATE TABLE IF NOT EXISTS audit.audit_dead_letter (
    id                  UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    original_event_id   UUID NOT NULL,
    event_type          VARCHAR(50) NOT NULL,
    original_payload    JSONB NOT NULL,
    error_message       TEXT NOT NULL,
    error_stack         TEXT,
    failed_at           TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    replayed            BOOLEAN DEFAULT FALSE,
    replayed_at         TIMESTAMP WITH TIME ZONE,
    replay_count        INTEGER DEFAULT 0,
    investigation_status VARCHAR(50) DEFAULT 'PENDING' CHECK (investigation_status IN ('PENDING', 'INVESTIGATING', 'RESOLVED', 'IGNORED')),
    resolution_notes    TEXT,
    created_at          TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP
);

-- Dead letter queue indexes
CREATE INDEX idx_dead_letter_event_type ON audit.audit_dead_letter(event_type);
CREATE INDEX idx_dead_letter_investigation_status ON audit.audit_dead_letter(investigation_status);
CREATE INDEX idx_dead_letter_failed_at ON audit.audit_dead_letter(failed_at DESC);
CREATE INDEX idx_dead_letter_replayed ON audit.audit_dead_letter(replayed, failed_at);

-- Prevent modification of audit logs (soft immutability - can be bypassed by dedicated role)
CREATE OR REPLACE FUNCTION audit.prevent_audit_modification()
RETURNS TRIGGER AS $$
BEGIN
    -- Recursion guard: skip if we're already in an audit override
    IF current_setting('app.audit_override', true) = 'on' THEN
        RETURN NEW;
    END IF;
    
    -- Check if caller has dedicated audit maintainer role or maintenance mode
    IF EXISTS (
        SELECT 1 
        FROM pg_roles 
        WHERE rolname = 'audit_maintainer' 
        AND pg_has_role(current_user, 'audit_maintainer', 'MEMBER')
    ) OR current_setting('app.maintenance_mode', true) = 'on' THEN
        -- Log the override for compliance
        INSERT INTO audit.system_audit_logs (
            outbox_event_id,
            request_id,
            trace_id,
            actor_id,
            actor_name,
            actor_email,
            action,
            module,
            description,
            severity,
            success,
            event_version
        ) VALUES (
            uuid_generate_v4(),
            uuid_generate_v4(),
            uuid_generate_v4(),
            NULL, -- current_user is text, not UUID
            current_user,
            NULL,
            'AUDIT_OVERRIDE',
            'AUDIT',
            'Audit log modification by ' || current_user,
            'CRITICAL',
            true,
            1
        );
        RETURN NEW;
    END IF;
    
    RAISE EXCEPTION 'Audit logs are immutable. Cannot modify % records', TG_TABLE_NAME;
END;
$$ LANGUAGE plpgsql;

-- Insert into outbox (called within transaction)
CREATE OR REPLACE FUNCTION audit.enqueue_audit_event(
    p_request_id UUID,
    p_trace_id UUID,
    p_event_type VARCHAR(50),
    p_payload JSONB
) RETURNS UUID AS $$
DECLARE
    event_id UUID;
    payload_size INTEGER;
BEGIN
    -- Calculate payload size
    payload_size := octet_length(p_payload::text);
    
    INSERT INTO audit.audit_outbox
        (request_id, trace_id, event_type, payload, payload_size)
    VALUES
        (p_request_id, p_trace_id, p_event_type, p_payload, payload_size)
    RETURNING id INTO event_id;
    RETURN event_id;
END;
$$ LANGUAGE plpgsql;

-- Recover stuck processing events (called by cron)
CREATE OR REPLACE FUNCTION audit.recover_stuck_processing_events()
RETURNS INTEGER AS $$
DECLARE
    recovered_count INTEGER;
BEGIN
    UPDATE audit.audit_outbox
    SET status = 'PENDING',
        updated_at = CURRENT_TIMESTAMP
    WHERE status = 'PROCESSING'
      AND updated_at < CURRENT_TIMESTAMP - INTERVAL '5 minutes';
      
    GET DIAGNOSTICS recovered_count = ROW_COUNT;
    RETURN recovered_count;
END;
$$ LANGUAGE plpgsql;

-- Move permanently failed events to DLQ
CREATE OR REPLACE FUNCTION audit.move_to_dead_letter()
RETURNS INTEGER AS $$
DECLARE
    moved_count INTEGER := 0;
    event_record RECORD;
BEGIN
    -- Move failed events to dead letter queue
    FOR event_record IN
        SELECT id, event_type, payload, error_message
        FROM audit.audit_outbox
        WHERE status = 'FAILED'
          AND retry_count >= max_retries
    LOOP
        INSERT INTO audit.audit_dead_letter (
            original_event_id,
            event_type,
            original_payload,
            error_message,
            failed_at
        ) VALUES (
            event_record.id,
            event_record.event_type,
            event_record.payload,
            event_record.error_message,
            CURRENT_TIMESTAMP
        );
        
        -- Mark as FAILED_PERMANENT in outbox
        UPDATE audit.audit_outbox
        SET status = 'FAILED_PERMANENT',
            failed_at = CURRENT_TIMESTAMP
        WHERE id = event_record.id;
        
        moved_count := moved_count + 1;
    END LOOP;
    
    RETURN moved_count;
END;
$$ LANGUAGE plpgsql;

-- Function to drop old partitions
CREATE OR REPLACE FUNCTION audit.drop_old_audit_partitions(months_to_keep INTEGER DEFAULT 24)
RETURNS INTEGER AS $$
DECLARE
    dropped_count INTEGER := 0;
    partition_name TEXT;
BEGIN
    -- Drop old system_audit_logs partitions
    FOR partition_name IN
        SELECT tablename
        FROM pg_tables
        WHERE schemaname = 'audit'
          AND tablename LIKE 'system_audit_logs_%'
          AND tablename < ('system_audit_logs_' || to_char(CURRENT_DATE - (months_to_keep || ' months')::interval, 'YYYY_MM'))
    LOOP
        EXECUTE 'DROP TABLE IF EXISTS audit.' || partition_name;
        dropped_count := dropped_count + 1;
    END LOOP;

    -- Drop old entity_change_logs partitions
    FOR partition_name IN
        SELECT tablename
        FROM pg_tables
        WHERE schemaname = 'audit'
          AND tablename LIKE 'entity_change_logs_%'
          AND tablename < ('entity_change_logs_' || to_char(CURRENT_DATE - (months_to_keep || ' months')::interval, 'YYYY_MM'))
    LOOP
        EXECUTE 'DROP TABLE IF EXISTS audit.' || partition_name;
        dropped_count := dropped_count + 1;
    END LOOP;

    RETURN dropped_count;
END;
$$ LANGUAGE plpgsql;

-- Function to cleanup audit_outbox (retention policy)
CREATE OR REPLACE FUNCTION audit.cleanup_audit_outbox()
RETURNS INTEGER AS $$
DECLARE
    deleted_count INTEGER;
BEGIN
    -- Delete COMPLETED events older than 7 days
    DELETE FROM audit.audit_outbox
    WHERE status = 'COMPLETED'
      AND processed_at < CURRENT_TIMESTAMP - INTERVAL '7 days';
    
    GET DIAGNOSTICS deleted_count = ROW_COUNT;
    RETURN deleted_count;
END;
$$ LANGUAGE plpgsql;

-- Payload versioning strategy for backward compatibility
CREATE OR REPLACE FUNCTION audit.upgrade_payload_version(p_payload JSONB, target_version INTEGER DEFAULT 2)
RETURNS JSONB AS $$
DECLARE
    current_version INTEGER;
BEGIN
    -- Get current version
    current_version := COALESCE(p_payload->'meta'->>'version', 1)::INTEGER;
    
    -- If already at target version, return as-is
    IF current_version >= target_version THEN
        RETURN p_payload;
    END IF;
    
    -- Upgrade v1 to v2
    IF current_version = 1 AND target_version >= 2 THEN
        -- Example: add new field
        p_payload := jsonb_set(
            p_payload,
            '{meta,audit_timestamp}',
            to_jsonb(CURRENT_TIMESTAMP)
        );
        
        -- Update version
        p_payload := jsonb_set(
            p_payload,
            '{meta,version}',
            to_jsonb(2)
        );
    END IF;
    
    RETURN p_payload;
END;
$$ LANGUAGE plpgsql;

-- Comments
COMMENT ON SCHEMA audit IS 'Audit schema for RBAC, backup, and compliance';
COMMENT ON TABLE audit.system_audit_logs IS 'Comprehensive audit log for all system operations (immutable)';
COMMENT ON TABLE audit.entity_change_logs IS 'Entity change tracking with snapshot diff (immutable)';
COMMENT ON TABLE audit.audit_outbox IS 'Staging table for transactional audit processing';
COMMENT ON TABLE audit.audit_dead_letter IS 'Dead letter queue for failed audit events';
COMMENT ON TYPE audit_severity IS 'Severity levels for audit events: LOW, MEDIUM, HIGH, CRITICAL';

-- LISTEN/NOTIFY Pattern (Realtime Processing)
-- Create channel for notifications (PostgreSQL automatically creates channels on first LISTEN/NOTIFY)
-- No explicit CREATE CHANNEL needed in PostgreSQL

-- Trigger to notify on new events
CREATE OR REPLACE FUNCTION audit.notify_outbox_insert()
RETURNS TRIGGER AS $$
BEGIN
    NOTIFY audit_channel;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER trigger_notify_outbox
    AFTER INSERT ON audit.audit_outbox
    FOR EACH STATEMENT EXECUTE FUNCTION audit.notify_outbox_insert();

-- SQL Function to create next month partition (for cron automation)
CREATE OR REPLACE FUNCTION audit.create_next_month_partition()
RETURNS TEXT AS $$
DECLARE
    partition_name TEXT;
    start_date DATE;
    end_date DATE;
    table_name TEXT;
BEGIN
    start_date := date_trunc('month', CURRENT_DATE + INTERVAL '1 month');
    end_date := start_date + INTERVAL '1 month';

    -- Create system_audit_logs partition
    table_name := 'system_audit_logs';
    partition_name := table_name || '_' || to_char(start_date, 'YYYY_MM');

    EXECUTE format('
        CREATE TABLE IF NOT EXISTS audit.%I PARTITION OF audit.%I
        FOR VALUES FROM (%L) TO (%L)
    ', partition_name, table_name, start_date, end_date);

    -- Create entity_change_logs partition
    table_name := 'entity_change_logs';
    partition_name := table_name || '_' || to_char(start_date, 'YYYY_MM');

    EXECUTE format('
        CREATE TABLE IF NOT EXISTS audit.%I PARTITION OF audit.%I
        FOR VALUES FROM (%L) TO (%L)
    ', partition_name, table_name, start_date, end_date);

    RETURN 'Created partitions for month: ' || to_char(start_date, 'YYYY-MM');
END;
$$ LANGUAGE plpgsql;

-- Auth Audit Logs (partitioned by month)
CREATE TABLE IF NOT EXISTS audit.auth_audit_logs (
    id                  UUID DEFAULT uuid_generate_v4(),
    outbox_event_id     UUID,
    request_id          UUID,
    trace_id            UUID,
    user_id             UUID REFERENCES public.users(id) ON DELETE SET NULL,
    action              VARCHAR(50) NOT NULL,
    ip_address          INET,
    user_agent          TEXT,
    device_info         JSONB,
    success             BOOLEAN DEFAULT TRUE,
    failure_reason      TEXT,
    session_id          UUID,
    event_version       INTEGER DEFAULT 1,
    created_at          TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    PRIMARY KEY (id, created_at)
) PARTITION BY RANGE (created_at);

-- Create partitions for auth_audit_logs (dynamic)
DO $$
DECLARE
    i INTEGER;
    start_date DATE;
    end_date DATE;
    partition_name TEXT;
BEGIN
    FOR i IN 0..3 LOOP
        start_date := date_trunc('month', CURRENT_DATE + (i || ' months')::interval);
        end_date := start_date + INTERVAL '1 month';
        partition_name := 'auth_audit_logs_' || to_char(start_date, 'YYYY_MM');

        EXECUTE format(
            'CREATE TABLE IF NOT EXISTS audit.%I PARTITION OF audit.auth_audit_logs
             FOR VALUES FROM (%L) TO (%L)',
            partition_name,
            start_date,
            end_date
        );
        RAISE NOTICE 'Created partition: audit.%', partition_name;
    END LOOP;
END $$;

-- Default partition for safety
CREATE TABLE audit.auth_audit_logs_default PARTITION OF audit.auth_audit_logs DEFAULT;

-- Refresh Token Audit Logs (partitioned by month)
CREATE TABLE IF NOT EXISTS audit.refresh_token_audit_logs (
    id                  UUID DEFAULT uuid_generate_v4(),
    outbox_event_id     UUID,
    request_id          UUID,
    trace_id            UUID,
    user_id             UUID REFERENCES public.users(id) ON DELETE RESTRICT,
    session_id          UUID,
    action              VARCHAR(50) NOT NULL,
    ip_address          INET,
    user_agent          TEXT,
    platform            VARCHAR(50) CHECK (platform IN ('web', 'mobile', 'zalo_miniapp', 'internal')),
    device_id           VARCHAR(255),
    previous_token_hash VARCHAR(255),
    new_token_hash      VARCHAR(255),
    reason              TEXT,
    event_version       INTEGER DEFAULT 1,
    created_at          TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
    PRIMARY KEY (id, created_at)
) PARTITION BY RANGE (created_at);

-- Create partitions for refresh_token_audit_logs (dynamic)
DO $$
DECLARE
    i INTEGER;
    start_date DATE;
    end_date DATE;
    partition_name TEXT;
BEGIN
    FOR i IN 0..3 LOOP
        start_date := date_trunc('month', CURRENT_DATE + (i || ' months')::interval);
        end_date := start_date + INTERVAL '1 month';
        partition_name := 'refresh_token_audit_logs_' || to_char(start_date, 'YYYY_MM');

        EXECUTE format(
            'CREATE TABLE IF NOT EXISTS audit.%I PARTITION OF audit.refresh_token_audit_logs
             FOR VALUES FROM (%L) TO (%L)',
            partition_name,
            start_date,
            end_date
        );
        RAISE NOTICE 'Created partition: audit.%', partition_name;
    END LOOP;
END $$;

-- Default partition for safety
CREATE TABLE audit.refresh_token_audit_logs_default PARTITION OF audit.refresh_token_audit_logs DEFAULT;

-- Refresh token audit logs indexes
CREATE INDEX idx_refresh_token_audit_user_id ON audit.refresh_token_audit_logs(user_id);
CREATE INDEX idx_refresh_token_audit_action ON audit.refresh_token_audit_logs(action);
CREATE INDEX idx_refresh_token_audit_created_at ON audit.refresh_token_audit_logs(created_at DESC);
CREATE INDEX idx_refresh_token_audit_session_id ON audit.refresh_token_audit_logs(session_id);

-- BRIN index for time-series scale (append-only optimization)
CREATE INDEX idx_system_audit_created_brin
ON audit.system_audit_logs USING BRIN(created_at);

CREATE INDEX idx_entity_change_created_brin
ON audit.entity_change_logs USING BRIN(created_at);

CREATE INDEX idx_auth_audit_created_brin
ON audit.auth_audit_logs USING BRIN(created_at);

CREATE INDEX idx_refresh_token_audit_created_brin
ON audit.refresh_token_audit_logs USING BRIN(created_at);







