-- Migration: Add training_status and active_training_task_id to training_datasets
-- Description: Track training status separately from dataset build status

-- Add training_status column
ALTER TABLE training_datasets
ADD COLUMN IF NOT EXISTS training_status VARCHAR(20) DEFAULT NULL;

-- Add active_training_task_id column
ALTER TABLE training_datasets
ADD COLUMN IF NOT EXISTS active_training_task_id UUID DEFAULT NULL;

-- Create index for training_status
CREATE INDEX IF NOT EXISTS idx_training_datasets_training_status
ON training_datasets(training_status);

-- Create index for active_training_task_id
CREATE INDEX IF NOT EXISTS idx_training_datasets_active_training_task_id
ON training_datasets(active_training_task_id);

-- Update existing datasets that have been used in completed training tasks to 'trained' status
UPDATE training_datasets d
SET status = 'trained'
WHERE d.status = 'ready'
AND EXISTS (
    SELECT 1 FROM training_tasks t
    WHERE t.dataset_id = d.dataset_id
    AND t.status = 'completed'
);