-- Fix phase tables to use experiment_id instead of composite key
-- This aligns the schema with application expectations

BEGIN;

-- 1) Add experiment_id column to all phase tables and remove NOT NULL constraints from old columns
DO $$
BEGIN
    -- Soaking table
    IF NOT EXISTS (
        SELECT 1 FROM information_schema.columns 
        WHERE table_schema = 'public' AND table_name = 'soaking' AND column_name = 'experiment_id'
    ) THEN
        ALTER TABLE public.soaking ADD COLUMN experiment_id UUID;
    END IF;
    
    -- Remove NOT NULL constraints from old columns to allow new data insertion
    ALTER TABLE public.soaking ALTER COLUMN experiment_number DROP NOT NULL;
    ALTER TABLE public.soaking ALTER COLUMN experiment_phase_id DROP NOT NULL;

    -- Airdrying table
    IF NOT EXISTS (
        SELECT 1 FROM information_schema.columns 
        WHERE table_schema = 'public' AND table_name = 'airdrying' AND column_name = 'experiment_id'
    ) THEN
        ALTER TABLE public.airdrying ADD COLUMN experiment_id UUID;
    END IF;
    
    ALTER TABLE public.airdrying ALTER COLUMN experiment_number DROP NOT NULL;
    ALTER TABLE public.airdrying ALTER COLUMN experiment_phase_id DROP NOT NULL;

    -- Cracking table
    IF NOT EXISTS (
        SELECT 1 FROM information_schema.columns 
        WHERE table_schema = 'public' AND table_name = 'cracking' AND column_name = 'experiment_id'
    ) THEN
        ALTER TABLE public.cracking ADD COLUMN experiment_id UUID;
    END IF;
    
    ALTER TABLE public.cracking ALTER COLUMN experiment_number DROP NOT NULL;
    ALTER TABLE public.cracking ALTER COLUMN experiment_phase_id DROP NOT NULL;

    -- Shelling table
    IF NOT EXISTS (
        SELECT 1 FROM information_schema.columns 
        WHERE table_schema = 'public' AND table_name = 'shelling' AND column_name = 'experiment_id'
    ) THEN
        ALTER TABLE public.shelling ADD COLUMN experiment_id UUID;
    END IF;
    
    ALTER TABLE public.shelling ALTER COLUMN experiment_number DROP NOT NULL;
    ALTER TABLE public.shelling ALTER COLUMN experiment_phase_id DROP NOT NULL;
END $$;

-- 2) Backfill experiment_id from composite key for all phase tables (only if old data exists)
-- This migration is designed to work with existing data that has the old schema
-- For fresh data, the seed files will populate experiment_id directly
DO $$
BEGIN
    -- Only backfill if there are records with the old schema (experiment_number is NOT NULL)
    -- and experiment_id is NULL (meaning they haven't been migrated yet)
    IF EXISTS (SELECT 1 FROM public.soaking WHERE experiment_id IS NULL AND experiment_number IS NOT NULL) THEN
        UPDATE public.soaking s
        SET experiment_id = e.id
        FROM public.experiments e
        WHERE s.experiment_id IS NULL
          AND e.experiment_number = s.experiment_number
          AND e.phase_id = s.experiment_phase_id;
    END IF;

    IF EXISTS (SELECT 1 FROM public.airdrying WHERE experiment_id IS NULL AND experiment_number IS NOT NULL) THEN
        UPDATE public.airdrying a
        SET experiment_id = e.id
        FROM public.experiments e
        WHERE a.experiment_id IS NULL
          AND e.experiment_number = a.experiment_number
          AND e.phase_id = a.experiment_phase_id;
    END IF;

    IF EXISTS (SELECT 1 FROM public.cracking WHERE experiment_id IS NULL AND experiment_number IS NOT NULL) THEN
        UPDATE public.cracking c
        SET experiment_id = e.id
        FROM public.experiments e
        WHERE c.experiment_id IS NULL
          AND e.experiment_number = c.experiment_number
          AND e.phase_id = c.experiment_phase_id;
    END IF;

    IF EXISTS (SELECT 1 FROM public.shelling WHERE experiment_id IS NULL AND experiment_number IS NOT NULL) THEN
        UPDATE public.shelling s
        SET experiment_id = e.id
        FROM public.experiments e
        WHERE s.experiment_id IS NULL
          AND e.experiment_number = s.experiment_number
          AND e.phase_id = s.experiment_phase_id;
    END IF;
END $$;

-- 3) Add foreign key constraints to experiments(id)
ALTER TABLE public.soaking
    ADD CONSTRAINT soaking_experiment_id_fkey
        FOREIGN KEY (experiment_id) REFERENCES public.experiments(id) ON DELETE CASCADE;

