#' Add Country Information to a Dataset Based on Latitude/Longitude
#'
#' This function adds a new column \code{country} by performing a fast spatial
#' join between point coordinates (\code{lat}, \code{lon}) and world country
#' polygons from \pkg{rnaturalearth}. Geometry is handled via \pkg{sf}.
#'
#' The input data must contain numeric columns named \code{lat} and \code{lon}
#' in WGS84 (EPSG:4326).  Rows with \code{NA} coordinates receive
#' \code{country = NA}.
#'
#' After assigning countries, Germany is automatically split into
#' \code{"Germany-North"} and \code{"Germany-South"} using a latitude threshold:
#' \itemize{
#'   \item lat >= 50 → \code{"Germany-North"}
#'   \item lat  < 50 → \code{"Germany-South"}
#' }
#'
#' If the input already contains a \code{country} column it is removed and
#' replaced by the spatially derived one.
#'
#' @param dt A \code{data.frame} or \code{data.table} containing columns
#'   \code{lat} and \code{lon}.
#' @param keep_geometry Logical. If \code{TRUE}, the sf geometry column is kept.
#'   Default is \code{FALSE}.
#'
#' @return A \code{data.table} with a new column \code{country}.
#' @author Matthias Templ
#'
#' @examples
#' \donttest{
#' small <- data.table(
#'   lat = c(48.2, 51.0, 47.5),
#'   lon = c(16.3, 10.1, 8.5)
#' )
#'
#' pep_with_country <- add_country(small)
#' pep_with_country
#' }
#'
#' @export
add_country <- function(dt, keep_geometry = FALSE) {

  # Check input
  if (!("lat" %in% names(dt)) || !("lon" %in% names(dt))) {
    stop("Input data must contain columns 'lat' and 'lon'.", call. = FALSE)
  }

  # Check required packages
  if (!requireNamespace("sf", quietly = TRUE))
    stop("Package 'sf' is required. Install it with install.packages('sf').",
         call. = FALSE)
  if (!requireNamespace("rnaturalearth", quietly = TRUE))
    stop("Package 'rnaturalearth' is required. Install it with install.packages('rnaturalearth').",
         call. = FALSE)
  if (!requireNamespace("rnaturalearthdata", quietly = TRUE))
    stop("Package 'rnaturalearthdata' is required. Install it with install.packages('rnaturalearthdata').",
         call. = FALSE)

  message("Assigning countries based on coordinates ...")

  # Convert to data.table (copy to avoid modifying original)
  dt <- data.table::copy(data.table::as.data.table(dt))

  # Remove pre-existing country column to avoid duplicates
  if ("country" %in% names(dt)) {
    dt[, country := NULL]
  }

  # Handle NA coordinates: set aside rows, join valid rows only
  has_na_coords <- is.na(dt$lon) | is.na(dt$lat)
  n_na <- sum(has_na_coords)
  if (n_na > 0) {
    message(sprintf("Note: %d row(s) have NA coordinates and will get country = NA.", n_na))
  }

  if (all(has_na_coords)) {
    dt[, country := NA_character_]
    return(dt)
  }

  # Split into valid and NA-coordinate rows
  if (n_na > 0) {
    dt[, .row_idx := .I]
    dt_valid <- dt[!has_na_coords]
    dt_na <- dt[has_na_coords]
    dt_na[, country := NA_character_]
  } else {
    dt_valid <- dt
  }

  # Convert to sf POINT object with WGS84 coordinates
  pep_sf <- sf::st_as_sf(
    dt_valid,
    coords = c("lon", "lat"),
    crs = 4326,
    remove = FALSE
  )

  # Load world polygons
  world <- rnaturalearth::ne_countries(scale = "medium", returnclass = "sf")[, c("name")]
  world <- sf::st_make_valid(world)

  # Reproject both datasets (faster join in planar CRS)
  pep_sf <- sf::st_transform(pep_sf, 3857)
  world  <- sf::st_transform(world, 3857)

  sf::st_agr(world) <- "constant"

  # Spatial join
  pep_joined <- sf::st_join(pep_sf, world, left = TRUE)

  # Convert back to data.table
  dt_result <- data.table::as.data.table(pep_joined)

  # Remove geometry unless requested
  if (!keep_geometry) {
    dt_result[, geometry := NULL]
  }

  data.table::setnames(dt_result, "name", "country")

  # Recombine with NA-coordinate rows
  if (n_na > 0) {
    dt_result <- data.table::rbindlist(list(dt_result, dt_na),
                                       use.names = TRUE, fill = TRUE)
    data.table::setorderv(dt_result, ".row_idx")
    dt_result[, .row_idx := NULL]
  }

  # Germany split
  dt_result[
    country == "Germany" & lat >= 50,
    country := "Germany-North"
  ]

  dt_result[
    country == "Germany" & lat < 50,
    country := "Germany-South"
  ]

  dt_result
}
