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

# Subscribe to one MQTT message for FJR country and append a two-letter ISO country
# code to fjr_country.json only when it differs from the last JSON record.
#
# Cron example, every 5 minutes:
# */5 * * * * /home/hbrain/source/hasgd/bin/fjr_country_mqtt_logger.sh
#
# Optional environment variables:
#   MQTT_HOST          MQTT broker host, default: 192.168.0.203
#   MQTT_PORT          MQTT broker port, default: 1883
#   MQTT_TOPIC         MQTT topic, default: fjr/country_code
#   MQTT_USER          MQTT username, optional
#   MQTT_PASSWORD      MQTT password, optional
#   MQTT_TIMEOUT                 Seconds to wait for one message, default: 30
#   FJR_COUNTRY_JSON_DIR         Output directory, default: repo root
#   FJR_COUNTRY_JSON_FILENAME    Output filename, default: fjr_country.json
#   FJR_COUNTRY_JSON             Full output JSON path; overrides DIR/FILENAME

SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
REPO_DIR="$(cd -- "${SCRIPT_DIR}/.." && pwd)"

MQTT_HOST="${MQTT_HOST:-192.168.0.203}"
MQTT_PORT="${MQTT_PORT:-1883}"
MQTT_TOPIC="${MQTT_TOPIC:-fjr/country_code}"
MQTT_TIMEOUT="${MQTT_TIMEOUT:-30}"
JSON_DIR="${FJR_COUNTRY_JSON_DIR:-${REPO_DIR}}"
JSON_FILENAME="${FJR_COUNTRY_JSON_FILENAME:-fjr_country.json}"
JSON_FILE="${FJR_COUNTRY_JSON:-${JSON_DIR%/}/${JSON_FILENAME}}"

command -v python3 >/dev/null 2>&1 || {
  echo "ERROR: python3 not found." >&2
  exit 1
}

mkdir -p -- "$(dirname -- "${JSON_FILE}")"

# Create the JSON file on first run.
if [[ ! -e "${JSON_FILE}" ]]; then
  printf '[]\n' >"${JSON_FILE}"
fi

python3 - "${JSON_FILE}" <<'PY'
import json
import sys

path = sys.argv[1]
try:
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
except Exception as exc:
    print(f"ERROR: {path} is not valid JSON: {exc}", file=sys.stderr)
    sys.exit(1)

if not isinstance(data, list):
    print(f"ERROR: {path} exists but is not a JSON array.", file=sys.stderr)
    sys.exit(1)
PY

# Avoid concurrent cron runs modifying the file at the same time.
exec 9>"${JSON_FILE}.lock"
if ! flock -n 9; then
  echo "Another fjr_country logger run is active; exiting."
  exit 0
fi

payload="$(
  MQTT_HOST="${MQTT_HOST}" \
  MQTT_PORT="${MQTT_PORT}" \
  MQTT_TOPIC="${MQTT_TOPIC}" \
  MQTT_TIMEOUT="${MQTT_TIMEOUT}" \
  MQTT_USER="${MQTT_USER:-}" \
  MQTT_PASSWORD="${MQTT_PASSWORD:-}" \
  python3 - <<'PY'
import os
import random
import socket
import sys

host = os.environ["MQTT_HOST"]
port = int(os.environ["MQTT_PORT"])
topic = os.environ["MQTT_TOPIC"]
timeout = float(os.environ["MQTT_TIMEOUT"])
username = os.environ.get("MQTT_USER") or None
password = os.environ.get("MQTT_PASSWORD") or None
client_id = f"fjr_country_logger_{random.randrange(65536):04x}"


def enc_len(value):
    out = bytearray()
    while True:
        digit = value % 128
        value //= 128
        if value:
            digit |= 128
        out.append(digit)
        if not value:
            return bytes(out)


def field(value):
    data = str(value).encode("utf-8")
    return len(data).to_bytes(2, "big") + data


def read_remaining(sock):
    value = 0
    multiplier = 1
    while True:
        data = sock.recv(1)
        if not data:
            raise EOFError("connection closed")
        digit = data[0]
        value += (digit & 127) * multiplier
        if not digit & 128:
            return value
        multiplier *= 128


def recv_packet(sock):
    header = sock.recv(1)
    if not header:
        raise EOFError("connection closed")
    remaining = read_remaining(sock)
    data = bytearray()
    while len(data) < remaining:
        chunk = sock.recv(remaining - len(data))
        if not chunk:
            raise EOFError("connection closed")
        data.extend(chunk)
    return header[0], bytes(data)


