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 via API (requires TOWEROPS_KEY env var)
      mix geoip.import ~/Downloads/GeoLite2-City-CSV_20260127/ --production

  ## What it does

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

  ### Production Mode (--production flag)
  1. Reads and parses CSV files locally on your machine
  2. Sends processed data to API in batches of 5,000 records
  3. API endpoint handles database inserts
  4. Requires TOWEROPS_KEY environment variable (superuser API token)

  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 (first batch only in production mode)
  - Expands CIDR ranges to start/end IP integers for fast binary search
  - Production mode does not upload CSV files - only sends processed JSON data
  - Default production API URL: https://towerops.net (override with TOWEROPS_URL env var)
  """

  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] ->
        if opts[:production] do
          # Production: Send batches to API
          api_key = System.get_env("TOWEROPS_KEY")

          if !api_key do
            Mix.shell().error("Error: TOWEROPS_KEY environment variable must be set for production import")
            exit({:shutdown, 1})
          end

          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> [--production]

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

          # Import to production via API (requires TOWEROPS_KEY env var)
          mix geoip.import ~/Downloads/GeoLite2-City-CSV_20260127/ --production
        """)
    end
  end

  defp parse_args(args) do
    {opts, remaining, _invalid} =
      OptionParser.parse(args,
        strict: [production: :boolean],
        aliases: [p: :production]
      )

    {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