#!/usr/bin/env bash
set -euo pipefail

TARGET_MAIL_IP="50.116.60.173"
TARGET_MX_PRIORITY="10"
TARGET_SPF="v=spf1 ip4:50.116.60.173 -all"
TARGET_DMARC="v=dmarc1; p=quarantine; adkim=s; aspf=s"
TARGET_TTL="300"

ENV_FILE=".env"
MODE="dry-run"
ZONE_FILTER=""
JSON_OUT=""

usage() {
  cat <<USAGE
Usage: $(basename "$0") [options]

Plan or apply WHM DNS cleanup/standardization for discovered mail domains in each zone.

For each discovered domain (apex/addon/subdomain), enforce:
  - mail.<domain> A ${TARGET_MAIL_IP}
  - <domain> MX ${TARGET_MX_PRIORITY} mail.<domain>
  - <domain> TXT "${TARGET_SPF}"
  - _dmarc.<domain> TXT "${TARGET_DMARC}"

Default mode is dry-run (no changes).

Options:
  --env FILE          Path to env file (default: .env)
  --zone DOMAIN       Only process one zone
  --json-out FILE     Write planned actions JSON
  --dry-run           Dry-run mode (default)
  --apply             Apply planned changes
  -h, --help          Show this help

Required env vars in ENV file:
  user
  whm_api_key
  whm_api_url
USAGE
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    --env)
      ENV_FILE="$2"
      shift 2
      ;;
    --zone)
      ZONE_FILTER="$2"
      shift 2
      ;;
    --json-out)
      JSON_OUT="$2"
      shift 2
      ;;
    --dry-run)
      MODE="dry-run"
      shift
      ;;
    --apply)
      MODE="apply"
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown option: $1" >&2
      usage >&2
      exit 1
      ;;
  esac
done

require_cmd() {
  if ! command -v "$1" >/dev/null 2>&1; then
    echo "Missing required command: $1" >&2
    exit 1
  fi
}

require_cmd curl
require_cmd jq

if [[ ! -f "$ENV_FILE" ]]; then
  echo "Env file not found: $ENV_FILE" >&2
  exit 1
fi

set -a
# shellcheck disable=SC1090
source "$ENV_FILE"
set +a

WHM_USER="${user:-}"
WHM_API_KEY="${whm_api_key:-}"
WHM_API_URL="${whm_api_url:-}"

if [[ -z "$WHM_USER" || -z "$WHM_API_KEY" || -z "$WHM_API_URL" ]]; then
  echo "Missing one or more required env vars: user, whm_api_key, whm_api_url" >&2
  exit 1
fi

WHM_BASE="${WHM_API_URL%/}"
if [[ "$WHM_BASE" != */json-api ]]; then
  WHM_BASE="${WHM_BASE}/json-api"
fi

AUTH_HEADER="Authorization: whm ${WHM_USER}:${WHM_API_KEY}"

