defmodule Microwaveprop.Repo.Migrations.PartitionHrrrProfiles do
  use Ecto.Migration

  @disable_ddl_transaction true
  @disable_migration_lock true

  def up do
    # Clean up any partial previous run: drop the partitioned table and orphaned
    # partition tables, then rename _old back so we can start fresh.
    execute """
    DO $$
    DECLARE
      tbl text;
    BEGIN
      IF EXISTS (SELECT 1 FROM pg_tables WHERE tablename = 'hrrr_profiles_old' AND schemaname = 'public') THEN
        DROP TABLE IF EXISTS hrrr_profiles CASCADE;
        -- Drop orphaned partition tables left behind after CASCADE
        FOR tbl IN
          SELECT tablename FROM pg_tables
          WHERE schemaname = 'public' AND tablename LIKE 'hrrr_profiles_%'
            AND tablename != 'hrrr_profiles_old'
        LOOP
          EXECUTE 'DROP TABLE IF EXISTS ' || tbl || ' CASCADE';
        END LOOP;
        ALTER TABLE hrrr_profiles_old RENAME TO hrrr_profiles;
      END IF;
    END $$
    """

    # 1. Rename old table
    execute "ALTER TABLE hrrr_profiles RENAME TO hrrr_profiles_old"

    # 2. Create partitioned table with identical schema
    execute """
    CREATE TABLE hrrr_profiles (
      id uuid NOT NULL DEFAULT gen_random_uuid(),
      valid_time timestamp(0) WITHOUT TIME ZONE NOT NULL,
      lat double precision NOT NULL,
      lon double precision NOT NULL,
      run_time timestamp(0) WITHOUT TIME ZONE,
      profile jsonb[] DEFAULT ARRAY[]::jsonb[],
      hpbl_m double precision,
      pwat_mm double precision,
      surface_temp_c double precision,
      surface_dewpoint_c double precision,
      surface_pressure_mb double precision,
      surface_refractivity double precision,
      min_refractivity_gradient double precision,
      ducting_detected boolean DEFAULT false,
      duct_characteristics jsonb[],
      inserted_at timestamp(0) WITHOUT TIME ZONE NOT NULL,
      updated_at timestamp(0) WITHOUT TIME ZONE NOT NULL
    ) PARTITION BY RANGE (valid_time)
    """

    # 3. Create partitions
    partitions = [
      {"2016-01-01", "2017-01-01"},
      {"2017-01-01", "2018-01-01"},
      {"2018-01-01", "2018-07-01"},
      {"2018-07-01", "2019-01-01"},
      {"2019-01-01", "2019-07-01"},
      {"2019-07-01", "2019-10-01"},
      {"2019-10-01", "2020-01-01"},
      {"2020-01-01", "2020-07-01"},
      {"2020-07-01", "2020-10-01"},
      {"2020-10-01", "2021-01-01"},
      {"2021-01-01", "2021-07-01"},
      {"2021-07-01", "2021-10-01"},
      {"2021-10-01", "2022-01-01"},
      {"2022-01-01", "2022-07-01"},
      {"2022-07-01", "2022-10-01"},
      {"2022-10-01", "2023-01-01"},
      {"2023-01-01", "2023-07-01"},
      {"2023-07-01", "2023-10-01"},
      {"2023-10-01", "2024-01-01"},
      {"2024-01-01", "2024-07-01"},
      {"2024-07-01", "2024-10-01"},
      {"2024-10-01", "2025-01-01"},
      {"2025-01-01", "2025-07-01"},
      {"2025-07-01", "2026-01-01"},
      {"2026-01-01", "2026-04-01"},
      {"2026-04-01", "2026-07-01"},
      {"2026-07-01", "2026-10-01"},
      {"2026-10-01", "2027-01-01"},
      {"2027-01-01", "2027-07-01"},
      {"2027-07-01", "2028-01-01"}
    ]

    for {from_date, to_date} <- partitions do
      suffix = from_date |> String.replace("-", "_") |> String.slice(0..6)

      execute """
      CREATE TABLE hrrr_profiles_#{suffix} PARTITION OF hrrr_profiles
        FOR VALUES FROM ('#{from_date}') TO ('#{to_date}')
      """
    end

    # 4. Indexes on partitioned table
    execute "DROP INDEX IF EXISTS hrrr_profiles_lat_lon_valid_time_index"
    execute "DROP INDEX IF EXISTS hrrr_profiles_valid_time_index"

    execute "CREATE UNIQUE INDEX hrrr_profiles_lat_lon_valid_time_index ON hrrr_profiles (lat, lon, valid_time)"
    execute "CREATE INDEX hrrr_profiles_valid_time_index ON hrrr_profiles (valid_time)"

    # 5. Migrate data
    execute "INSERT INTO hrrr_profiles SELECT * FROM hrrr_profiles_old"

    # 6. Drop old table
    execute "DROP TABLE hrrr_profiles_old"
  end

  def down do
    execute """
    CREATE TABLE hrrr_profiles_flat (
      id uuid NOT NULL,
      valid_time timestamp(0) WITHOUT TIME ZONE NOT NULL,
      lat double precision NOT NULL,
      lon double precision NOT NULL,
      run_time timestamp(0) WITHOUT TIME ZONE,
      profile jsonb[] DEFAULT ARRAY[]::jsonb[],
      hpbl_m double precision,
      pwat_mm double precision,
      surface_temp_c double precision,
      surface_dewpoint_c double precision,
      surface_pressure_mb double precision,
      surface_refractivity double precision,
      min_refractivity_gradient double precision,
      ducting_detected boolean DEFAULT false,
      duct_characteristics jsonb[],
      inserted_at timestamp(0) WITHOUT TIME ZONE NOT NULL,
      updated_at timestamp(0) WITHOUT TIME ZONE NOT NULL,
      PRIMARY KEY (id)
    )
    """

    execute "INSERT INTO hrrr_profiles_flat SELECT * FROM hrrr_profiles"
    execute "DROP TABLE hrrr_profiles"
    execute "ALTER TABLE hrrr_profiles_flat RENAME TO hrrr_profiles"

    execute "CREATE UNIQUE INDEX hrrr_profiles_lat_lon_valid_time_index ON hrrr_profiles (lat, lon, valid_time)"
    execute "CREATE INDEX hrrr_profiles_valid_time_index ON hrrr_profiles (valid_time)"
  end
end