flags = 0x02  # clean session
payload = field(client_id)
if username is not None:
    flags |= 0x80
    payload += field(username)
if password is not None:
    flags |= 0x40
    payload += field(password)

variable_header = b"\x00\x04MQTT\x04" + bytes([flags]) + b"\x00\x3c"
connect_packet = b"\x10" + enc_len(len(variable_header) + len(payload)) + variable_header + payload
subscribe_payload = field(topic) + b"\x00"
subscribe_packet = b"\x82" + enc_len(2 + len(subscribe_payload)) + b"\x00\x01" + subscribe_payload

try:
    with socket.create_connection((host, port), timeout=timeout) as sock:
        sock.settimeout(timeout)
        sock.sendall(connect_packet)
        packet_type, data = recv_packet(sock)
        if packet_type >> 4 != 2 or len(data) < 2 or data[1] != 0:
            code = data[1] if len(data) > 1 else "?"
            print(f"ERROR: MQTT connection failed, CONNACK code {code}", file=sys.stderr)
            raise SystemExit(1)

        sock.sendall(subscribe_packet)
        while True:
            packet_type, data = recv_packet(sock)
            if packet_type >> 4 != 3:
                continue
            topic_len = int.from_bytes(data[:2], "big")
            idx = 2 + topic_len
            qos = (packet_type >> 1) & 0x03
            if qos:
                idx += 2
            sys.stdout.write(data[idx:].decode("utf-8", "replace"))
            break
except socket.timeout:
    print(f"ERROR: no MQTT message received from {topic} within {timeout:g}s.", file=sys.stderr)
    raise SystemExit(1)
except Exception as exc:
    print(f"ERROR: MQTT subscribe failed: {exc}", file=sys.stderr)
    raise SystemExit(1)
PY
)"

# Accept plain text country names/codes, a JSON string, or a JSON object containing
# country_code/country/state/value/address.country(_code). Store ISO alpha-2 code.
country_code="$(PAYLOAD="${payload}" python3 - <<'PY'
import json
import os
import re
import sys
import unicodedata

payload = os.environ.get("PAYLOAD", "").strip()

# Common FJR/Europe mappings plus common English/local country names.
COUNTRIES = {
    "albania": "AL", "shqiperia": "AL", "shqipëria": "AL",
    "andorra": "AD",
    "austria": "AT", "osterreich": "AT", "österreich": "AT",
    "belarus": "BY",
    "belgium": "BE", "belgique": "BE", "belgie": "BE", "belgië": "BE",
    "bosnia and herzegovina": "BA", "bosnia": "BA", "bosna i hercegovina": "BA",
    "bulgaria": "BG", "balgariya": "BG", "bǎlgariya": "BG", "българия": "BG",
    "croatia": "HR", "hrvatska": "HR",
    "cyprus": "CY",
    "czechia": "CZ", "czech republic": "CZ", "cesko": "CZ", "česko": "CZ",
    "denmark": "DK", "danmark": "DK",
    "estonia": "EE", "eesti": "EE",
    "finland": "FI", "suomi": "FI",
    "france": "FR",
    "germany": "DE", "deutschland": "DE",
    "greece": "GR", "hellas": "GR", "ellada": "GR", "ελλάδα": "GR",
    "hungary": "HU", "magyarorszag": "HU", "magyarország": "HU",
    "iceland": "IS", "island": "IS", "ísland": "IS",
    "ireland": "IE", "eire": "IE", "éire": "IE",
    "italy": "IT", "italia": "IT",
    "kosovo": "XK",
    "latvia": "LV", "latvija": "LV",
    "liechtenstein": "LI",
    "lithuania": "LT", "lietuva": "LT",
    "luxembourg": "LU", "letzebuerg": "LU", "lëtzebuerg": "LU",
    "malta": "MT",
    "moldova": "MD",
    "monaco": "MC",
    "montenegro": "ME", "crna gora": "ME",
    "netherlands": "NL", "the netherlands": "NL", "nederland": "NL", "holland": "NL",
    "north macedonia": "MK", "macedonia": "MK", "severna makedonija": "MK",
    "norway": "NO", "norge": "NO", "noreg": "NO",
    "poland": "PL", "polska": "PL",
    "portugal": "PT",
    "romania": "RO", "românia": "RO",
    "san marino": "SM",
    "serbia": "RS", "srbija": "RS", "србија": "RS",
    "slovakia": "SK", "slovensko": "SK",
    "slovenia": "SI", "slovenija": "SI",
    "spain": "ES", "espana": "ES", "españa": "ES",
    "sweden": "SE", "sverige": "SE",
    "switzerland": "CH", "schweiz": "CH", "suisse": "CH", "svizzera": "CH",
    "turkey": "TR", "turkiye": "TR", "türkiye": "TR",
    "ukraine": "UA", "ukraina": "UA", "україна": "UA",
    "united kingdom": "GB", "great britain": "GB", "uk": "GB",
    "vatican city": "VA", "holy see": "VA",
}