whm_api_call() {
  local func="$1"
  shift || true

  local url="${WHM_BASE}/${func}"
  local -a args
  args=(
    -fsS
    --get
    -H "$AUTH_HEADER"
    --data-urlencode "api.version=1"
  )

  while [[ $# -gt 0 ]]; do
    args+=(--data-urlencode "$1")
    shift
  done

  local body
  if ! body="$(curl "${args[@]}" "$url")"; then
    echo "WHM API transport error: ${func} (${url})" >&2
    return 1
  fi
  printf '%s\n' "$body"
}

check_api_ok() {
  local body="$1"
  local ok
  ok="$(jq -r '.metadata.result // 0' <<<"$body" 2>/dev/null || echo 0)"
  [[ "$ok" == "1" ]]
}

get_zones() {
  local body
  body="$(whm_api_call listzones)" || return 1
  if ! check_api_ok "$body"; then
    echo "WHM listzones failed:" >&2
    jq -r '.metadata.reason // . | tostring' <<<"$body" >&2 || true
    return 1
  fi

  jq -r '
    [
      .data.zone[]?
      | if type == "string" then .
        elif type == "object" then (.domain // .zone // empty)
        else empty end
    ]
    | map(select(type == "string" and length > 0))
    | unique
    | .[]
  ' <<<"$body"
}

plan_zone_actions() {
  local zone="$1"
  local body
  body="$(whm_api_call parse_dns_zone "zone=${zone}")" || {
    echo "WARN: parse_dns_zone transport failed for ${zone}" >&2
    echo '{"zone":"'"$zone"'","actions":[]}'
    return 0
  }

  if ! check_api_ok "$body"; then
    echo "WARN: parse_dns_zone failed for ${zone}: $(jq -r '.metadata.reason // "unknown error"' <<<"$body" 2>/dev/null || echo "unknown")" >&2
    echo '{"zone":"'"$zone"'","actions":[]}'
    return 0
  fi

  jq -c \
    --arg zone "$zone" \
    --arg target_mail_ip "$TARGET_MAIL_IP" \
    --arg target_mx_priority "$TARGET_MX_PRIORITY" \
    --arg target_spf "$TARGET_SPF" \
    --arg target_dmarc "$TARGET_DMARC" \
    --arg target_ttl "$TARGET_TTL" '
    def s: if . == null then "" else tostring end;
    def norm: ascii_downcase | sub("\\.$"; "");
    def b64d: try @base64d catch .;
    def canon_txt:
      s
      | gsub("^\"|\"$"; "")
      | gsub("\\s+"; " ")
      | sub("^ "; "")
      | sub(" $"; "")
      | ascii_downcase;
    def fqdn($name; $zone):
      ($name | s | norm) as $n
      | ($zone | norm) as $z
      | if $n == "" then ""
        elif $n == "@" then $z
        elif ($n == $z) or ($n | endswith("." + $z)) then $n
        else $n + "." + $z
        end;
    def rec_type:
      (if (.type // "") == "record" and (.record_type // "") != "" then .record_type else .type end)
      | s | ascii_upcase;
    def rec_name: (.dname_raw // .name // .dname // "" | s);
    def rec_line:
      if (.line_index != null) then
        (.line_index
          | if type == "number" then (. + 1)
            elif type == "string" and test("^[0-9]+$") then ((tonumber) + 1)
            else -1
            end)
      else
        ((.line // .Line // .linenum // -1)
          | if type == "number" then .
            elif type == "string" and test("^[0-9]+$") then tonumber
            else -1
            end)
      end;
    def rec_txt:
      (if (.data_b64 | type) == "array" and (.data_b64 | length) > 0 then
         (.data_b64 | map(tostring | b64d) | join(""))
       else
         (.txtdata // .record // .data // .value // "" | s)
       end);
    def rec_a:
      (if (.data_b64 | type) == "array" and (.data_b64 | length) > 0 then
         (.data_b64[0] | tostring | b64d)
       else
         (.address // .record // .value // "" | s)
       end)
      | norm;
    def rec_cname:
      (if (.data_b64 | type) == "array" and (.data_b64 | length) > 0 then
         (.data_b64[0] | tostring | b64d)
       else
         (.cname // .target // .record // .value // "" | s)
       end)
      | norm;
    def rec_mx_pref:
      (if (.data_b64 | type) == "array" and (.data_b64 | length) > 0 then
         (.data_b64[0] | tostring | b64d)
       else
         (.preference // .pref // .pri // .priority // "" | s)
       end);
    def rec_mx_host:
      (if (.data_b64 | type) == "array" and (.data_b64 | length) > 1 then
         (.data_b64[1] | tostring | b64d)
       else
         (.exchange // .target // .record // .value // "" | s)
       end)
      | norm;

    def mk_add($zone; $cat; $name; $rtype; $content; $params):
      {op:"add", zone:$zone, category:$cat, name:$name, type:$rtype, content:$content, params:$params};
    def mk_rm($zone; $cat; $r):
      {op:"remove", zone:$zone, category:$cat, line:$r.line, name:$r.name, type:$r.type, content:$r.content};

    [
      .data.payload[]?
      | . as $r
      | ($r | rec_type) as $t
      | select($t != "" and $t != "COMMENT" and $t != "CONTROL")
      | ($r | rec_name | fqdn(.; $zone)) as $n
      | select($n != "")
      | {
          line: ($r | rec_line),
          name: $n,
          type: $t,
          txt: ($r | rec_txt),
          txt_canon: ($r | rec_txt | canon_txt),
          a: ($r | rec_a),
          cname: ($r | rec_cname),
          mx_pref: ($r | rec_mx_pref),
          mx_host: ($r | rec_mx_host)
        }
      | . + {
          content:
            (if .type == "TXT" then .txt
             elif .type == "A" then .a
             elif .type == "CNAME" then .cname
             elif .type == "MX" then (.mx_pref + " " + .mx_host)
             else ""
             end)
        }
    ] as $recs
    |
    # Candidate mail domains discovered from existing mail-related records.
    (
      [ $recs[] | select(.type == "MX") | .name ]
      + [ $recs[] | select(.type == "TXT" and (.txt_canon | contains("v=spf1"))) | .name ]
      + [ $recs[] | select(.type == "TXT" and (.name | startswith("_dmarc.")) and (.txt_canon | contains("v=dmarc1"))) | (.name | sub("^_dmarc\\."; "")) ]
      + [ $recs[] | select((.type == "A" or .type == "CNAME") and (.name | startswith("mail."))) | (.name | sub("^mail\\."; "")) ]
    ) as $domain_candidates
    |
    ($zone | norm) as $z
    |
    (
      ($domain_candidates | map(select(type == "string" and length > 0 and (. == $z or endswith("." + $z)))) | unique)
      + [$z]
      | unique
    ) as $domains
    |
    {
      zone: $zone,
      actions:
      (
        [
          $domains[] as $d
          | ("mail." + $d) as $mailhost
          | ("_dmarc." + $d) as $dmarc_name

          | ($recs | map(select(.type == "A" and .name == $mailhost))) as $mail_a
          | ($recs | map(select(.type == "CNAME" and .name == $mailhost))) as $mail_cname
          | ($recs | map(select(.type == "MX" and .name == $d))) as $mx
          | ($recs | map(select(.type == "TXT" and .name == $d and (.txt_canon | contains("v=spf1"))))) as $spf
          | ($recs | map(select(.type == "TXT" and .name == $dmarc_name and (.txt_canon | contains("v=dmarc1"))))) as $dmarc

          | ($mail_a | map(select(.a == ($target_mail_ip | norm))) | sort_by(.line) | .[0]) as $keep_mail_a
          | ($mx | map(select(.mx_pref == $target_mx_priority and .mx_host == $mailhost)) | sort_by(.line) | .[0]) as $keep_mx
          | ($spf | map(select(.txt_canon == ($target_spf | canon_txt))) | sort_by(.line) | .[0]) as $keep_spf
          | ($dmarc | map(select(.txt_canon == ($target_dmarc | canon_txt))) | sort_by(.line) | .[0]) as $keep_dmarc

          | (
              # remove: mail A non-target or duplicate target
              [ $mail_a[] | select((.a != ($target_mail_ip | norm)) or (.line != ($keep_mail_a.line // -999999))) | mk_rm($zone; "mail_a"; .) ]
              # remove: mail CNAME always (mail host must be A)
              + [ $mail_cname[] | mk_rm($zone; "mail_cname"; .) ]
              # remove: MX not target or duplicate target
              + [ $mx[] | select((.line != ($keep_mx.line // -999999))) | mk_rm($zone; "mx"; .) ]
              # remove: SPF not target or duplicate target
              + [ $spf[] | select((.line != ($keep_spf.line // -999999))) | mk_rm($zone; "spf_txt"; .) ]
              # remove: DMARC not target or duplicate target
              + [ $dmarc[] | select((.line != ($keep_dmarc.line // -999999))) | mk_rm($zone; "dmarc_txt"; .) ]
              # add missing target mail A
              + (if $keep_mail_a == null then
                   [mk_add($zone; "mail_a"; $mailhost; "A"; $target_mail_ip; {
                     zone:$zone, name:($mailhost + "."), type:"A", address:$target_mail_ip, ttl:$target_ttl
                   })]
                 else [] end)
              # add missing target MX
              + (if $keep_mx == null then
                   [mk_add($zone; "mx"; $d; "MX"; ($target_mx_priority + " " + $mailhost); {
                     zone:$zone, name:($d + "."), type:"MX", preference:$target_mx_priority, exchange:($mailhost + "."), ttl:$target_ttl
                   })]
                 else [] end)
              # add missing target SPF
              + (if $keep_spf == null then
                   [mk_add($zone; "spf_txt"; $d; "TXT"; $target_spf; {
                     zone:$zone, name:($d + "."), type:"TXT", txtdata:$target_spf, ttl:$target_ttl
                   })]
                 else [] end)
              # add missing target DMARC
              + (if $keep_dmarc == null then
                   [mk_add($zone; "dmarc_txt"; $dmarc_name; "TXT"; $target_dmarc; {
                     zone:$zone, name:($dmarc_name + "."), type:"TXT", txtdata:$target_dmarc, ttl:$target_ttl
                   })]
                 else [] end)
            )
        ] | add // []
      )
    }
  ' <<<"$body"
}

remove_record_line() {
  local zone="$1"
  local line="$2"
  local body
  if ! body="$(whm_api_call removezonerecord "zone=${zone}" "line=${line}")"; then
    echo "APPLY FAIL remove zone=${zone} line=${line} reason=transport_error" >&2
    return 1
  fi
  if check_api_ok "$body"; then
    echo "APPLY OK  remove zone=${zone} line=${line}"
  else
    echo "APPLY FAIL remove zone=${zone} line=${line} reason=$(jq -r '.metadata.reason // "unknown"' <<<"$body" 2>/dev/null || echo "unknown")" >&2
    return 1
  fi
}

add_record() {
  local zone="$1"
  local name="$2"
  local type="$3"
  local ttl="$4"
  local body

  case "$type" in
    A)
      local address="$5"
      if ! body="$(whm_api_call addzonerecord "zone=${zone}" "name=${name}" "type=A" "address=${address}" "ttl=${ttl}")"; then
        echo "APPLY FAIL add zone=${zone} type=A name=${name} reason=transport_error" >&2
        return 1
      fi
      ;;
    MX)
      local preference="$5"
      local exchange="$6"
      if ! body="$(whm_api_call addzonerecord "zone=${zone}" "name=${name}" "type=MX" "preference=${preference}" "exchange=${exchange}" "ttl=${ttl}")"; then
        echo "APPLY FAIL add zone=${zone} type=MX name=${name} reason=transport_error" >&2
        return 1
      fi
      ;;
    TXT)
      local txtdata="$5"
      if ! body="$(whm_api_call addzonerecord "zone=${zone}" "name=${name}" "type=TXT" "txtdata=${txtdata}" "ttl=${ttl}")"; then
        echo "APPLY FAIL add zone=${zone} type=TXT name=${name} reason=transport_error" >&2
        return 1
      fi
      ;;
    *)
      echo "Unsupported add type: $type" >&2
      return 1
      ;;
  esac

  if check_api_ok "$body"; then
    echo "APPLY OK  add zone=${zone} type=${type} name=${name}"
  else
    echo "APPLY FAIL add zone=${zone} type=${type} name=${name} reason=$(jq -r '.metadata.reason // "unknown"' <<<"$body" 2>/dev/null || echo "unknown")" >&2
    return 1
  fi
}

main() {
  local -a zones
  if [[ -n "$ZONE_FILTER" ]]; then
    zones=("$ZONE_FILTER")
  else
    local zone_list
    if ! zone_list="$(get_zones)"; then
      exit 1
    fi
    mapfile -t zones <<<"$zone_list"
  fi

  if [[ ${#zones[@]} -eq 0 ]]; then
    echo "No zones found." >&2
    exit 1
  fi

  local actions_file
  actions_file="$(mktemp)"
  trap "rm -f '$actions_file'" EXIT
  local zone
  for zone in "${zones[@]}"; do
    local plan
    plan="$(plan_zone_actions "$zone")"
    jq -c '.actions[]?' <<<"$plan" >> "$actions_file"
  done
  local actions
  if [[ -s "$actions_file" ]]; then
    actions="$(jq -sc '.' "$actions_file")"
  else
    actions='[]'
  fi

  local total
  total="$(jq 'length' <<<"$actions")"

  echo "Mode: $MODE"
  echo "Desired mail A IP: $TARGET_MAIL_IP"
  echo "Desired MX: $TARGET_MX_PRIORITY mail.<domain>"
  echo "Desired SPF: $TARGET_SPF"
  echo "Desired DMARC: $TARGET_DMARC"
  echo "Zones scanned: ${#zones[@]}"
  echo "Actions queued: $total"

  if [[ "$total" -gt 0 ]]; then
    echo
    echo "Planned actions:"
    jq -r '
      sort_by(.zone, .op, .category, (.line // 0), .name)
      | .[]
      | if .op == "remove" then
          "- REMOVE zone=\(.zone) line=\(.line) category=\(.category) name=\(.name) type=\(.type) content=\(.content)"
        else
          "- ADD    zone=\(.zone) category=\(.category) name=\(.name) type=\(.type) content=\(.content)"
        end
    ' <<<"$actions"

    echo
    echo "Per-zone count:"
    jq -r '
      group_by(.zone)
      | map({zone: .[0].zone, count: length})
      | sort_by(.zone)
      | .[]
      | "- \(.zone): \(.count)"
    ' <<<"$actions"
  fi

  if [[ -n "$JSON_OUT" ]]; then
    jq '.' <<<"$actions" > "$JSON_OUT"
    echo
    echo "Wrote action JSON: $JSON_OUT"
  fi

  if [[ "$MODE" == "dry-run" ]]; then
    echo
    echo "Dry-run only. No DNS records were changed."
    return 0
  fi

  if [[ "$total" -eq 0 ]]; then
    echo
    echo "Nothing to apply."
    return 0
  fi

  echo
  echo "Applying removals first..."
  local failures=0

  while IFS=$'\t' read -r z l; do
    if ! remove_record_line "$z" "$l"; then
      failures=$((failures + 1))
    fi
  done < <(jq -r '[.[] | select(.op == "remove")] | sort_by(.zone, .line) | reverse | .[] | [.zone, (.line|tostring)] | @tsv' <<<"$actions")

  echo
  echo "Applying additions..."
  while IFS=$'\t' read -r z t n ttl a1 a2; do
    case "$t" in
      A)
        if ! add_record "$z" "$n" "$t" "$ttl" "$a1"; then
          failures=$((failures + 1))
        fi
        ;;
      MX)
        if ! add_record "$z" "$n" "$t" "$ttl" "$a1" "$a2"; then
          failures=$((failures + 1))
        fi
        ;;
      TXT)
        if ! add_record "$z" "$n" "$t" "$ttl" "$a1"; then
          failures=$((failures + 1))
        fi
        ;;
      *)
        echo "APPLY FAIL add zone=${z} reason=unsupported_type_${t}" >&2
        failures=$((failures + 1))
        ;;
    esac
  done < <(jq -r '
    [.[] | select(.op == "add")]
    | sort_by(.zone, .category, .name)
    | .[]
    | [
        .zone,
        .type,
        .params.name,
        (.params.ttl // "300"),
        (.params.address // .params.preference // .params.txtdata // ""),
        (.params.exchange // "")
      ]
    | @tsv
  ' <<<"$actions")

  echo
  if [[ "$failures" -gt 0 ]]; then
    echo "Apply complete with failures: $failures" >&2
    exit 1
  fi

  echo "Apply complete: all queued actions processed successfully."
  trap - EXIT
  rm -f "$actions_file"
}

main