defmodule Microwaveprop.Weather.NexradClient do
  @moduledoc """
  Fetches and processes IEM CONUS n0q composite reflectivity PNGs.

  The n0q product is a palettized 8-bit PNG, 12200 x 5400 pixels,
  covering CONUS at 0.005 degrees/pixel (-126W to -65W, 23N to 50N).
  Available every 5 minutes from the IEM archive.

  Pixel value 0 = no echo; values 1-255 map linearly to roughly
  -30 to +95 dBZ. For each point of interest we extract a ~25 km
  box (~50x50 pixels) and compute summary statistics.
  """

  alias Microwaveprop.Weather.NexradCache

  require Logger

  # Image geometry
  @lon_min -126.0
  @lat_max 50.0
  @deg_per_pixel 0.005

  # Default box half-size in pixels (~25 km at mid-latitudes)
  @default_box_half 25

  @base_url "https://mesonet.agron.iastate.edu/archive/data"

  @doc "Fetch and process one n0q frame. Returns list of observation maps for points of interest."
  @spec fetch_frame(DateTime.t(), [{float(), float()}]) :: {:ok, [map()]} | {:error, term()}
  def fetch_frame(timestamp, points_of_interest) do
    rounded = round_to_5min(timestamp)
    url = frame_url(rounded)

    req_opts = Application.get_env(:microwaveprop, :nexrad_req_options, [])

    Microwaveprop.Instrument.span([:nexrad, :fetch_frame], %{}, fn ->
      case http_get(url, [receive_timeout: 60_000, retry: false] ++ req_opts) do
        {:ok, %{status: 200, body: body}} when is_binary(body) ->
          process_frame(body, rounded, points_of_interest)

        {:ok, %{status: status}} ->
          Logger.warning("NexradClient: HTTP #{status} for #{url}")
          {:error, "NEXRAD n0q HTTP #{status}"}

        {:error, reason} ->
          Logger.warning("NexradClient: transport error for #{url}: #{inspect(reason)}")
          {:error, reason}
      end
    end)
  end

  @doc "Round a DateTime down to the nearest 5-minute boundary."
  @spec round_to_5min(DateTime.t()) :: DateTime.t()
  def round_to_5min(%DateTime{} = dt) do
    rounded_minute = div(dt.minute, 5) * 5
    %{dt | minute: rounded_minute, second: 0, microsecond: {0, 0}}
  end

  @doc "Build the IEM archive URL for an n0q frame at the given (already-rounded) timestamp."
  @spec frame_url(DateTime.t()) :: String.t()
  def frame_url(%DateTime{} = dt) do
    date_path = Calendar.strftime(dt, "%Y/%m/%d")
    file_ts = Calendar.strftime(dt, "%Y%m%d%H%M")
    "#{@base_url}/#{date_path}/GIS/uscomp/n0q_#{file_ts}.png"
  end

  @doc """
  Fetch the latest n0q frame and extract all rain cells with reflectivity
  above `min_dbz` within `radius_km` of the given point. Returns
  `{:ok, [{lat, lon, dbz}]}` or `{:error, reason}`.

  Samples every ~5 km (10 pixels) within the bounding box for efficiency.
  The decoded PNG pixel buffer is cached in `NexradCache` for the current
  5-minute window, so concurrent clicks and repeat lookups skip the HTTP +
  PNG decode entirely.
  """
  @spec fetch_rain_cells(float(), float(), float(), float()) :: {:ok, [{float(), float(), float()}]} | {:error, term()}
  def fetch_rain_cells(lat, lon, radius_km \\ 300.0, min_dbz \\ 25.0) do
    case fetch_decoded_frame(DateTime.utc_now()) do
      {:ok, pixels, width} ->
        {:ok, extract_rain_cells(pixels, width, lat, lon, radius_km, min_dbz)}

      error ->
        error
    end
  end

  @doc """
  Fetch and decode the n0q PNG frame nearest the given timestamp.
  Returns `{:ok, pixels, width}` — a flat binary of palette-index bytes —
  or `{:error, reason}`. Goes through `NexradCache`, keyed on the 5-min
  rounded timestamp — concurrent radar workers (and any backfill
  neighbors that share a 5-min window) reuse the same decoded buffer
  instead of each paying the full ~5 MB download + ~1.7 s decode.
  """
  @spec fetch_decoded_frame(DateTime.t()) ::
          {:ok, binary(), non_neg_integer()} | {:error, term()}
  def fetch_decoded_frame(%DateTime{} = dt) do
    rounded = round_to_5min(dt)

    case NexradCache.fetch(rounded) do
      {:ok, pixels, width} ->
        {:ok, pixels, width}

      :miss ->
        with {:ok, pixels, width} <- fetch_and_decode_frame(rounded) do
          NexradCache.put(rounded, pixels, width)
          NexradCache.prune_older_than(DateTime.add(rounded, -600, :second))
          {:ok, pixels, width}
        end
    end
  end

  defp fetch_and_decode_frame(rounded) do
    url = frame_url(rounded)

    Microwaveprop.Instrument.span([:nexrad, :fetch_frame], %{}, fn ->
      handle_decode_frame_response(url, do_fetch_png(url))
    end)
  end

  defp do_fetch_png(url) do
    req_opts = Application.get_env(:microwaveprop, :nexrad_req_options, [])
    http_get(url, [receive_timeout: 60_000, retry: false] ++ req_opts)
  end

  defp handle_decode_frame_response(_url, {:ok, %{status: 200, body: body}}) when is_binary(body) do
    Microwaveprop.Instrument.span([:nexrad, :decode_png], %{bytes: byte_size(body)}, fn ->
      decode_png_to_pixels(body)
    end)
  end

  defp handle_decode_frame_response(url, {:ok, %{status: status}}) do
    Logger.warning("NexradClient: HTTP #{status} for #{url}")
    {:error, "NEXRAD n0q HTTP #{status}"}
  end

  defp handle_decode_frame_response(url, {:error, reason}) do
    Logger.warning("NexradClient: transport error for #{url}: #{inspect(reason)}")
    {:error, reason}
  end

  defp extract_rain_cells(pixels, width, center_lat, center_lon, radius_km, min_dbz) do
    height = div(byte_size(pixels), width)

    # Convert radius to approximate pixel count (~0.005 deg/px, ~0.5 km/px at mid-lat)
    dlat = radius_km / 111.0
    dlon = radius_km / (111.0 * :math.cos(:math.pi() * center_lat / 180.0))

    lat_min = center_lat - dlat
    lat_max = center_lat + dlat
    lon_min = center_lon - dlon
    lon_max = center_lon + dlon

    {x_min, y_max_px} = latlon_to_pixel(lat_min, lon_min)
    {x_max, y_min_px} = latlon_to_pixel(lat_max, lon_max)

    x_min = max(x_min, 0)
    x_max = min(x_max, width - 1)
    y_min_px = max(y_min_px, 0)
    y_max_px = min(y_max_px, height - 1)

    # Sample every 10 pixels (~5 km) for efficiency
    step = 10

    for y <- y_min_px..y_max_px//step,
        x <- x_min..x_max//step,
        offset = y * width + x,
        offset >= 0,
        offset < byte_size(pixels),
        <<_::binary-size(^offset), pixel_val::8, _::binary>> = pixels,
        pixel_val > 0,
        dbz = pixel_to_dbz(pixel_val),
        dbz >= min_dbz do
      cell_lat = @lat_max - y * @deg_per_pixel
      cell_lon = @lon_min + x * @deg_per_pixel
      {Float.round(cell_lat, 3), Float.round(cell_lon, 3), Float.round(dbz, 1)}
    end
  end

  @doc """
  Convert (lat, lon) to pixel coordinates in the n0q image.

  Returns `{x, y}` where x is the column (0 = leftmost = -126W)
  and y is the row (0 = topmost = 50N).
  """
  @spec latlon_to_pixel(float(), float()) :: {integer(), integer()}
  def latlon_to_pixel(lat, lon) do
    x = round((lon - @lon_min) / @deg_per_pixel)
    y = round((@lat_max - lat) / @deg_per_pixel)
    {x, y}
  end

  @doc """
  Convert a pixel value (0-255) to dBZ.

  Value 0 = no echo (returns 0.0). Values 1-255 map linearly to
  approximately -30 to +95 dBZ.
  """
  @spec pixel_to_dbz(integer()) :: float()
  def pixel_to_dbz(0), do: 0.0

  def pixel_to_dbz(value) when value >= 1 and value <= 255 do
    # Linear mapping: 1 -> -30 dBZ, 255 -> 95 dBZ
    -30.0 + (value - 1) / 254.0 * 125.0
  end

  @doc """
  Extract summary statistics from a box of pixels centered at (cx, cy).

  `pixels` is a flat binary where each byte is one pixel value.
  `width` is the image width (number of columns per row).
  `half` is the half-size of the extraction box in pixels.

  Returns a map with :mean_reflectivity_dbz, :max_reflectivity_dbz,
  :texture_variance, and :pixel_count.
  """
  @spec extract_box(binary(), integer(), integer(), integer(), integer()) :: map()
  def extract_box(pixels, width, cx, cy, half \\ @default_box_half) do
    height = div(byte_size(pixels), width)

    x_min = max(cx - half, 0)
    x_max = min(cx + half, width - 1)
    y_min = max(cy - half, 0)
    y_max = min(cy + half, height - 1)

    # Collect non-zero dBZ values. `//1` forces an ascending range so
    # a box whose center is outside the image (min > max after
    # clamping) becomes empty rather than iterating backwards.
    dbz_values =
      for y <- y_min..y_max//1,
          x <- x_min..x_max//1,
          offset = y * width + x,
          offset >= 0,
          offset < byte_size(pixels),
          <<_::binary-size(^offset), pixel_val::8, _::binary>> = pixels,
          pixel_val > 0 do
        pixel_to_dbz(pixel_val)
      end

    count = length(dbz_values)

    compute_box_stats(dbz_values, count)
  end

  # -- Private --

  defp http_get(url, opts) do
    runner = Application.get_env(:microwaveprop, :nexrad_http_get, &Req.get/2)
    runner.(url, opts)
  end

  defp compute_box_stats(_values, 0) do
    %{mean_reflectivity_dbz: 0.0, max_reflectivity_dbz: 0.0, texture_variance: 0.0, pixel_count: 0}
  end

  defp compute_box_stats(values, count) do
    mean = Enum.sum(values) / count
    max_val = Enum.max(values)
    variance = sample_variance(values, mean, count)

    %{
      mean_reflectivity_dbz: Float.round(mean, 2),
      max_reflectivity_dbz: Float.round(max_val, 2),
      texture_variance: Float.round(variance, 2),
      pixel_count: count
    }
  end

  defp sample_variance(_values, _mean, count) when count < 2, do: 0.0

  defp sample_variance(values, mean, count) do
    values
    |> Enum.map(fn v -> (v - mean) * (v - mean) end)
    |> Enum.sum()
    |> Kernel./(count - 1)
  end

  defp process_frame(png_binary, observed_at, points) do
    case decode_png_to_pixels(png_binary) do
      {:ok, pixels, width} ->
        observations =
          Enum.map(points, fn {lat, lon} ->
            {cx, cy} = latlon_to_pixel(lat, lon)
            stats = extract_box(pixels, width, cx, cy)

            Map.merge(stats, %{
              observed_at: observed_at,
              lat: lat,
              lon: lon
            })
          end)

        {:ok, observations}

      {:error, reason} ->
        {:error, reason}
    end
  end

  # Minimal PNG decoder for indexed-color (palette) 8-bit images.
  #
  # Extracts raw palette index bytes (not RGB), which is what we need
  # since the n0q product encodes dBZ as palette index values.
  #
  # PNG structure:
  #   8-byte signature
  #   Chunks: [length(4) | type(4) | data(length) | crc(4)]
  #   IHDR chunk has width, height, bit_depth, color_type
  #   IDAT chunks contain zlib-compressed scanline data
  #   Each scanline: filter_byte | pixel_data[width]
  defp decode_png_to_pixels(png_binary) do
    # Validate PNG signature
    <<137, 80, 78, 71, 13, 10, 26, 10, rest::binary>> = png_binary

    # Parse chunks
    chunks = parse_chunks(rest, [])

    # Get dimensions from IHDR
    [{_type, ihdr_data} | _] = chunks
    <<width::32, height::32, _bit_depth::8, _color_type::8, _rest::binary>> = ihdr_data

    # Concatenate all IDAT data and decompress
    idat_data =
      chunks
      |> Enum.filter(fn {type, _} -> type == "IDAT" end)
      |> Enum.map(fn {_, data} -> data end)
      |> IO.iodata_to_binary()

    decompressed = :zlib.uncompress(idat_data)

    # Undo per-row PNG filters, extract raw pixel index bytes
    pixels = unfilter_rows(decompressed, width, height)

    {:ok, pixels, width}
  rescue
    e ->
      {:error, "PNG decode failed: #{inspect(e)}"}
  end

  defp parse_chunks(<<>>, acc), do: Enum.reverse(acc)

  defp parse_chunks(<<length::32, type::binary-size(4), rest::binary>>, acc) do
    <<data::binary-size(^length), _crc::32, remaining::binary>> = rest
    parse_chunks(remaining, [{type, data} | acc])
  end

  # Undo PNG row filters for 8-bit indexed images (1 byte per pixel).
  # Returns a flat binary of pixel index values.
  defp unfilter_rows(data, width, height) do
    # bytes_per_pixel = 1 for indexed color
    # Each row is: 1 filter byte + width pixel bytes
    stride = width

    {rows, _} =
      Enum.reduce(1..height, {[], data}, fn _row_idx, {acc, remaining} ->
        <<filter::8, row_data::binary-size(^stride), rest::binary>> = remaining
        prev_row = List.first(acc, :binary.copy(<<0>>, stride))
        filtered = apply_filter(filter, row_data, prev_row, 1)
        {[filtered | acc], rest}
      end)

    rows
    |> Enum.reverse()
    |> IO.iodata_to_binary()
  end

  # PNG filter types for 8-bit indexed (bpp=1)
  defp apply_filter(0, row, _prev, _bpp), do: row

  defp apply_filter(1, row, _prev, bpp) do
    # Sub: each byte is diff from byte bpp positions to the left
    unfilter_sub(row, bpp)
  end

  defp apply_filter(2, row, prev, _bpp) do
    # Up: each byte is diff from byte directly above
    unfilter_up(row, prev)
  end

  defp apply_filter(3, row, prev, bpp) do
    # Average: each byte is diff from average of left and above
    unfilter_average(row, prev, bpp)
  end

  defp apply_filter(4, row, prev, bpp) do
    # Paeth: each byte uses Paeth predictor
    unfilter_paeth(row, prev, bpp)
  end

  # Indexed-color PNG only ever has bpp = 1 (one palette index byte per
  # pixel), so "left" is simply the previously-decoded output byte. We
  # walk the input binary once and accumulate output bytes into an
  # iolist, which is O(n) — the original `acc ++ [byte]` plus
  # `Enum.at(acc, idx - bpp)` was O(n²) per row and dominated decode
  # time for the 12200×5400 n0q frame.
  defp unfilter_sub(row, 1), do: unfilter_sub_iter(row, 0, [])

  defp unfilter_sub_iter(<<>>, _left, acc), do: acc |> Enum.reverse() |> :binary.list_to_bin()

  defp unfilter_sub_iter(<<byte, rest::binary>>, left, acc) do
    val = rem(byte + left, 256)
    unfilter_sub_iter(rest, val, [val | acc])
  end

  defp unfilter_up(row, prev), do: unfilter_up_iter(row, prev, [])

  defp unfilter_up_iter(<<>>, _prev, acc), do: acc |> Enum.reverse() |> :binary.list_to_bin()

  defp unfilter_up_iter(<<byte, row_rest::binary>>, <<above, prev_rest::binary>>, acc) do
    val = rem(byte + above, 256)
    unfilter_up_iter(row_rest, prev_rest, [val | acc])
  end

  # `prev` may be shorter than `row` if a row is shorter than the
  # previous one — for indexed-color this never happens but we
  # defensively treat missing prior bytes as zero.
  defp unfilter_up_iter(<<byte, row_rest::binary>>, <<>>, acc) do
    val = rem(byte, 256)
    unfilter_up_iter(row_rest, <<>>, [val | acc])
  end

  defp unfilter_average(row, prev, 1), do: unfilter_average_iter(row, prev, 0, [])

  defp unfilter_average_iter(<<>>, _prev, _left, acc), do: acc |> Enum.reverse() |> :binary.list_to_bin()

  defp unfilter_average_iter(<<byte, row_rest::binary>>, <<above, prev_rest::binary>>, left, acc) do
    val = rem(byte + div(left + above, 2), 256)
    unfilter_average_iter(row_rest, prev_rest, val, [val | acc])
  end

  defp unfilter_average_iter(<<byte, row_rest::binary>>, <<>>, left, acc) do
    val = rem(byte + div(left, 2), 256)
    unfilter_average_iter(row_rest, <<>>, val, [val | acc])
  end

  defp unfilter_paeth(row, prev, 1), do: unfilter_paeth_iter(row, prev, 0, 0, [])

  defp unfilter_paeth_iter(<<>>, _prev, _left, _upper_left, acc), do: acc |> Enum.reverse() |> :binary.list_to_bin()

  defp unfilter_paeth_iter(<<byte, row_rest::binary>>, <<above, prev_rest::binary>>, left, upper_left, acc) do
    pr = paeth_predictor(left, above, upper_left)
    val = rem(byte + pr, 256)
    unfilter_paeth_iter(row_rest, prev_rest, val, above, [val | acc])
  end

  defp unfilter_paeth_iter(<<byte, row_rest::binary>>, <<>>, left, _upper_left, acc) do
    pr = paeth_predictor(left, 0, 0)
    val = rem(byte + pr, 256)
    unfilter_paeth_iter(row_rest, <<>>, val, 0, [val | acc])
  end

  defp paeth_predictor(a, b, c) do
    p = a + b - c
    pa = abs(p - a)
    pb = abs(p - b)
    pc = abs(p - c)

    cond do
      pa <= pb and pa <= pc -> a
      pb <= pc -> b
      true -> c
    end
  end
end