defmodule Mix.Tasks.HrrrClimatology do
  @shortdoc "Build surface temperature climatology from hrrr_profiles"
  @moduledoc """
  Aggregates `hrrr_profiles.surface_temp_c` by (lat, lon, month, hour)
  into `hrrr_climatology` for use by the temperature-anomaly feature.

  Discovers which (month, hour) combos have data first, then processes
  only those batches. Idempotent via ON CONFLICT.

      mix hrrr_climatology                  # build from all grid-point profiles
      mix hrrr_climatology --min-samples 5  # require at least 5 observations per cell
  """
  use Mix.Task

  alias Microwaveprop.Repo

  @impl Mix.Task
  def run(argv) do
    Mix.Task.run("app.start")
    Oban.pause_all_queues(Oban)

    {opts, _, _} = OptionParser.parse(argv, switches: [min_samples: :integer])
    min_samples = Keyword.get(opts, :min_samples, 3)

    # Discover which (month, hour) combos actually have data
    %{rows: combos} =
      Repo.query!(
        """
        SELECT EXTRACT(MONTH FROM valid_time)::int AS month,
               EXTRACT(HOUR FROM valid_time)::int AS hour
        FROM hrrr_profiles
        WHERE surface_temp_c IS NOT NULL AND is_grid_point = true
        GROUP BY 1, 2
        ORDER BY 1, 2
        """,
        [],
        timeout: 120_000
      )

    Mix.shell().info("Building climatology (min_samples=#{min_samples}, #{length(combos)} batches)...")

    total =
      combos
      |> Enum.with_index(1)
      |> Enum.reduce(0, fn {[month, hour], idx}, acc ->
        %{num_rows: count} =
          Repo.query!(
            """
            INSERT INTO hrrr_climatology (id, lat, lon, month, hour,
                                          mean_surface_temp_c, stddev_surface_temp_c, sample_count)
            SELECT gen_random_uuid(), lat, lon,
                   $2 AS month,
                   $3 AS hour,
                   AVG(surface_temp_c),
                   STDDEV_SAMP(surface_temp_c),
                   COUNT(*)
            FROM hrrr_profiles
            WHERE surface_temp_c IS NOT NULL
              AND is_grid_point = true
              AND EXTRACT(MONTH FROM valid_time)::int = $2
              AND EXTRACT(HOUR FROM valid_time)::int = $3
            GROUP BY lat, lon
            HAVING COUNT(*) >= $1
            ON CONFLICT (lat, lon, month, hour)
            DO UPDATE SET
              mean_surface_temp_c = EXCLUDED.mean_surface_temp_c,
              stddev_surface_temp_c = EXCLUDED.stddev_surface_temp_c,
              sample_count = EXCLUDED.sample_count
            """,
            [min_samples, month, hour],
            timeout: 300_000
          )

        Mix.shell().info("  [#{idx}/#{length(combos)}] month=#{month} hour=#{hour}: #{count} rows")
        acc + count
      end)

    Mix.shell().info("Upserted #{total} climatology records total.")
  end
end
