BEGIN;

ALTER TABLE user_broker ADD COLUMN IF NOT EXISTS api_secret TEXT;
ALTER TABLE user_broker ADD COLUMN IF NOT EXISTS auth_state TEXT;

UPDATE user_broker
SET auth_state = 'UNKNOWN'
WHERE auth_state IS NULL;

DO $$
BEGIN
  IF EXISTS (SELECT 1 FROM pg_constraint WHERE conname = 'chk_strategy_run_status') THEN
    ALTER TABLE strategy_run DROP CONSTRAINT chk_strategy_run_status;
  END IF;
  ALTER TABLE strategy_run
    ADD CONSTRAINT chk_strategy_run_status
    CHECK (status IN ('RUNNING','STOPPED','ERROR','PAUSED_AUTH_EXPIRED'));
END $$;

CREATE OR REPLACE FUNCTION enforce_strategy_run_transition()
RETURNS trigger AS $$
BEGIN
  IF TG_OP = 'UPDATE' THEN
    IF OLD.status = NEW.status THEN
      RETURN NEW;
    END IF;

    IF OLD.status = 'ERROR' THEN
      RAISE EXCEPTION 'run status cannot transition from % to %', OLD.status, NEW.status;
    END IF;

    IF OLD.status IN ('STOPPED','PAUSED_AUTH_EXPIRED') AND NEW.status <> 'RUNNING' THEN
      RAISE EXCEPTION 'run status cannot transition from % to %', OLD.status, NEW.status;
    END IF;

    IF OLD.status = 'RUNNING' AND NEW.status NOT IN ('STOPPED','ERROR','PAUSED_AUTH_EXPIRED') THEN
      RAISE EXCEPTION 'run status cannot transition from % to %', OLD.status, NEW.status;
    END IF;
  END IF;
  RETURN NEW;
END;
$$ LANGUAGE plpgsql;

DROP INDEX IF EXISTS uq_one_running_run_per_user;

COMMIT;