DO $$
DECLARE
  column_type text;
BEGIN
  SELECT data_type
  INTO column_type
  FROM information_schema.columns
  WHERE table_schema = 'public'
    AND table_name = 'base_model_catalog'
    AND column_name = 'model_type';

  IF column_type IS DISTINCT FROM 'jsonb' THEN
    DROP INDEX IF EXISTS idx_base_model_catalog_type;

    ALTER TABLE base_model_catalog
      ADD COLUMN IF NOT EXISTS model_type_next jsonb NOT NULL DEFAULT '[]'::jsonb;

    WITH source_types AS (
      SELECT id,
             CASE
               WHEN jsonb_typeof(capabilities->'originalTypes') = 'array' THEN capabilities->'originalTypes'
               WHEN jsonb_typeof(metadata->'originalTypes') = 'array' THEN metadata->'originalTypes'
               ELSE jsonb_build_array(model_type)
             END AS model_types
      FROM base_model_catalog
    ),
    normalized_types AS (
      SELECT source_types.id, normalized.model_type
      FROM source_types
      CROSS JOIN LATERAL jsonb_array_elements_text(source_types.model_types) AS raw_type(model_type)
      CROSS JOIN LATERAL (
        SELECT 'text_generate' AS model_type
        WHERE raw_type.model_type IN ('chat', 'text', 'responses')
        UNION ALL SELECT 'image_generate'
        WHERE raw_type.model_type = 'image'
        UNION ALL SELECT 'image_edit'
        WHERE raw_type.model_type = 'image'
        UNION ALL SELECT 'image_generate'
        WHERE raw_type.model_type = 'images.generations'
        UNION ALL SELECT 'image_edit'
        WHERE raw_type.model_type = 'images.edits'
        UNION ALL SELECT 'video_generate'
        WHERE raw_type.model_type IN ('video', 'videos.generations')
        UNION ALL SELECT raw_type.model_type
        WHERE trim(raw_type.model_type) <> ''
          AND raw_type.model_type NOT IN (
            'chat', 'text', 'responses', 'image', 'images.generations', 'images.edits', 'video', 'videos.generations'
          )
      ) AS normalized
    )
    UPDATE base_model_catalog target
    SET model_type_next = COALESCE((
      SELECT jsonb_agg(DISTINCT normalized_types.model_type ORDER BY normalized_types.model_type)
      FROM normalized_types
      WHERE normalized_types.id = target.id
    ), '[]'::jsonb);

    ALTER TABLE base_model_catalog DROP COLUMN model_type;
    ALTER TABLE base_model_catalog RENAME COLUMN model_type_next TO model_type;
  END IF;
END $$;

UPDATE base_model_catalog
SET model_type = CASE
  WHEN jsonb_typeof(model_type) = 'array' THEN model_type
  WHEN jsonb_typeof(model_type) = 'string' THEN jsonb_build_array(model_type #>> '{}')
  ELSE '[]'::jsonb
END;

UPDATE base_model_catalog
SET default_snapshot = default_snapshot || jsonb_build_object('modelType', model_type)
WHERE catalog_type = 'system'
  AND COALESCE(default_snapshot, '{}'::jsonb) <> '{}'::jsonb;

ALTER TABLE base_model_catalog
  ALTER COLUMN model_type SET NOT NULL,
  ALTER COLUMN model_type SET DEFAULT '[]'::jsonb;

CREATE INDEX IF NOT EXISTS idx_base_model_catalog_type
  ON base_model_catalog(provider_key, status);

CREATE INDEX IF NOT EXISTS idx_base_model_catalog_model_type
  ON base_model_catalog USING gin(model_type);

CREATE OR REPLACE FUNCTION fill_system_base_model_default_snapshot()
RETURNS trigger AS $$
BEGIN
  IF NEW.catalog_type = 'system' AND NEW.default_snapshot IS NULL THEN
    NEW.default_snapshot = jsonb_build_object(
      'providerKey', NEW.provider_key,
      'canonicalModelKey', NEW.canonical_model_key,
      'providerModelName', NEW.provider_model_name,
      'modelType', NEW.model_type,
      'modelAlias', NEW.display_name,
      'capabilities', NEW.capabilities,
      'baseBillingConfig', NEW.base_billing_config,
      'defaultRateLimitPolicy', NEW.default_rate_limit_policy,
      'pricingRuleSetId', COALESCE(NEW.pricing_rule_set_id::text, ''),
      'runtimePolicySetId', COALESCE(NEW.runtime_policy_set_id::text, ''),
      'runtimePolicyOverride', NEW.runtime_policy_override,
      'metadata', NEW.metadata,
      'pricingVersion', NEW.pricing_version,
      'status', NEW.status
    );
  END IF;
  RETURN NEW;
END;
$$ LANGUAGE plpgsql;
