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.
  """

  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, [])

    case Req.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}} ->
        {:error, "NEXRAD n0q HTTP #{status}"}

      {:error, reason} ->
        {:error, reason}
    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.
  """
  @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
    now = DateTime.utc_now()
    rounded = round_to_5min(now)
    url = frame_url(rounded)

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

    case Req.get(url, [receive_timeout: 60_000, retry: false] ++ req_opts) do
      {:ok, %{status: 200, body: body}} when is_binary(body) ->
        case decode_png_to_pixels(body) do
          {:ok, pixels, width} ->
            cells = extract_rain_cells(pixels, width, lat, lon, radius_km, min_dbz)
            {:ok, cells}

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

      {:ok, %{status: status}} ->
        {:error, "NEXRAD n0q HTTP #{status}"}

      {:error, reason} ->
        {:error, reason}
    end
  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
    dbz_values =
      for y <- y_min..y_max,
          x <- x_min..x_max,
          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 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

  defp unfilter_sub(row, bpp) do
    row_bytes = :binary.bin_to_list(row)

    {result, _} =
      Enum.reduce(row_bytes, {[], 0}, fn byte, {acc, idx} ->
        left = if idx >= bpp, do: Enum.at(acc, idx - bpp), else: 0
        val = rem(byte + left, 256)
        {acc ++ [val], idx + 1}
      end)

    :binary.list_to_bin(result)
  end

  defp unfilter_up(row, prev) do
    row_bytes = :binary.bin_to_list(row)
    prev_bytes = :binary.bin_to_list(prev)

    row_bytes
    |> Enum.zip(prev_bytes)
    |> Enum.map(fn {a, b} -> rem(a + b, 256) end)
    |> :binary.list_to_bin()
  end

  defp unfilter_average(row, prev, bpp) do
    row_bytes = :binary.bin_to_list(row)
    prev_bytes = :binary.bin_to_list(prev)

    {result, _} =
      Enum.reduce(Enum.with_index(row_bytes), {[], prev_bytes}, fn {byte, idx}, {acc, prev_list} ->
        left = if idx >= bpp, do: Enum.at(acc, idx - bpp), else: 0
        above = Enum.at(prev_list, idx, 0)
        val = rem(byte + div(left + above, 2), 256)
        {acc ++ [val], prev_list}
      end)

    :binary.list_to_bin(result)
  end

  defp unfilter_paeth(row, prev, bpp) do
    row_bytes = :binary.bin_to_list(row)
    prev_bytes = :binary.bin_to_list(prev)

    {result, _} =
      Enum.reduce(Enum.with_index(row_bytes), {[], prev_bytes}, fn {byte, idx}, {acc, prev_list} ->
        left = if idx >= bpp, do: Enum.at(acc, idx - bpp), else: 0
        above = Enum.at(prev_list, idx, 0)
        upper_left = if idx >= bpp, do: Enum.at(prev_list, idx - bpp, 0), else: 0
        pr = paeth_predictor(left, above, upper_left)
        val = rem(byte + pr, 256)
        {acc ++ [val], prev_list}
      end)

    :binary.list_to_bin(result)
  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