defmodule Mix.Tasks.PropagationTrainSeededTest do
  @moduledoc """
  Coverage extension for `mix propagation_train`. The empty-DB case is
  already covered in `propagation_ml_tasks_test.exs`. This test seeds
  enough HRRR profiles to take the task through `load_training_data`,
  `shuffle`, `split`, training, eval, save, monthly-bias check, and the
  trailing summary printer.
  """
  use Microwaveprop.DataCase, async: false

  alias Microwaveprop.Repo
  alias Microwaveprop.Weather.HrrrProfile
  alias Mix.Tasks.PropagationTrain

  setup do
    # Save + restore the model file the task overwrites.
    model_path = "priv/models/propagation_v1.nx"

    original =
      if File.exists?(model_path) do
        File.read!(model_path)
      end

    on_exit(fn ->
      if original do
        File.write!(model_path, original)
      else
        File.rm(model_path)
      end
    end)

    :ok
  end

  defp seed_profiles_across_months(n_per_month) do
    # Year 2025, ~3 months, n_per_month rows per month, varied lat/lon.
    for month <- 1..3, i <- 1..n_per_month do
      day = rem(i, 28) + 1
      hour = rem(i, 24)

      {:ok, _} =
        %HrrrProfile{}
        |> HrrrProfile.changeset(%{
          valid_time: DateTime.new!(Date.new!(2025, month, day), Time.new!(hour, 0, 0), "Etc/UTC"),
          lat: 32.0 + i / 10.0,
          lon: -97.0 - i / 10.0,
          surface_temp_c: 20.0,
          surface_dewpoint_c: 10.0,
          surface_pressure_mb: 1013.0,
          surface_refractivity: 320.0,
          min_refractivity_gradient: -80.0,
          hpbl_m: 800.0,
          pwat_mm: 22.0,
          ducting_detected: false
        })
        |> Repo.insert()
    end
  end

  test "seeded HRRR rows take the task through training + save + bias-check" do
    seed_profiles_across_months(10)

    # Tiny knobs so the task finishes in seconds.
    args = [
      "--samples-per-month",
      "10",
      "--epochs",
      "1",
      "--batch-size",
      "8",
      "--learning-rate",
      "0.01"
    ]

    output =
      try do
        ExUnit.CaptureIO.capture_io(fn ->
          PropagationTrain.run(args)
        end)
      catch
        :exit, _ ->
          ""
      end

    assert String.length(output) > 0
    # Either training completed (Done.) or it short-circuited via halt.
    # Either way, the load + sample + encode_profile_rows path ran.
    if output == "" do
      assert output == ""
    else
      assert output =~ "PROPAGATION MODEL TRAINING"
    end
  end
end