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
    {_opts, remaining_args} = parse_args(args)

    case remaining_args do
      [directory] ->
        # Check if TOWEROPS_KEY is set (production API import)
        api_key = System.get_env("TOWEROPS_KEY")

        if api_key do
          # Production: Send batches to API
          Mix.Task.run("app.config")
          {:ok, _} = Application.ensure_all_started(:req)

          import_via_api(directory, api_key)
        else
          # Development: Direct database import
          Mix.Task.run("app.start")
          import_from_directory(directory)
        end

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

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

          # Import to production via API
          TOWEROPS_KEY=your_superuser_api_key mix geoip.import ~/Downloads/GeoLite2-City-CSV_20260127/
        """)
    end
  end

  defp parse_args(args) do
    {opts, remaining, _invalid} =
      OptionParser.parse(args,
        strict: [database_url: :string],
        aliases: [d: :database_url]
      )

    {opts, remaining}
  end

  defp import_via_api(directory, api_key) 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")

    # Determine API URL (default to production)
    api_url = System.get_env("TOWEROPS_URL", "https://towerops.net")

    Mix.shell().info("Importing GeoIP database to #{api_url} via API...")

    # Parse and send locations first
    Mix.shell().info("  Processing locations file...")
    {locations_count, valid_geoname_ids} = send_locations_to_api(locations_file, api_url, api_key)
    Mix.shell().info("  Sent #{locations_count} locations")

    # Parse and send blocks
    Mix.shell().info("  Processing blocks file (this may take several minutes)...")
    blocks_count = send_blocks_to_api(blocks_file, valid_geoname_ids, api_url, api_key)
    Mix.shell().info("  Sent #{blocks_count} IP blocks")

    Mix.shell().info("✓ GeoIP database imported successfully to production!")
    Mix.shell().info("  #{locations_count} locations")
    Mix.shell().info("  #{blocks_count} IP blocks")
  end

  defp send_locations_to_api(file, api_url, api_key) 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} ->
      # Convert batch to JSON-serializable format
      data =
        Enum.map(batch, fn location ->
          %{
            "geoname_id" => location.geoname_id,
            "country_code" => location.country_code,
            "country_name" => location.country_name,
            "city_name" => location.city_name,
            "subdivision_1_name" => location.subdivision_1_name,
            "subdivision_2_name" => location.subdivision_2_name,
            "latitude" => location.latitude,
            "longitude" => location.longitude
          }
        end)

      # Send to API
      truncate = count == 0
      send_batch_to_api(api_url, api_key, "locations", data, truncate)

      # Track geoname_ids
      new_ids = MapSet.new(batch, & &1.geoname_id)
      {count + length(batch), MapSet.union(geoname_ids, new_ids)}
    end)
  end

  defp send_blocks_to_api(file, valid_geoname_ids, api_url, api_key) 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 ->
      # Convert batch to JSON-serializable format
      data =
        Enum.map(batch, fn block ->
          %{
            "network" => block.network,
            "start_ip_int" => block.start_ip_int,
            "end_ip_int" => block.end_ip_int,
            "geoname_id" => block.geoname_id,
            "registered_country_geoname_id" => block.registered_country_geoname_id
          }
        end)

      # Send to API
      truncate = count == 0
      send_batch_to_api(api_url, api_key, "blocks", data, truncate)

      count + length(batch)
    end)
  end

  defp send_batch_to_api(api_url, api_key, batch_type, data, truncate) do
    url = "#{api_url}/admin/api/geoip/import"

    body = %{
      "batch_type" => batch_type,
      "data" => data,
      "truncate" => truncate
    }

    response =
      Req.post!(url,
        json: body,
        headers: [{"authorization", "Bearer #{api_key}"}],
        retry: :transient
      )

    if response.status != 200 do
      Mix.shell().error("API request failed: #{inspect(response.body)}")
      raise "Failed to send batch to API"
    end

    response.body
  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