-- Align experiment_repetitions schema with application expectations
-- Adds experiment_id and scheduled_date, maintains existing data, and updates constraints

-- 1) Add columns if missing and remove NOT NULL constraints from old columns
DO $$
BEGIN
    IF NOT EXISTS (
        SELECT 1 FROM information_schema.columns 
        WHERE table_schema = 'public' AND table_name = 'experiment_repetitions' AND column_name = 'experiment_id'
    ) THEN
        ALTER TABLE public.experiment_repetitions
        ADD COLUMN experiment_id UUID;
    END IF;

    IF NOT EXISTS (
        SELECT 1 FROM information_schema.columns 
        WHERE table_schema = 'public' AND table_name = 'experiment_repetitions' AND column_name = 'scheduled_date'
    ) THEN
        ALTER TABLE public.experiment_repetitions
        ADD COLUMN scheduled_date TIMESTAMPTZ NULL;
    END IF;

    IF NOT EXISTS (
        SELECT 1 FROM information_schema.columns 
        WHERE table_schema = 'public' AND table_name = 'experiment_repetitions' AND column_name = 'completion_status'
    ) THEN
        ALTER TABLE public.experiment_repetitions
        ADD COLUMN completion_status BOOLEAN NOT NULL DEFAULT false;
    END IF;

    IF NOT EXISTS (
        SELECT 1 FROM information_schema.columns 
        WHERE table_schema = 'public' AND table_name = 'experiment_repetitions' AND column_name = 'is_locked'
    ) THEN
        ALTER TABLE public.experiment_repetitions
        ADD COLUMN is_locked BOOLEAN NOT NULL DEFAULT false,
        ADD COLUMN locked_at TIMESTAMPTZ NULL,
        ADD COLUMN locked_by UUID NULL REFERENCES public.user_profiles(id);
    END IF;
    
    -- Remove NOT NULL constraints from old columns to allow new data insertion
    ALTER TABLE public.experiment_repetitions ALTER COLUMN experiment_number DROP NOT NULL;
    ALTER TABLE public.experiment_repetitions ALTER COLUMN experiment_phase_id DROP NOT NULL;
END $$;

-- 2) Backfill experiment_id by joining on (experiment_number, experiment_phase_id) -> experiments(id)
UPDATE public.experiment_repetitions er
SET experiment_id = e.id
FROM public.experiments e
WHERE er.experiment_id IS NULL
  AND e.experiment_number = er.experiment_number
  AND e.phase_id = er.experiment_phase_id;

-- 3) Create trigger to auto-populate experiment_id for inserts
CREATE OR REPLACE FUNCTION public.set_experiment_id_on_repetition()
RETURNS TRIGGER AS $func$
BEGIN
    IF NEW.experiment_id IS NULL THEN
        SELECT e.id INTO NEW.experiment_id
        FROM public.experiments e
        WHERE e.experiment_number = NEW.experiment_number
          AND e.phase_id = NEW.experiment_phase_id;
        
        -- If still NULL, raise an error with helpful message
        IF NEW.experiment_id IS NULL THEN
            RAISE EXCEPTION 'Could not find experiment with experiment_number=% and phase_id=%', 
                NEW.experiment_number, NEW.experiment_phase_id;
        END IF;
    END IF;
    RETURN NEW;
END;
$func$ LANGUAGE plpgsql;

DROP TRIGGER IF EXISTS trg_set_experiment_id_on_repetition ON public.experiment_repetitions;
CREATE TRIGGER trg_set_experiment_id_on_repetition
    BEFORE INSERT OR UPDATE OF experiment_number, experiment_phase_id ON public.experiment_repetitions
    FOR EACH ROW
    EXECUTE FUNCTION public.set_experiment_id_on_repetition();

-- 4) Add FK and not null once backfilled and trigger is in place
ALTER TABLE public.experiment_repetitions
    ADD CONSTRAINT experiment_repetitions_experiment_id_fkey
        FOREIGN KEY (experiment_id) REFERENCES public.experiments(id) ON DELETE CASCADE;

ALTER TABLE public.experiment_repetitions
    ALTER COLUMN experiment_id SET NOT NULL;

-- 5) Create indexes to support queries used in app
CREATE INDEX IF NOT EXISTS idx_experiment_repetitions_experiment_id ON public.experiment_repetitions(experiment_id);
CREATE INDEX IF NOT EXISTS idx_experiment_repetitions_repetition_number ON public.experiment_repetitions(repetition_number);
CREATE INDEX IF NOT EXISTS idx_experiment_repetitions_scheduled_date ON public.experiment_repetitions(scheduled_date);

-- 6) Maintain uniqueness: unique repetition_number per experiment_id
DO $$
BEGIN
    IF NOT EXISTS (
        SELECT 1 FROM pg_constraint 
        WHERE conname = 'uniq_experiment_repetition_number_per_experiment_id'
    ) THEN
        ALTER TABLE public.experiment_repetitions
        ADD CONSTRAINT uniq_experiment_repetition_number_per_experiment_id
        UNIQUE (experiment_id, repetition_number);
    END IF;
END $$;

-- 6) Optional: keep legacy uniqueness on (experiment_number, experiment_phase_id, repetition_number) if desired
-- This keeps backward compatibility with any existing references

-- 7) RLS already enabled; no change to policies necessary for added columns