-- ================================================================================================
-- Migration 046: Batch Job Tracking Schema
-- Support for batch job orchestration with parent-child job hierarchy and progress aggregation
-- ================================================================================================

-- ================================================================================================
-- BATCH JOBS TABLE
-- Parent jobs that orchestrate multiple child training jobs
-- ================================================================================================
CREATE TABLE IF NOT EXISTS batch_jobs (
    -- Primary identifiers
    id UUID PRIMARY KEY,
    
    -- Batch job metadata
    name VARCHAR(255) NOT NULL,
    description TEXT,
    
    -- Status tracking
    status VARCHAR(50) NOT NULL DEFAULT 'Pending',
    
    -- Progress aggregation
    total_jobs INTEGER NOT NULL DEFAULT 0,
    pending_jobs INTEGER NOT NULL DEFAULT 0,
    running_jobs INTEGER NOT NULL DEFAULT 0,
    completed_jobs INTEGER NOT NULL DEFAULT 0,
    failed_jobs INTEGER NOT NULL DEFAULT 0,
    overall_progress DOUBLE PRECISION NOT NULL DEFAULT 0.0,
    
    -- Configuration
    config_json JSONB NOT NULL DEFAULT '{}'::jsonb,
    
    -- Audit timestamps
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    started_at TIMESTAMPTZ,
    completed_at TIMESTAMPTZ,
    
    -- Constraints
    CONSTRAINT chk_batch_job_status CHECK (
        status IN ('Pending', 'Running', 'Completed', 'Failed', 'Stopped', 'Paused')
    ),
    CONSTRAINT chk_batch_job_counts CHECK (
        total_jobs >= 0 AND
        pending_jobs >= 0 AND
        running_jobs >= 0 AND
        completed_jobs >= 0 AND
        failed_jobs >= 0 AND
        (pending_jobs + running_jobs + completed_jobs + failed_jobs) <= total_jobs
    ),
    CONSTRAINT chk_batch_job_progress CHECK (
        overall_progress >= 0.0 AND overall_progress <= 100.0
    )
);

-- Add table comment
COMMENT ON TABLE batch_jobs IS 'Parent batch jobs that orchestrate multiple child training jobs with progress aggregation';

-- Add column comments
COMMENT ON COLUMN batch_jobs.id IS 'Unique batch job identifier';
COMMENT ON COLUMN batch_jobs.name IS 'Human-readable batch job name';
COMMENT ON COLUMN batch_jobs.status IS 'Current batch job status (Pending, Running, Completed, Failed, Stopped, Paused)';
COMMENT ON COLUMN batch_jobs.total_jobs IS 'Total number of child jobs in this batch';
COMMENT ON COLUMN batch_jobs.pending_jobs IS 'Number of pending child jobs';
COMMENT ON COLUMN batch_jobs.running_jobs IS 'Number of running child jobs';
COMMENT ON COLUMN batch_jobs.completed_jobs IS 'Number of completed child jobs';
COMMENT ON COLUMN batch_jobs.failed_jobs IS 'Number of failed child jobs';
COMMENT ON COLUMN batch_jobs.overall_progress IS 'Weighted overall progress (0.0-100.0)';
COMMENT ON COLUMN batch_jobs.config_json IS 'Batch job configuration and metadata';

-- ================================================================================================
-- CHILD JOBS TABLE
-- Individual training jobs that belong to a batch
-- ================================================================================================
CREATE TABLE IF NOT EXISTS child_jobs (
    -- Primary identifiers
    id UUID PRIMARY KEY,
    batch_id UUID NOT NULL REFERENCES batch_jobs(id) ON DELETE CASCADE,
    
    -- Job metadata
    model_type VARCHAR(50) NOT NULL,
    model_weight DOUBLE PRECISION NOT NULL DEFAULT 0.25,
    
    -- Status tracking
    status VARCHAR(50) NOT NULL DEFAULT 'Pending',
    
    -- Progress tracking
    current_epoch INTEGER NOT NULL DEFAULT 0,
    total_epochs INTEGER NOT NULL DEFAULT 0,
    progress_pct DOUBLE PRECISION NOT NULL DEFAULT 0.0,
    
    -- Configuration
    config_json JSONB NOT NULL DEFAULT '{}'::jsonb,
    
    -- Error tracking
    error_message TEXT,
    
    -- Audit timestamps
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    started_at TIMESTAMPTZ,
    completed_at TIMESTAMPTZ,
    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    
    -- Constraints
    CONSTRAINT chk_child_job_status CHECK (
        status IN ('Pending', 'Running', 'Completed', 'Failed', 'Stopped', 'Paused')
    ),
    CONSTRAINT chk_child_job_model_type CHECK (
        model_type IN ('DQN', 'PPO', 'MAMBA-2', 'TFT', 'TLOB')
    ),
    CONSTRAINT chk_child_job_weight CHECK (
        model_weight >= 0.0 AND model_weight <= 1.0
    ),
    CONSTRAINT chk_child_job_progress CHECK (
        progress_pct >= 0.0 AND progress_pct <= 100.0
    ),
    CONSTRAINT chk_child_job_epochs CHECK (
        current_epoch >= 0 AND
        total_epochs >= 0 AND
        current_epoch <= total_epochs
    )
);

