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.HTTP
  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
          run_production_import(directory)
        else
          # Development: Direct database import
          Mix.Task.run("app.start")
          import_from_directory(directory)
        end

      _ ->
        show_usage_help()
    end
  end

  defp run_production_import(directory) 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)
  end

  defp show_usage_help do
    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

  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
    }

    case HTTP.post(__MODULE__, url,
           json: body,
           headers: [{"authorization", "Bearer #{api_key}"}],
           retry: :transient
         ) do
      {:ok, %{status: 200, body: response_body}} ->
        response_body

      {:ok, %{body: response_body}} ->
        Mix.shell().error("API request failed: #{inspect(response_body)}")
        raise "Failed to send batch to API"

      {:error, reason} ->
        Mix.shell().error("API request failed: #{inspect(reason)}")
        raise "Failed to send batch to API"
    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 ->
          # Generate UUID using Ecto's autogenerate mechanism
          {:ok, uuid} = Ecto.Type.cast(:binary_id, Ecto.UUID.generate())

          block
          |> Map.put(:id, uuid)
          |> 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

        build_block_entry(network, final_geoname_id, registered_id_int)

      _ ->
        nil
    end
  end

  defp build_block_entry(_network, nil, _registered_id_int), do: nil

  defp build_block_entry(network, final_geoname_id, registered_id_int) 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

  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