towerops/lib/mix/tasks/geoip.import.ex

defmodule Mix.Tasks.Geoip.Import do
  @shortdoc "Import MaxMind GeoLite2 City CSV database"

  @moduledoc """
  Import GeoLite2 City CSV database into PostgreSQL for country and city lookups.

  ## Usage

      # Import to local database (development)
      mix geoip.import ~/Downloads/GeoLite2-City-CSV_20260127/

      # Import to production database using DATABASE_URL environment variable
      DATABASE_URL=postgresql://user:pass@host/db mix geoip.import ~/Downloads/GeoLite2-City-CSV_20260127/

      # Import to production database using command-line flag
      mix geoip.import ~/Downloads/GeoLite2-City-CSV_20260127/ --database-url postgresql://user:pass@host/db

  ## What it does

  1. Reads GeoLite2-City-Locations-en.csv (geoname_id -> country, city, region)
  2. Inserts locations into geoip_locations table
  3. Reads GeoLite2-City-Blocks-IPv4.csv (IP ranges -> geoname_id)
  4. Inserts IP blocks into geoip_blocks table

  The resulting tables are used by Towerops.GeoIP for fast country and city lookups.

  ## Notes

  - Only processes IPv4 addresses (IPv6 support could be added)
  - Uses bulk inserts for performance (5,000 rows at a time)
  - Truncates existing data before import
  - Expands CIDR ranges to start/end IP integers for fast binary search
  - Can import to any PostgreSQL database using DATABASE_URL env var or --database-url flag
  """

  use Mix.Task

  alias Towerops.GeoIP.Block
  alias Towerops.GeoIP.Location
  alias Towerops.Repo

  # PostgreSQL has a 65,535 parameter limit
  # With ~8 fields per row: 65,535 / 8 = ~8,191 rows max
  # Using 5,000 for safety margin
  @batch_size 5_000

  @impl Mix.Task
  def run(args) do
    case args do
      [directory] ->
        Mix.Task.run("app.start")
        import_from_directory(directory)

      _ ->
        Mix.shell().error("""
        Usage: mix geoip.import <directory>

        Example:
          mix geoip.import ~/Downloads/GeoLite2-City-CSV_20260127/
        """)
    end
  end

  defp import_from_directory(directory) do
    directory = Path.expand(directory)
    blocks_file = Path.join(directory, "GeoLite2-City-Blocks-IPv4.csv")
    locations_file = Path.join(directory, "GeoLite2-City-Locations-en.csv")

    Mix.shell().info("Loading GeoLite2 City database from #{directory}...")

    # Truncate existing data
    Mix.shell().info("  Truncating existing GeoIP data...")
    Repo.delete_all(Block)
    Repo.delete_all(Location)

    # Load locations (geoname_id -> country, city, region)
    Mix.shell().info("  Reading locations file...")
    {locations_count, valid_geoname_ids} = load_locations(locations_file)
    Mix.shell().info("  Loaded #{locations_count} locations")

    # Load blocks and convert to IP ranges
    Mix.shell().info("  Reading blocks file (this may take a few minutes)...")
    blocks_count = load_blocks(blocks_file, valid_geoname_ids)
    Mix.shell().info("  Processed #{blocks_count} IP ranges with location data")

    Mix.shell().info("✓ GeoIP database imported successfully!")
    Mix.shell().info("  #{locations_count} locations (countries, cities, regions)")
    Mix.shell().info("  #{blocks_count} IP ranges covering IPv4 address space")
  end

  @doc false
  def load_locations(file) do
    file
    |> File.stream!()
    |> Stream.drop(1)
    |> Stream.map(&parse_location_line/1)
    |> Stream.reject(&is_nil/1)
    |> Stream.chunk_every(@batch_size)
    |> Enum.reduce({0, MapSet.new()}, fn batch, {count, geoname_ids} ->
      now = DateTime.truncate(DateTime.utc_now(), :second)

      entries =
        Enum.map(batch, fn location ->
          location
          |> Map.put(:inserted_at, now)
          |> Map.put(:updated_at, now)
        end)

      {inserted, _} = Repo.insert_all(Location, entries)

      # Track all geoname_ids for foreign key validation
      new_ids = MapSet.new(batch, & &1.geoname_id)
      {count + inserted, MapSet.union(geoname_ids, new_ids)}
    end)
  end

  defp parse_location_line(line) do
    # CSV format: geoname_id,locale_code,continent_code,continent_name,country_iso_code,country_name,subdivision_1_iso_code,subdivision_1_name,subdivision_2_iso_code,subdivision_2_name,city_name,metro_code,time_zone,is_in_european_union
    parts = String.split(line, ",")

    case parts do
      [
        geoname_id,
        _locale,
        _continent_code,
        _continent_name,
        country_code,
        country_name,
        _subdiv1_iso,
        subdiv1_name,
        _subdiv2_iso,
        subdiv2_name,
        city_name | _rest
      ] ->
        if geoname_id != "" and country_code != "" do
          %{
            geoname_id: String.to_integer(geoname_id),
            country_code: country_code,
            country_name: if(country_name == "", do: nil, else: country_name),
            city_name: if(city_name == "", do: nil, else: city_name),
            subdivision_1_name: if(subdiv1_name == "", do: nil, else: subdiv1_name),
            subdivision_2_name: if(subdiv2_name == "", do: nil, else: subdiv2_name),
            latitude: nil,
            longitude: nil
          }
        end

      _ ->
        nil
    end
  end

  @doc false
  def load_blocks(file, valid_geoname_ids) do
    file
    |> File.stream!()
    |> Stream.drop(1)
    |> Stream.map(&parse_block_line/1)
    |> Stream.reject(&is_nil/1)
    |> Stream.filter(fn block -> MapSet.member?(valid_geoname_ids, block.geoname_id) end)
    |> Stream.chunk_every(@batch_size)
    |> Enum.reduce(0, fn batch, count ->
      now = DateTime.truncate(DateTime.utc_now(), :second)

      entries =
        Enum.map(batch, fn block ->
          block
          |> Map.put(:id, Ecto.UUID.generate())
          |> Map.put(:inserted_at, now)
          |> Map.put(:updated_at, now)
        end)

      {inserted, _} = Repo.insert_all(Block, entries)
      count + inserted
    end)
  end

  defp parse_block_line(line) do
    # Format: network,geoname_id,registered_country_geoname_id,...
    parts = String.split(line, ",")

    case parts do
      [network, geoname_id, registered_country_id | _rest] ->
        # Try geoname_id first, fall back to registered_country_id
        geoname_id_int = parse_geoname_id(geoname_id)
        registered_id_int = parse_geoname_id(registered_country_id)

        final_geoname_id = geoname_id_int || registered_id_int

        if final_geoname_id do
          case parse_cidr(network) do
            {:ok, start_int, end_int} ->
              %{
                network: network,
                start_ip_int: start_int,
                end_ip_int: end_int,
                geoname_id: final_geoname_id,
                registered_country_geoname_id: registered_id_int
              }

            :error ->
              nil
          end
        end

      _ ->
        nil
    end
  end

  defp parse_geoname_id(""), do: nil

  defp parse_geoname_id(geoname_id_str) do
    case Integer.parse(geoname_id_str) do
      {geoname_id, _} -> geoname_id
      :error -> nil
    end
  end

  defp parse_cidr(cidr) do
    case String.split(cidr, "/") do
      [ip_str, prefix_str] ->
        with {:ok, {a, b, c, d}} <- parse_ip_address(ip_str),
             {prefix_len, ""} <- Integer.parse(prefix_str) do
          start_int = a * 256 * 256 * 256 + b * 256 * 256 + c * 256 + d
          # Calculate network mask and end IP
          host_bits = 32 - prefix_len
          num_addresses = Integer.pow(2, host_bits)
          end_int = start_int + num_addresses - 1

          {:ok, start_int, end_int}
        else
          _ -> :error
        end

      _ ->
        :error
    end
  end

  defp parse_ip_address(ip_str) do
    case :inet.parse_address(String.to_charlist(ip_str)) do
      {:ok, ip_tuple} -> {:ok, ip_tuple}
      _ -> :error
    end
  end
end