-- Add table comment
COMMENT ON TABLE child_jobs IS 'Individual child training jobs that belong to a parent batch job';

-- Add column comments
COMMENT ON COLUMN child_jobs.id IS 'Unique child job identifier';
COMMENT ON COLUMN child_jobs.batch_id IS 'Parent batch job identifier';
COMMENT ON COLUMN child_jobs.model_type IS 'Type of ML model (DQN, PPO, MAMBA-2, TFT, TLOB)';
COMMENT ON COLUMN child_jobs.model_weight IS 'Weight for progress aggregation (0.0-1.0)';
COMMENT ON COLUMN child_jobs.status IS 'Current job status (Pending, Running, Completed, Failed, Stopped, Paused)';
COMMENT ON COLUMN child_jobs.current_epoch IS 'Current training epoch (0-based)';
COMMENT ON COLUMN child_jobs.total_epochs IS 'Total number of training epochs';
COMMENT ON COLUMN child_jobs.progress_pct IS 'Individual job progress (0.0-100.0)';
COMMENT ON COLUMN child_jobs.config_json IS 'Child job configuration and metadata';

-- ================================================================================================
-- HIGH-PERFORMANCE INDEXES
-- ================================================================================================

-- Index for batch job queries
CREATE INDEX IF NOT EXISTS idx_batch_jobs_status
    ON batch_jobs(status);

CREATE INDEX IF NOT EXISTS idx_batch_jobs_created_at
    ON batch_jobs(created_at DESC);

-- Index for child job queries
CREATE INDEX IF NOT EXISTS idx_child_jobs_batch_id
    ON child_jobs(batch_id);

CREATE INDEX IF NOT EXISTS idx_child_jobs_status
    ON child_jobs(status);

CREATE INDEX IF NOT EXISTS idx_child_jobs_model_type
    ON child_jobs(model_type);

-- Composite index for batch summary queries
CREATE INDEX IF NOT EXISTS idx_child_jobs_batch_status
    ON child_jobs(batch_id, status);

-- GIN indexes for JSONB queries
CREATE INDEX IF NOT EXISTS idx_batch_jobs_config_gin
    ON batch_jobs USING GIN (config_json);

CREATE INDEX IF NOT EXISTS idx_child_jobs_config_gin
    ON child_jobs USING GIN (config_json);

-- ================================================================================================
-- TRIGGER FUNCTIONS FOR DATA INTEGRITY
-- ================================================================================================

-- Function to update child job updated_at timestamp
CREATE OR REPLACE FUNCTION update_child_jobs_timestamp()
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at := NOW();
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- Trigger to auto-update updated_at on modification
CREATE TRIGGER tg_child_jobs_update_timestamp
    BEFORE UPDATE ON child_jobs
    FOR EACH ROW
    EXECUTE FUNCTION update_child_jobs_timestamp();

-- Function to update batch job counts and progress
CREATE OR REPLACE FUNCTION update_batch_job_aggregates()
RETURNS TRIGGER AS $$
DECLARE
    v_batch_id UUID;
    v_total_jobs INTEGER;
    v_pending_jobs INTEGER;
    v_running_jobs INTEGER;
    v_completed_jobs INTEGER;
    v_failed_jobs INTEGER;
    v_weighted_progress DOUBLE PRECISION;
    v_new_status VARCHAR(50);