ALTER TABLE public.airdrying
    ADD CONSTRAINT airdrying_experiment_id_fkey
        FOREIGN KEY (experiment_id) REFERENCES public.experiments(id) ON DELETE CASCADE;

ALTER TABLE public.cracking
    ADD CONSTRAINT cracking_experiment_id_fkey
        FOREIGN KEY (experiment_id) REFERENCES public.experiments(id) ON DELETE CASCADE;

ALTER TABLE public.shelling
    ADD CONSTRAINT shelling_experiment_id_fkey
        FOREIGN KEY (experiment_id) REFERENCES public.experiments(id) ON DELETE CASCADE;

-- 4) Create triggers to auto-populate experiment_id for phase tables
CREATE OR REPLACE FUNCTION public.set_experiment_id_on_soaking()
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;
    END IF;
    RETURN NEW;
END;
$func$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION public.set_experiment_id_on_airdrying()
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;
    END IF;
    RETURN NEW;
END;
$func$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION public.set_experiment_id_on_cracking()
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;
    END IF;
    RETURN NEW;
END;
$func$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION public.set_experiment_id_on_shelling()
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;
    END IF;
    RETURN NEW;
END;
$func$ LANGUAGE plpgsql;

-- Create triggers
DROP TRIGGER IF EXISTS trg_set_experiment_id_on_soaking ON public.soaking;
CREATE TRIGGER trg_set_experiment_id_on_soaking
    BEFORE INSERT OR UPDATE OF experiment_number, experiment_phase_id ON public.soaking
    FOR EACH ROW
    EXECUTE FUNCTION public.set_experiment_id_on_soaking();

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

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

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

-- 5) Make experiment_id NOT NULL after backfilling and triggers are in place
-- Only do this if there are no NULL values
DO $$
BEGIN
    -- Check if all records have experiment_id populated before making it NOT NULL
    IF NOT EXISTS (SELECT 1 FROM public.soaking WHERE experiment_id IS NULL) THEN
        ALTER TABLE public.soaking ALTER COLUMN experiment_id SET NOT NULL;
    END IF;
    
    IF NOT EXISTS (SELECT 1 FROM public.airdrying WHERE experiment_id IS NULL) THEN
        ALTER TABLE public.airdrying ALTER COLUMN experiment_id SET NOT NULL;
    END IF;
    
    IF NOT EXISTS (SELECT 1 FROM public.cracking WHERE experiment_id IS NULL) THEN
        ALTER TABLE public.cracking ALTER COLUMN experiment_id SET NOT NULL;
    END IF;
    
    IF NOT EXISTS (SELECT 1 FROM public.shelling WHERE experiment_id IS NULL) THEN
        ALTER TABLE public.shelling ALTER COLUMN experiment_id SET NOT NULL;
    END IF;
END $$;

-- 6) Create indexes for experiment_id
CREATE INDEX IF NOT EXISTS idx_soaking_experiment_id ON public.soaking(experiment_id);
CREATE INDEX IF NOT EXISTS idx_airdrying_experiment_id ON public.airdrying(experiment_id);
CREATE INDEX IF NOT EXISTS idx_cracking_experiment_id ON public.cracking(experiment_id);
CREATE INDEX IF NOT EXISTS idx_shelling_experiment_id ON public.shelling(experiment_id);

-- 7) Update unique constraints to use experiment_id instead of composite key
-- Drop old unique constraints
ALTER TABLE public.soaking DROP CONSTRAINT IF EXISTS unique_soaking_per_experiment;
ALTER TABLE public.airdrying DROP CONSTRAINT IF EXISTS unique_airdrying_per_experiment;
ALTER TABLE public.cracking DROP CONSTRAINT IF EXISTS unique_cracking_per_experiment;
ALTER TABLE public.shelling DROP CONSTRAINT IF EXISTS unique_shelling_per_experiment;

-- Add new unique constraints using experiment_id
ALTER TABLE public.soaking ADD CONSTRAINT unique_soaking_per_experiment UNIQUE (experiment_id);
ALTER TABLE public.airdrying ADD CONSTRAINT unique_airdrying_per_experiment UNIQUE (experiment_id);
ALTER TABLE public.cracking ADD CONSTRAINT unique_cracking_per_experiment UNIQUE (experiment_id);
ALTER TABLE public.shelling ADD CONSTRAINT unique_shelling_per_experiment UNIQUE (experiment_id);

COMMIT;