def normalize(value):
    value = str(value or "").strip()
    value = re.sub(r"\s+", " ", value)
    return value


def ascii_key(value):
    value = normalize(value).lower()
    value = unicodedata.normalize("NFKD", value)
    value = "".join(ch for ch in value if not unicodedata.combining(ch))
    return value


def as_code(value):
    value = normalize(value)
    if not value:
        return ""

    # Already ISO alpha-2, including XK for Kosovo.
    if re.fullmatch(r"[A-Za-z]{2}", value):
        return value.upper()

    lower = value.lower()
    if lower in COUNTRIES:
        return COUNTRIES[lower]

    key = ascii_key(value)
    if key in COUNTRIES:
        return COUNTRIES[key]

    return ""


def get_path(obj, *path):
    cur = obj
    for key in path:
        if not isinstance(cur, dict) or key not in cur:
            return None
        cur = cur[key]
    return cur

try:
    parsed = json.loads(payload)
except Exception:
    code = as_code(payload)
    if code:
        print(code)
        raise SystemExit(0)
    print(f"ERROR: cannot map country to ISO code: {payload}", file=sys.stderr)
    raise SystemExit(2)

candidates = []
if isinstance(parsed, str):
    candidates.append(parsed)
elif isinstance(parsed, dict):
    for path in [
        ("country_code",),
        ("address", "country_code"),
        ("country",),
        ("address", "country"),
        ("state",),
        ("value",),
    ]:
        value = get_path(parsed, *path)
        if value is not None:
            candidates.append(value)
else:
    candidates.append(parsed)

for candidate in candidates:
    code = as_code(candidate)
    if code:
        print(code)
        break
else:
    print(f"ERROR: cannot map country to ISO code from payload: {payload}", file=sys.stderr)
    raise SystemExit(2)
PY
)"

if [[ -z "${country_code}" || "${country_code}" == "null" ]]; then
  echo "ERROR: MQTT payload did not contain a usable country value." >&2
  exit 1
fi

# IMPORTANT: only inspect the last record in the JSON file.
last_country_code="$(python3 - "${JSON_FILE}" <<'PY'
import json
import sys

with open(sys.argv[1], "r", encoding="utf-8") as f:
    data = json.load(f)

if not data:
    print("")
    raise SystemExit(0)

last = data[-1]
if isinstance(last, dict):
    print(str(last.get("country_code", last.get("code", last.get("country", last.get("value", ""))))).upper())
else:
    print(str(last).upper())
PY
)"

if [[ "${country_code}" == "${last_country_code}" ]]; then
  echo "No change: ${country_code}"
  exit 0
fi

timestamp="$(date -Is)"
tmp_file="$(mktemp "${JSON_FILE}.tmp.XXXXXX")"
cleanup() {
  [[ -n "${tmp_file:-}" && -e "${tmp_file}" ]] && rm -f -- "${tmp_file}"
}
trap cleanup EXIT

python3 - "${JSON_FILE}" "${tmp_file}" "${timestamp}" "${country_code}" <<'PY'
import json
import sys

src, dst, timestamp, country_code = sys.argv[1:5]
with open(src, "r", encoding="utf-8") as f:
    data = json.load(f)

# Append only the new current value; do not scan older history.
data.append({"timestamp": timestamp, "country_code": country_code})

with open(dst, "w", encoding="utf-8") as f:
    json.dump(data, f, ensure_ascii=False, indent=2)
    f.write("\n")
PY

mv -- "${tmp_file}" "${JSON_FILE}"
tmp_file=""

echo "Appended: ${country_code}"