BEGIN
    -- Determine which batch to update
    v_batch_id := COALESCE(NEW.batch_id, OLD.batch_id);
    
    -- Count jobs by status
    SELECT
        COUNT(*),
        COUNT(*) FILTER (WHERE status = 'Pending'),
        COUNT(*) FILTER (WHERE status = 'Running'),
        COUNT(*) FILTER (WHERE status = 'Completed'),
        COUNT(*) FILTER (WHERE status = 'Failed')
    INTO
        v_total_jobs,
        v_pending_jobs,
        v_running_jobs,
        v_completed_jobs,
        v_failed_jobs
    FROM child_jobs
    WHERE batch_id = v_batch_id;
    
    -- Calculate weighted progress
    SELECT COALESCE(SUM(progress_pct * model_weight), 0.0)
    INTO v_weighted_progress
    FROM child_jobs
    WHERE batch_id = v_batch_id;
    
    -- Determine batch status
    IF v_failed_jobs > 0 AND v_running_jobs = 0 AND v_pending_jobs = 0 THEN
        v_new_status := 'Failed';
    ELSIF v_completed_jobs = v_total_jobs AND v_total_jobs > 0 THEN
        v_new_status := 'Completed';
    ELSIF v_running_jobs > 0 THEN
        v_new_status := 'Running';
    ELSE
        v_new_status := 'Pending';
    END IF;
    
    -- Update batch job
    UPDATE batch_jobs
    SET
        total_jobs = v_total_jobs,
        pending_jobs = v_pending_jobs,
        running_jobs = v_running_jobs,
        completed_jobs = v_completed_jobs,
        failed_jobs = v_failed_jobs,
        overall_progress = v_weighted_progress,
        status = v_new_status,
        started_at = CASE
            WHEN started_at IS NULL AND v_running_jobs > 0 THEN NOW()
            ELSE started_at
        END,
        completed_at = CASE
            WHEN v_new_status IN ('Completed', 'Failed') THEN NOW()
            ELSE NULL
        END
    WHERE id = v_batch_id;
    
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- Trigger to update batch aggregates on child job changes
CREATE TRIGGER tg_child_jobs_update_batch_aggregates
    AFTER INSERT OR UPDATE OR DELETE ON child_jobs
    FOR EACH ROW
    EXECUTE FUNCTION update_batch_job_aggregates();

-- ================================================================================================
-- ANALYTICAL VIEWS FOR REPORTING
-- ================================================================================================

-- View for active batch jobs
CREATE OR REPLACE VIEW v_active_batch_jobs AS
SELECT
    b.id,
    b.name,
    b.description,
    b.status,
    b.total_jobs,
    b.pending_jobs,
    b.running_jobs,
    b.completed_jobs,
    b.failed_jobs,
    b.overall_progress,
    b.created_at,
    b.started_at,
    b.completed_at,
    EXTRACT(EPOCH FROM (COALESCE(b.completed_at, NOW()) - b.created_at)) as duration_seconds
FROM batch_jobs b
WHERE b.status IN ('Pending', 'Running')
ORDER BY b.created_at DESC;

COMMENT ON VIEW v_active_batch_jobs IS 'Active batch jobs (Pending or Running) with progress metrics';

-- View for batch job details with child job breakdown
CREATE OR REPLACE VIEW v_batch_job_details AS
SELECT
    b.id as batch_id,
    b.name as batch_name,
    b.status as batch_status,
    b.overall_progress,
    c.id as child_job_id,
    c.model_type,
    c.model_weight,
    c.status as child_status,
    c.progress_pct,
    c.current_epoch,
    c.total_epochs,
    c.error_message
FROM batch_jobs b
LEFT JOIN child_jobs c ON b.id = c.batch_id
ORDER BY b.created_at DESC, c.model_type;

COMMENT ON VIEW v_batch_job_details IS 'Detailed view of batch jobs with all child jobs';

-- ================================================================================================
-- FINAL VALIDATION
-- ================================================================================================

DO $$
BEGIN
    IF NOT EXISTS (
        SELECT 1 FROM information_schema.tables
        WHERE table_name = 'batch_jobs'
    ) THEN
        RAISE EXCEPTION 'Migration 046 failed: batch_jobs table not created';
    END IF;
    
    IF NOT EXISTS (
        SELECT 1 FROM information_schema.tables
        WHERE table_name = 'child_jobs'
    ) THEN
        RAISE EXCEPTION 'Migration 046 failed: child_jobs table not created';
    END IF;
    
    RAISE NOTICE 'Migration 046 completed successfully: Batch job tracking schema created';
END $$;