#!/usr/bin/env python3
"""
Multi-engine CLI: fetch company suggestions and insert into global_companies_registry
via the Athena Gateway. When --country-code EE, also inserts into ee_companies (RIK rows).

Engines:
  - northdata — https://www.northdata.com/suggest.json (use with any ISO country filter)
  - kvk       — Dutch KVK search API (requires --country-code NL)

Examples:
  python scripts/scrape_global_registry.py --search-engine northdata --country-code EE \\
      --query "Finance" --no-insert

  python scripts/scrape_global_registry.py --search-engine kvk --country-code NL \\
      --query "secra" --output nl_global_seed.csv

  python scripts/scrape_global_registry.py --randomize-query --search-engine northdata \\
      --country-code NL --query fi --min-query-length 3 --max-query-length 8

  NL-only dual table (`nl_companies` + global) is `nl-scrape` (Rust), from repo root:

  cargo run --release --bin nl-scrape -- --query secra

  Global registry only (`dexter-scrape`):

  cargo run --release --bin dexter-scrape -- \\
      --athena-url="https://cluster.athena-db.com" \\
      --athena-key="$ATHENA_API_KEY" \\
      --athena-client="dexter" \\
      --country-code="NL" --search-engine="kvk" \\
      --max-query-length=3

  Or Python: scripts/example_kvk_nl_randomized.* wrap the above; legacy script:

  python scripts/scrape_global_registry.py \\
      --athena-url="https://cluster.athena-db.com" \\
      --athena-key="$ATHENA_API_KEY" \\
      --athena-client="dexter" \\
      --country-code="NL" --search-engine="kvk" \\
      --randomize-query --max-query-length=3

  Optional NL VAT enrichment from Bizzy (KVK flow only):

  python scripts/scrape_global_registry.py --search-engine kvk --country-code NL \\
      --query secra --vat-id-search-engine bizzy

  Randomized mode uses (1) names from the last API response, (2) random rows from
  nl_companies.csv when present, (3) scripts/words.txt — never busy-loops on duplicate
  queries like "B.V."; use --randomize-no-seed-csv to skip the CSV.

Environment:
  ATHENA_API_KEY, ATHENA_USER_ID, ATHENA_COMPANY_ID, ATHENA_ORGANIZATION_ID (optional)
  NORTHDATA_COOKIE — optional Cookie header if suggest.json returns blocked responses

Dependencies: aiohttp, colorama (same as scrape_nl_comps.py)
"""

from __future__ import annotations

import argparse
import asyncio
import csv
import hashlib
import io
import json
import os
import random
import re
import inspect
import subprocess
import sys
import time
from functools import wraps
from typing import Any, Literal
from urllib.parse import quote, unquote, urlparse

_scripts_dir = os.path.dirname(os.path.abspath(__file__))
if _scripts_dir not in sys.path:
    sys.path.insert(0, _scripts_dir)

try:
    import aiohttp
except ImportError:
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "aiohttp", "colorama"]
    )
    import aiohttp

from colorama import Fore, Style, init as colorama_init

# KVK fetch/parse (global-only insert here; no changes to scrape_nl_comps.py)
from scrape_nl_comps import fetch_all_kvk_companies

colorama_init(autoreset=True)

TELEMETRY_ENABLED = False

# Columns on public.global_companies_registry (omit id, created_at, auto defaults)
GLOBAL_REGISTRY_KEYS = frozenset(
    {
        "name",
        "company_number",
        "rights_form",
        "country_code",
        "status",
        "last_update",
        "description",
        "region",
        "rights_form_code",
        "state",
        "name_english",
        "language",
        "hash",
        "address",
        "global_company_id",
        "start_date",
        "legal_persons",
        "contact_data",
        "activities",
        "activity_code_table",
        "establishment_date",
        "establishment_type",
        "entity_type",
        "vat_id",
        "email",
        "iban",
        "phone",
        "postal_code",
        "house_number",
        "province",
        "city",
        "street",
        "country",
        "website",
    }
)

# public.ee_companies (omit id, created_at)
EE_COMPANIES_KEYS = frozenset(
    {
        "registrikood",
        "name",
        "legal_form",
        "status",
        "location",
        "vat_id",
        "register_reference",
        "source_url",
    }
)

_rate_scraped = 0
_rate_inserted = 0
_rate_last_time: float | None = None
_rate_last_scraped = 0
_rate_last_inserted = 0

# Athena gateway "API key policy unavailable" backoff (shared across inserts)
_GATEWAY_POLICY_PAUSE_SEC = 15
_GATEWAY_COOLDOWN_SEC = 150.0
_GATEWAY_MAX_CONCURRENT_INSERTS = 5
_ATHENA_INSERT_JITTER_SEC = (0.2, 0.5)

_gateway_backoff_lock = asyncio.Lock()
_gateway_cooldown_until = 0.0
_gateway_insert_cap = _GATEWAY_MAX_CONCURRENT_INSERTS


def _is_gateway_api_key_policy_error(detail: Any) -> bool:
    if isinstance(detail, dict):
        msg = str(detail.get("message", ""))
        if "API key policy unavailable" in msg:
            return True
    return "Gateway API key enforcement" in str(detail)


async def _coordinate_gateway_backoff(max_workers: int) -> None:
    """One 15s pause under lock; extend insert cool-down (throttled concurrency)."""
    global _gateway_cooldown_until, _gateway_insert_cap
    async with _gateway_backoff_lock:
        now = time.monotonic()
        if now < _gateway_cooldown_until:
            return
        ansi_warning(
            "Gateway API key policy unavailable; pausing "
            f"{_GATEWAY_POLICY_PAUSE_SEC}s and throttling inserts for "
            f"~{int(_GATEWAY_COOLDOWN_SEC)}s."
        )
        await asyncio.sleep(_GATEWAY_POLICY_PAUSE_SEC)
        _gateway_cooldown_until = time.monotonic() + _GATEWAY_COOLDOWN_SEC
        _gateway_insert_cap = max(
            1, min(max_workers, _GATEWAY_MAX_CONCURRENT_INSERTS)
        )


def _rate_add_scraped(n: int) -> None:
    global _rate_scraped
    _rate_scraped += n


def _rate_add_inserted(n: int = 1) -> None:
    global _rate_inserted
    _rate_inserted += n


def _rate_reset() -> None:
    global _rate_scraped, _rate_inserted
    _rate_scraped = 0
    _rate_inserted = 0


async def _rate_logger_loop() -> None:
    global _rate_last_time, _rate_last_scraped, _rate_last_inserted
    _rate_last_time = time.monotonic()
    _rate_last_scraped = _rate_scraped
    _rate_last_inserted = _rate_inserted
    while True:
        await asyncio.sleep(15)
        now = time.monotonic()
        elapsed = now - _rate_last_time
        if elapsed <= 0:
            continue
        delta_s = _rate_scraped - _rate_last_scraped
        delta_i = _rate_inserted - _rate_last_inserted
        scrapes_per_min = (delta_s / elapsed) * 60 if elapsed else 0
        inserts_per_min = (delta_i / elapsed) * 60 if elapsed else 0
        ansi_info(
            f"Rate (last {elapsed:.0f}s): scrapes {scrapes_per_min:.1f}/min, inserts {inserts_per_min:.1f}/min"
        )
        _rate_last_time = now
        _rate_last_scraped = _rate_scraped
        _rate_last_inserted = _rate_inserted


def ansi_info(msg: str) -> None:
    print(
        f"{Fore.MAGENTA}[GLOBAL-REG]{Style.RESET_ALL} {Fore.CYAN}[INFO]{Style.RESET_ALL} {msg}"
    )


def ansi_success(msg: str) -> None:
    print(
        f"{Fore.MAGENTA}[GLOBAL-REG]{Style.RESET_ALL} {Fore.GREEN}[SUCCESS]{Style.RESET_ALL} {msg}"
    )


def ansi_warning(msg: str) -> None:
    print(
        f"{Fore.MAGENTA}[GLOBAL-REG]{Style.RESET_ALL} {Fore.YELLOW}[WARNING]{Style.RESET_ALL} {msg}"
    )


def ansi_error(msg: str) -> None:
    print(
        f"{Fore.MAGENTA}[GLOBAL-REG]{Style.RESET_ALL} {Fore.RED}[ERROR]{Style.RESET_ALL} {msg}"
    )


def telemetry(func):
    if inspect.iscoroutinefunction(func):

        @wraps(func)
        async def async_wrapper(*args, **kwargs):
            if TELEMETRY_ENABLED:
                start = time.perf_counter()
                result = await func(*args, **kwargs)
                elapsed_ms = (time.perf_counter() - start) * 1000
                ansi_info(f"{func.__name__} took {elapsed_ms:.2f}ms")
                return result
            return await func(*args, **kwargs)

        return async_wrapper

    @wraps(func)
    def sync_wrapper(*args, **kwargs):
        if TELEMETRY_ENABLED:
            start = time.perf_counter()
            result = func(*args, **kwargs)
            elapsed_ms = (time.perf_counter() - start) * 1000
            ansi_info(f"{func.__name__} took {elapsed_ms:.2f}ms")
            return result
        return func(*args, **kwargs)

    return sync_wrapper


def filter_global_payload(d: dict[str, Any]) -> dict[str, Any]:
    out = {k: v for k, v in d.items() if k in GLOBAL_REGISTRY_KEYS and v is not None}
    out.pop("global_company_id", None)
    return out


def filter_ee_payload(d: dict[str, Any]) -> dict[str, Any]:
    return {k: v for k, v in d.items() if k in EE_COMPANIES_KEYS and v is not None}


def parse_ee_registrikood_from_segment(segment: str) -> str | None:
    """
    Extract Estonian äriregistri kood from Northdata path segment (e.g. 'RIK 11283797').
    Returns digits-only string or None.
    """
    s = (segment or "").strip()
    if not s:
        return None
    m = re.match(r"^RIK\s*(\d{7,9})\s*$", s, re.IGNORECASE)
    if m:
        return m.group(1)
    if re.fullmatch(r"\d{7,9}", s):
        return s
    return None


def infer_ee_legal_form(name: str) -> str | None:
    """Best-effort legal form from Estonian company name."""
    if not name:
        return None
    if re.search(r"\bOÜ\b|Osaühing", name, re.IGNORECASE):
        return "OÜ"
    if re.search(r"\bAS\b|Aktsiaselts", name, re.IGNORECASE):
        return "AS"
    if re.search(r"\bFIE\b", name, re.IGNORECASE):
        return "FIE"
    if re.search(r"\bMTÜ\b", name, re.IGNORECASE):
        return "MTÜ"
    if re.search(r"\bUÜ\b", name, re.IGNORECASE):
        return "UÜ"
    if re.search(r"\bTÜ\b|Täisühing", name, re.IGNORECASE):
        return "TÜ"
    return None


def map_global_row_to_ee_company(row: dict[str, Any]) -> dict[str, Any] | None:
    """
    Map a global_companies_registry-shaped row to ee_companies.
    Only for Estonian RIK identifiers; returns None if registrikood cannot be parsed.
    """
    if (row.get("country_code") or "").upper() != "EE":
        return None
    segment = (row.get("company_number") or "").strip()
    rk = parse_ee_registrikood_from_segment(segment)
    if not rk:
        return None
    name = (row.get("name") or "").strip()
    if not name:
        return None
    ee: dict[str, Any] = {
        "registrikood": rk,
        "name": name,
        "legal_form": infer_ee_legal_form(name),
        "status": row.get("status"),
        "location": row.get("description") or row.get("address"),
        "vat_id": row.get("vat_id"),
        "register_reference": segment if segment else None,
        "source_url": row.get("website"),
    }
    return filter_ee_payload(ee)


def entity_type_from_dutch_rights_form(rights_form: str | None) -> str | None:
    if not rights_form:
        return None
    mapping = {
        "Eenmanszaak": "EMZ",
        "Vennootschap Onder Firma": "VOF",
        "Maatschap": "MS",
        "Stichting": "ST",
        "Besloten Vennootschap": "BV",
    }
    return mapping.get(rights_form)


def northdata_register_label(segment: str) -> str | None:
    """Best-effort register label from the last path segment of a Northdata URL."""
    s = segment.strip()
    if not s:
        return None
    prefixes = (
        ("Companies House ", "Companies House"),
        ("Siren ", "Siren"),
        ("RIK ", "RIK"),
        ("CRO ", "CRO"),
        ("GEMI ", "GEMI"),
        ("ON ", "ON"),
        ("HRB ", "HRB"),
        ("UID ", "UID"),
        ("MCIT ", "MCIT"),
        ("Handelsregister ", "Handelsregister"),
    )
    for prefix, label in prefixes:
        if s.startswith(prefix):
            return label
    if s.upper().startswith("KRS"):
        return "KRS"
    return None


@telemetry
def northdata_company_number_from_url(url: str) -> str | None:
    try:
        path = urlparse(url).path.strip("/")
        if not path:
            return None
        last = path.split("/")[-1]
        decoded = unquote(last).strip()
        return decoded or None
    except Exception:
        return None


def northdata_status_from_title(title: str | None) -> str | None:
    if not title:
        return None
    if "\u271d" in title or "\u271d\ufe0e" in title or "✝" in title:
        return "dissolved"
    return None


@telemetry
def map_northdata_hit_to_global(
    hit: dict[str, Any], country_code: str
) -> dict[str, Any] | None:
    """
    Map a Northdata suggest result to global_companies_registry.
    Returns None if company_number cannot be derived.
    """
    if hit.get("type") != "company":
        return None
    name = hit.get("name") or hit.get("title")
    url = hit.get("url")
    if not name or not url:
        return None
    segment = northdata_company_number_from_url(url)
    if not segment:
        return None

    cc = country_code.upper()
    company_number = segment
    hash_input = f"{name}|{company_number}|{cc}"
    hash_val = hashlib.sha256(hash_input.encode("utf-8")).hexdigest()

    description = hit.get("description")
    title = hit.get("title")
    status = northdata_status_from_title(title)
    register = northdata_register_label(segment)
    activity_code_table = register

    global_row: dict[str, Any] = {
        "name": name,
        "company_number": company_number,
        "rights_form": None,
        "country_code": cc,
        "status": status,
        "last_update": None,
        "description": description,
        "region": None,
        "rights_form_code": None,
        "state": None,
        "name_english": None,
        "language": None,
        "hash": hash_val,
        "address": description,
        "start_date": None,
        "legal_persons": None,
        "contact_data": None,
        "activities": None,
        "activity_code_table": activity_code_table,
        "establishment_date": None,
        "establishment_type": None,
        "entity_type": None,
        "vat_id": None,
        "email": None,
        "iban": None,
        "phone": None,
        "postal_code": None,
        "house_number": None,
        "province": None,
        "city": None,
        "street": None,
        "country": None,
        "website": url,
    }
    return filter_global_payload(global_row)


@telemetry
def map_kvk_company_to_global(
    company: dict[str, Any], country_code: str
) -> dict[str, Any] | None:
    """Map a nl_companies-shaped dict to global_companies_registry (country_code from CLI)."""
    kvk = company.get("kvk_number")
    name = company.get("name")
    if kvk is None or not name:
        return None

    cc = country_code.upper()
    street = company.get("street")
    house_number = company.get("house_number")
    postal_code = company.get("postal_code")
    city = company.get("city")
    province = company.get("province")
    country = company.get("country") or "Nederland"

    address_parts = []
    if street:
        address_parts.append(str(street))
    if house_number is not None:
        address_parts.append(str(house_number))
    if postal_code:
        address_parts.append(str(postal_code))
    if city:
        address_parts.append(str(city))
    if province:
        address_parts.append(str(province))
    if country:
        address_parts.append(str(country))
    address = ", ".join(address_parts) if address_parts else None

    company_number = str(kvk)
    hash_input = f"{name}|{company_number}|{cc}"
    hash_val = hashlib.sha256(hash_input.encode("utf-8")).hexdigest()

    active = company.get("active")
    status = "active" if active else "inactive" if active is not None else None

    activities = company.get("activities")
    activity_code_table = "SBI" if company.get("main_sbi_code") else None
    establishment_type = (
        "main"
        if company.get("establishment") is True
        else "branch" if company.get("establishment") is False else None
    )

    hn = house_number
    if hn is not None:
        hn = str(hn)

    global_row: dict[str, Any] = {
        "name": name,
        "company_number": company_number,
        "rights_form": company.get("rights_form"),
        "country_code": cc,
        "status": status,
        "last_update": None,
        "description": company.get("main_activity_description")
        or company.get("main_activity"),
        "region": province,
        "rights_form_code": company.get("rights_form_code"),
        "state": province,
        "name_english": None,
        "language": "nl",
        "hash": hash_val,
        "address": address,
        "start_date": None,
        "legal_persons": None,
        "contact_data": None,
        "activities": activities,
        "activity_code_table": activity_code_table,
        "establishment_date": None,
        "establishment_type": establishment_type,
        "entity_type": None,
        "vat_id": company.get("vat_id"),
        "email": None,
        "iban": None,
        "phone": None,
        "postal_code": postal_code,
        "house_number": hn,
        "province": province,
        "city": city,
        "street": street,
        "country": country,
        "website": None,
    }
    return filter_global_payload(global_row)


def build_athena_headers(
    athena_key: str,
    athena_client: str,
    user_id=None,
    company_id=None,
    organization_id=None,
    strip_nulls=None,
) -> dict[str, str]:
    headers: dict[str, str] = {
        "Content-Type": "application/json",
        "X-Athena-Client": athena_client,
    }
    if athena_key:
        headers["apikey"] = athena_key
        headers["x-api-key"] = athena_key
    if user_id:
        headers["X-User-Id"] = user_id
    if company_id:
        headers["X-Company-Id"] = company_id
    if organization_id:
        headers["X-Organization-Id"] = organization_id
    if strip_nulls is not None:
        headers["X-Strip-Nulls"] = "true" if strip_nulls else "false"
    return headers


AthenaInsertStatus = Literal["ok", "duplicate", "error", "exception"]


async def athena_insert(
    session: aiohttp.ClientSession,
    athena_url: str,
    headers: dict[str, str],
    table_name: str,
    payload: dict[str, Any],
) -> tuple[AthenaInsertStatus, Any]:
    endpoint = f"{athena_url.rstrip('/')}/gateway/insert"
    if time.monotonic() < _gateway_cooldown_until:
        lo, hi = _ATHENA_INSERT_JITTER_SEC
        await asyncio.sleep(random.uniform(lo, hi))
    try:
        async with session.put(
            endpoint,
            headers=headers,
            data=json.dumps({"table_name": table_name, "insert_body": payload}),
            timeout=aiohttp.ClientTimeout(total=15),
        ) as resp:
            status = resp.status
            if status in (200, 201, 204):
                return "ok", None
            if status == 409:
                return "duplicate", await resp.text()
            try:
                return "error", await resp.json()
            except Exception:
                return "error", await resp.text()
    except Exception as exc:
        return "exception", str(exc)


async def _athena_insert_with_gateway_retry(
    session: aiohttp.ClientSession,
    athena_url: str,
    headers: dict[str, str],
    table_name: str,
    payload: dict[str, Any],
    batch_max_workers: int,
) -> tuple[AthenaInsertStatus, Any]:
    """PUT insert; on gateway policy error, coordinate backoff and retry once."""
    status, detail = await athena_insert(
        session, athena_url, headers, table_name, payload
    )
    if status == "error" and _is_gateway_api_key_policy_error(detail):
        await _coordinate_gateway_backoff(batch_max_workers)
        status, detail = await athena_insert(
            session, athena_url, headers, table_name, payload
        )
    return status, detail


def _athena_headers_or_build(
    athena_key, athena_client, user_id, company_id, organization_id, prebuilt=None
):
    if prebuilt is not None:
        return prebuilt
    return build_athena_headers(
        athena_key,
        athena_client,
        user_id=user_id,
        company_id=company_id,
        organization_id=organization_id,
    )


@telemetry
async def insert_global_company_via_athena(
    session,
    global_company: dict[str, Any],
    athena_url: str,
    athena_key: str,
    athena_client: str,
    hash_cache: set[str],
    user_id=None,
    company_id=None,
    organization_id=None,
    headers=None,
    batch_max_workers: int = 25,
) -> None:
    hkey = global_company.get("hash")
    if not hkey:
        return
    if hkey in hash_cache:
        return

    payload = {k: v for k, v in global_company.items() if v is not None}
    rights_form = global_company.get("rights_form")
    et = entity_type_from_dutch_rights_form(rights_form)
    if et:
        payload["entity_type"] = et

    for key in ("id", "created_at", "global_company_id"):
        payload.pop(key, None)

    h = _athena_headers_or_build(
        athena_key, athena_client, user_id, company_id, organization_id, headers
    )

    status, detail = await _athena_insert_with_gateway_retry(
        session,
        athena_url,
        h,
        "global_companies_registry",
        payload,
        batch_max_workers,
    )
    if status == "ok":
        _rate_add_inserted(1)
        ansi_success(
            f"Inserted: ({global_company.get('company_number')}) {global_company.get('name')} [global_companies_registry]"
        )
        hash_cache.add(hkey)
    elif status == "duplicate":
        hash_cache.add(hkey)
    elif status == "error":
        ansi_error(
            f"Failed to insert: {global_company.get('name')} ({global_company.get('company_number')}): {detail} [global_companies_registry]"
        )
    else:
        ansi_error(
            f"Exception during insert: {global_company.get('name')} ({global_company.get('company_number')}): {detail} [global_companies_registry]"
        )


@telemetry
async def insert_ee_company_via_athena(
    session,
    ee_row: dict[str, Any],
    athena_url: str,
    athena_key: str,
    athena_client: str,
    registrikood_cache: set[str],
    user_id=None,
    company_id=None,
    organization_id=None,
    headers=None,
    batch_max_workers: int = 25,
) -> None:
    """Insert into ee_companies (duplicate registrikood → 409, same as global hash)."""
    rk = ee_row.get("registrikood")
    if not rk:
        return
    if rk in registrikood_cache:
        return

    payload = {k: v for k, v in ee_row.items() if v is not None}
    for key in ("id", "created_at"):
        payload.pop(key, None)

    h = _athena_headers_or_build(
        athena_key, athena_client, user_id, company_id, organization_id, headers
    )

    status, detail = await _athena_insert_with_gateway_retry(
        session, athena_url, h, "ee_companies", payload, batch_max_workers
    )
    if status == "ok":
        _rate_add_inserted(1)
        ansi_success(f"Inserted: ({rk}) {ee_row.get('name')} [ee_companies]")
        registrikood_cache.add(rk)
    elif status == "duplicate":
        registrikood_cache.add(rk)
    elif status == "error":
        ansi_error(
            f"Failed to insert ee_companies: {ee_row.get('name')} ({rk}): {detail}"
        )
    else:
        ansi_error(
            f"Exception during ee_companies insert: {ee_row.get('name')} ({rk}): {detail}"
        )


NORTHDATA_DEFAULT_UA = (
    "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
    "(KHTML, like Gecko) Chrome/146.0.0.0 Safari/537.36"
)

BIZZY_COMPANY_URL_PREFIX = "https://bizzy.ai/nl/nl/"
BIZZY_VAT_HTML_SNIPPET_LEN = 8000
BIZZY_MAX_CONCURRENT_REQUESTS = 3
BIZZY_REQUEST_PAUSE_SEC = 0.4
BIZZY_HTTP_TIMEOUT_SEC = 20
BIZZY_VAT_PATTERN = re.compile(r"NL\s*(\d{9}B\d{2})", re.IGNORECASE)
_BIZZY_FETCH_SEM = asyncio.Semaphore(BIZZY_MAX_CONCURRENT_REQUESTS)


def _kvk_number_digits(kvk: Any) -> str | None:
    """Return KVK as digit string, or None if missing or non-numeric."""
    if kvk is None:
        return None
    if isinstance(kvk, int):
        kvk_s = str(kvk)
    else:
        kvk_s = str(kvk).strip()
    return kvk_s if kvk_s.isdigit() else None


def _parse_vat_from_bizzy_html(html: str) -> str | None:
    marker = 'data-cy="company-subheader-vat-number"'
    idx = html.find(marker)
    chunk = (
        html[idx : idx + BIZZY_VAT_HTML_SNIPPET_LEN] if idx != -1 else html
    )
    m = BIZZY_VAT_PATTERN.search(chunk)
    if not m:
        return None
    return "NL" + m.group(1).upper()


@telemetry
async def enrich_nl_companies_vat_from_bizzy(
    session: aiohttp.ClientSession,
    companies: list[dict[str, Any]],
) -> None:
    """Fill vat_id from Bizzy HTML when KVK did not return btwNummer."""
    headers = {
        "accept": "text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8",
        "accept-language": "en-US,en;q=0.9,nl;q=0.8",
        "user-agent": NORTHDATA_DEFAULT_UA,
    }
    timeout = aiohttp.ClientTimeout(total=BIZZY_HTTP_TIMEOUT_SEC)

    async def fetch_one(c: dict[str, Any]) -> None:
        if c.get("vat_id"):
            return
        kvk_s = _kvk_number_digits(c.get("kvk_number"))
        if not kvk_s:
            return
        async with _BIZZY_FETCH_SEM:
            url = f"{BIZZY_COMPANY_URL_PREFIX}{kvk_s}"
            try:
                async with session.get(
                    url,
                    headers=headers,
                    timeout=timeout,
                    allow_redirects=True,
                ) as resp:
                    if resp.status != 200:
                        return
                    html = await resp.text()
                vat = _parse_vat_from_bizzy_html(html)
                if vat:
                    c["vat_id"] = vat
            except (aiohttp.ClientError, asyncio.TimeoutError, UnicodeDecodeError):
                pass
            finally:
                await asyncio.sleep(BIZZY_REQUEST_PAUSE_SEC)

    await asyncio.gather(*(fetch_one(c) for c in companies))


def northdata_headers(cookie: str | None) -> dict[str, str]:
    h = {
        "accept": "application/json, text/javascript, */*; q=0.01",
        "accept-language": "en-US,en;q=0.9",
        "referer": "https://www.northdata.com/",
        "user-agent": NORTHDATA_DEFAULT_UA,
        "x-requested-with": "XMLHttpRequest",
    }
    if cookie:
        h["cookie"] = cookie
    return h


@telemetry
async def fetch_northdata_suggest(
    session: aiohttp.ClientSession,
    base_url: str,
    query: str,
    countries: str,
    cookie: str | None,
    max_retries: int = 4,
) -> list[dict[str, Any]]:
    """GET suggest.json; returns raw result dicts (may include non-company types)."""
    root = base_url.rstrip("/")
    q = quote(query, safe="")
    # countries: ISO2 or empty
    cc_param = quote(countries, safe="")
    url = f"{root}/suggest.json?query={q}&countries={cc_param}"

    headers = northdata_headers(cookie)
    delay = 1.0
    last_err = None
    for attempt in range(max_retries):
        try:
            async with session.get(
                url,
                headers=headers,
                timeout=aiohttp.ClientTimeout(total=30),
            ) as resp:
                if resp.status == 200:
                    data = await resp.json()
                    return list(data.get("results") or [])
                if resp.status in (429, 503):
                    last_err = f"HTTP {resp.status}"
                    await asyncio.sleep(delay)
                    delay = min(delay * 2, 30)
                    continue
                text = await resp.text()
                last_err = f"HTTP {resp.status}: {text[:200]}"
                ansi_warning(last_err)
                return []
        except Exception as e:
            last_err = str(e)
            await asyncio.sleep(delay)
            delay = min(delay * 2, 30)
    ansi_error(f"Northdata suggest failed after retries: {last_err}")
    return []


def save_global_rows_csv(
    rows: list[dict[str, Any]], filename: str, append: bool
) -> None:
    if not rows:
        ansi_warning("No rows to save to CSV.")
        return
    fieldnames = sorted({k for r in rows for k in r.keys()})
    mode = "a" if append else "w"
    file_exists = os.path.exists(filename) and os.path.getsize(filename) > 0
    with open(filename, mode=mode, newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
        if mode == "w" or not file_exists:
            w.writeheader()
        for r in rows:
            w.writerow({k: r.get(k) for k in fieldnames})
    ansi_success(f"Saved {len(rows)} row(s) to '{filename}'.")


def load_hash_cache_from_csv(filename: str) -> set[str]:
    if not os.path.exists(filename):
        return set()
    out: set[str] = set()
    with open(filename, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        if not reader.fieldnames or "hash" not in reader.fieldnames:
            return set()
        for row in reader:
            h = row.get("hash")
            if h:
                out.add(h)
    return out


# Tokens to skip when building search prefixes (legal forms / noise)
_SKIP_NAME_TOKENS = frozenset(
    {
        "b",
        "bv",
        "b.v",
        "b.v.",
        "nv",
        "vof",
        "cv",
        "ab",
        "kb",
        "as",
        "asa",
        "oy",
        "oyj",
        "oü",
        "ou",
        "ltd",
        "llc",
        "inc",
        "plc",
        "gmbh",
        "ag",
        "sas",
        "sa",
        "srl",
        "spa",
        "sl",
        "slu",
        "ooo",
        "zoo",
        "see",
    }
)

# Whole-query rejects (normalized lowercase, no spaces)
_BAD_QUERY_WHOLE = frozenset(
    {
        "b.v.",
        "bv",
        "nv",
        "vof",
        "ab",
        "oy",
        "ltd",
        "llc",
        "inc",
    }
)

# Randomized mode: auto-expand max query length until this ceiling (at least --max-query-length).
RANDOM_QUERY_MAX_LEN_CAP = 64


def _normalize_query_token(t: str) -> str:
    return re.sub(r"^[\W_]+|[\W_]+$", "", t, flags=re.UNICODE).lower().rstrip(".")


def meaningful_tokens_from_name(name: str) -> list[str]:
    """Split company name into tokens, dropping common legal-form noise."""
    if not name or not name.strip():
        return []
    parts = re.split(r"\s+", name.strip())
    out: list[str] = []
    for p in parts:
        if not p:
            continue
        norm = _normalize_query_token(p)
        if not norm or norm in _SKIP_NAME_TOKENS:
            continue
        if len(norm) == 1 and not norm.isalnum():
            continue
        out.append(p.strip())
    return out


def query_candidates_from_name(name: str, min_len: int, max_len: int) -> list[str]:
    """
    Build search strings from a company name: joined token prefixes and
    length-graded prefixes (longer first → slightly more specific, fewer collisions).
    """
    tokens = meaningful_tokens_from_name(name)
    seen_local: set[str] = set()
    cands: list[str] = []

    def add(s: str) -> None:
        s = s.strip()
        if not s:
            return
        if len(s) < min_len or len(s) > max_len:
            return
        low = s.lower().strip()
        if low in _BAD_QUERY_WHOLE:
            return
        if not any(ch.isalpha() for ch in s):
            return
        key = low
        if key not in seen_local:
            seen_local.add(key)
            cands.append(s)

    if tokens:
        for n in range(1, len(tokens) + 1):
            chunk = " ".join(tokens[:n]).strip()
            if len(chunk) > max_len:
                chunk = chunk[:max_len].rstrip()
            add(chunk)
        joined = " ".join(tokens)
        upper = min(max_len, len(joined))
        for L in range(upper, min_len - 1, -1):
            add(joined[:L])
    else:
        stripped = name.strip()
        upper = min(max_len, len(stripped))
        for L in range(upper, min_len - 1, -1):
            add(stripped[:L])
    return cands


def query_is_acceptable(q: str, min_len: int, max_len: int) -> bool:
    q = q.strip()
    if len(q) < min_len or len(q) > max_len:
        return False
    if q.lower() in _BAD_QUERY_WHOLE:
        return False
    return any(c.isalpha() for c in q)


def load_words_file(path: str) -> list[str]:
    if not path or not os.path.isfile(path):
        return []
    words: list[str] = []
    with open(path, encoding="utf-8", errors="replace") as f:
        for line in f:
            w = line.strip()
            if w and not w.startswith("#"):
                words.append(w)
    return words


def build_csv_line_byte_offsets(path: str) -> list[int]:
    """Start byte offset of each line (line 0 = first line, usually header)."""
    offsets: list[int] = []
    pos = 0
    with open(path, "rb") as f:
        for line in f:
            offsets.append(pos)
            pos += len(line)
    return offsets


def random_csv_row_at_line_index(
    path: str, offsets: list[int], fieldnames: list[str]
) -> dict[str, Any] | None:
    if len(offsets) < 2:
        return None
    j = random.randint(1, len(offsets) - 1)
    with open(path, "rb") as f:
        f.seek(offsets[j])
        data_b = f.readline()
    try:
        data = data_b.decode("utf-8")
    except UnicodeDecodeError:
        return None
    if not data.strip():
        return None
    reader = csv.DictReader(io.StringIO(data), fieldnames=fieldnames)
    try:
        return next(reader)
    except StopIteration:
        return None


def prepare_seed_csv_access(path: str | None) -> tuple[list[int], list[str]] | None:
    if not path or not os.path.isfile(path):
        return None
    try:
        with open(path, "rb") as f:
            header_line_b = f.readline()
        if not header_line_b:
            return None
        header_s = header_line_b.decode("utf-8", errors="replace")
        fieldnames = next(csv.reader([header_s]))
        offsets = build_csv_line_byte_offsets(path)
        if len(offsets) < 2:
            return None
        ansi_info(
            f"Randomize seed CSV: {path} ({len(offsets) - 1} data lines) for query suggestions."
        )
        return (offsets, fieldnames)
    except OSError as e:
        ansi_warning(f"Could not open --randomize-seed-csv {path}: {e}")
        return None


def pick_next_search_query(
    seen_queries: set[str],
    last_result_names: list[str],
    words: list[str],
    seed_csv: tuple[list[int], list[str]] | None,
    seed_csv_path: str | None,
    text_seed: str,
    min_len: int,
    max_len: int,
    csv_samples: int = 40,
) -> tuple[str, bool]:
    """
    Returns (query, from_guaranteed_fallback).

    Steps 1–4 prefer textual queries not in seen_queries. Step 5 is numeric/time fallback;
    from_guaranteed_fallback is True only for step 5 so the caller can widen max_len first.
    """
    # 1) Last batch names (breadth: shuffle token-derivations)
    pool: list[str] = []
    for name in last_result_names:
        pool.extend(query_candidates_from_name(name, min_len, max_len))
    random.shuffle(pool)
    for c in dict.fromkeys(pool):
        if c not in seen_queries and query_is_acceptable(c, min_len, max_len):
            return c, False

    # 2) Random rows from seed CSV (nl_companies, etc.)
    if seed_csv and seed_csv_path:
        offsets, fieldnames = seed_csv
        for _ in range(csv_samples):
            row = random_csv_row_at_line_index(seed_csv_path, offsets, fieldnames)
            if not row:
                continue
            name = (row.get("name") or row.get("naam") or "").strip()
            if not name:
                continue
            cands = query_candidates_from_name(name, min_len, max_len)
            random.shuffle(cands)
            for c in cands:
                if c not in seen_queries and query_is_acceptable(c, min_len, max_len):
                    return c, False

    # 3) Word list
    bag = list(words)
    random.shuffle(bag)
    for w in bag:
        c = w.strip()[:max_len].strip()
        if c not in seen_queries and query_is_acceptable(c, min_len, max_len):
            return c, False

    # 4) Text seed
    for c in query_candidates_from_name(text_seed, min_len, max_len):
        if c not in seen_queries:
            return c, False

    # 5) Guaranteed unique (never spin)
    for _ in range(1000):
        n = random.randint(10**5, 10**9)
        c = str(n)[:max_len]
        if len(c) >= min_len and c not in seen_queries:
            ansi_info(f"No novel text queries left; using numeric fallback '{c}'.")
            return c, True
    c = f"q{time.time_ns()}"[:max_len]
    if len(c) < min_len:
        c = (c + "0" * min_len)[:max_len]
    ansi_info(f"Using time-based fallback query '{c}'.")
    return c, True


@telemetry
async def process_global_rows(
    session,
    rows: list[dict[str, Any]],
    athena_url: str,
    athena_key: str,
    athena_client: str,
    insert: bool,
    max_workers: int,
    hash_cache: set[str],
    country_code: str,
    ee_registrikood_cache: set[str] | None = None,
    athena_user_id=None,
    athena_company_id=None,
    athena_organization_id=None,
) -> None:
    if not insert or not rows:
        if insert and not rows:
            ansi_warning("No rows to insert.")
        return

    cc = (country_code or "").upper()
    ee_cache = ee_registrikood_cache if cc == "EE" else None

    headers = build_athena_headers(
        athena_key,
        athena_client,
        user_id=athena_user_id,
        company_id=athena_company_id,
        organization_id=athena_organization_id,
    )
    now = time.monotonic()
    sem_n = (
        min(max_workers, _gateway_insert_cap)
        if now < _gateway_cooldown_until
        else max_workers
    )
    if sem_n < max_workers:
        ansi_info(
            f"Insert concurrency capped at {sem_n} (--max-workers {max_workers}) during gateway cool-down."
        )
    semaphore = asyncio.Semaphore(sem_n)

    async def worker(row: dict[str, Any]) -> None:
        async with semaphore:
            try:
                tasks = [
                    insert_global_company_via_athena(
                        session,
                        row,
                        athena_url,
                        athena_key,
                        athena_client,
                        hash_cache,
                        user_id=athena_user_id,
                        company_id=athena_company_id,
                        organization_id=athena_organization_id,
                        headers=headers,
                        batch_max_workers=max_workers,
                    )
                ]
                if ee_cache is not None:
                    ee_row = map_global_row_to_ee_company(row)
                    if ee_row:
                        tasks.append(
                            insert_ee_company_via_athena(
                                session,
                                ee_row,
                                athena_url,
                                athena_key,
                                athena_client,
                                ee_cache,
                                user_id=athena_user_id,
                                company_id=athena_company_id,
                                organization_id=athena_organization_id,
                                headers=headers,
                                batch_max_workers=max_workers,
                            )
                        )
                await asyncio.gather(*tasks)
            except Exception as e:
                ansi_error(f"Exception during insert: {e}")

    await asyncio.gather(*(asyncio.create_task(worker(r)) for r in rows))


async def randomize_query_loop(
    session,
    engine: str,
    country_code: str,
    query_seed: str,
    output_file: str,
    athena_url: str,
    athena_key: str,
    athena_client: str,
    insert: bool,
    max_workers: int,
    min_query_length: int,
    max_query_length: int,
    kvk_api_url: str,
    northdata_base_url: str,
    northdata_cookie: str | None,
    words_file: str | None,
    seed_csv_path: str | None,
    athena_user_id=None,
    athena_company_id=None,
    athena_organization_id=None,
    vat_id_search_engine: str = "none",
) -> None:
    seen_queries: set[str] = set()
    hash_cache = load_hash_cache_from_csv(output_file)
    ee_registrikood_cache: set[str] | None = (
        set() if country_code.upper() == "EE" else None
    )
    last_result_names: list[str] = []
    words = load_words_file(words_file or "")
    if words_file and not words:
        ansi_warning(f"No words loaded from '{words_file}' (missing or empty).")
    seed_access = prepare_seed_csv_access(seed_csv_path)

    if max_query_length < min_query_length:
        ansi_warning(
            f"--max-query-length ({max_query_length}) < --min-query-length ({min_query_length}); using min for both."
        )
        max_query_length = min_query_length

    text_seed = query_seed.strip() or "fi"
    effective_max = max_query_length
    expand_cap = max(RANDOM_QUERY_MAX_LEN_CAP, max_query_length)

    while True:
        while True:
            query, from_fallback = pick_next_search_query(
                seen_queries,
                last_result_names,
                words,
                seed_access,
                seed_csv_path,
                text_seed,
                min_query_length,
                effective_max,
            )
            if not from_fallback:
                break
            if effective_max >= expand_cap:
                break
            prev = effective_max
            effective_max += 1
            ansi_info(
                f"No novel text queries up to length {prev}; raising max query length to {effective_max}."
            )
        seen_queries.add(query)
        ansi_info(
            f"Randomized query: '{query}' (seen {len(seen_queries)} unique queries)"
        )

        if engine == "northdata":
            raw = await fetch_northdata_suggest(
                session, northdata_base_url, query, country_code, northdata_cookie
            )
            mapped = []
            for hit in raw:
                g = map_northdata_hit_to_global(hit, country_code)
                if g:
                    mapped.append(g)
            rows = mapped
        else:
            companies = await fetch_all_kvk_companies(session, kvk_api_url, query)
            if vat_id_search_engine == "bizzy":
                await enrich_nl_companies_vat_from_bizzy(session, companies)
            rows = []
            for company in companies:
                g = map_kvk_company_to_global(company, country_code)
                if g:
                    rows.append(g)

        _rate_add_scraped(len(rows))
        ansi_info(f"Fetched {len(rows)} mapped row(s) for query '{query}'.")

        await process_global_rows(
            session,
            rows,
            athena_url,
            athena_key,
            athena_client,
            insert,
            max_workers,
            hash_cache,
            country_code,
            ee_registrikood_cache,
            athena_user_id=athena_user_id,
            athena_company_id=athena_company_id,
            athena_organization_id=athena_organization_id,
        )

        if rows:
            save_global_rows_csv(rows, output_file, append=True)
            for r in rows:
                hk = r.get("hash")
                if hk:
                    hash_cache.add(hk)
                if ee_registrikood_cache is not None:
                    ee = map_global_row_to_ee_company(r)
                    if ee and ee.get("registrikood"):
                        ee_registrikood_cache.add(ee["registrikood"])
            names = []
            for r in rows:
                n = (r.get("name") or "").strip()
                if n:
                    names.append(n)
            # Prefer names from the latest successful response for the next query expansion
            last_result_names = names[:120]
        else:
            ansi_warning("No rows to save.")


async def main_async() -> None:
    global TELEMETRY_ENABLED

    parser = argparse.ArgumentParser(
        description="Fetch companies (Northdata or KVK), insert into global_companies_registry; "
        "with --country-code EE also insert into ee_companies (RIK-parsed rows)."
    )
    parser.add_argument(
        "--country-code",
        type=str,
        required=True,
        help="ISO 3166-1 alpha-2 country code (e.g. NL, EE)",
    )
    parser.add_argument(
        "--search-engine",
        type=str,
        choices=("northdata", "kvk"),
        required=True,
        help="Data source: northdata (suggest.json) or kvk (Dutch API)",
    )
    parser.add_argument(
        "--query",
        type=str,
        default="",
        help="Search query (required for non-randomized run; used as seed for --randomize-query)",
    )
    parser.add_argument(
        "--output",
        type=str,
        default="global_registry_scrape.csv",
        help="CSV path for mapped global rows",
    )
    parser.add_argument(
        "--no-insert",
        action="store_true",
        help="Do not call Athena; only fetch/map (and optionally CSV)",
    )
    parser.add_argument("--max-workers", type=int, default=25)
    parser.add_argument(
        "--kvk-api-url",
        type=str,
        default="https://web-api.kvk.nl/zoeken/v3/search",
    )
    parser.add_argument(
        "--northdata-base-url",
        type=str,
        default="https://www.northdata.com",
    )
    parser.add_argument(
        "--northdata-cookie-file",
        type=str,
        default=None,
        help="Path to file whose contents are sent as Cookie (else NORTHDATA_COOKIE env)",
    )
    parser.add_argument(
        "--athena-url",
        type=str,
        default="https://mirror3.athena-db.com",
    )
    parser.add_argument(
        "--athena-key",
        type=str,
        default=os.getenv("ATHENA_API_KEY", ""),
    )
    parser.add_argument(
        "--athena-client",
        type=str,
        default="railway_xbp",
    )
    parser.add_argument(
        "--athena-user-id", type=str, default=os.getenv("ATHENA_USER_ID")
    )
    parser.add_argument(
        "--athena-company-id", type=str, default=os.getenv("ATHENA_COMPANY_ID")
    )
    parser.add_argument(
        "--athena-organization-id",
        type=str,
        default=os.getenv("ATHENA_ORGANIZATION_ID"),
    )
    parser.add_argument("--randomize-query", action="store_true")
    parser.add_argument(
        "--max-query-length",
        type=int,
        default=8,
        help="Max length for randomized search queries (default: 8)",
    )
    parser.add_argument(
        "--min-query-length",
        type=int,
        default=3,
        help="Min length for randomized search queries (default: 3)",
    )
    parser.add_argument(
        "--randomize-words-file",
        type=str,
        default=None,
        help="Word list for randomized queries (default: scripts/words.txt beside this script)",
    )
    parser.add_argument(
        "--randomize-seed-csv",
        type=str,
        default=None,
        metavar="PATH",
        help="CSV with name/naam column for random seeds (default: scripts/nl_companies.csv if it exists)",
    )
    parser.add_argument(
        "--randomize-no-seed-csv",
        action="store_true",
        help="Do not load any seed CSV for query suggestions (words + last results only)",
    )
    parser.add_argument("--telemetry", action="store_true")
    parser.add_argument(
        "--append-csv",
        action="store_true",
        help="Append to output CSV instead of overwriting on single run",
    )
    parser.add_argument(
        "--vat-id-search-engine",
        type=str,
        choices=("none", "bizzy"),
        default="none",
        help="NL + KVK only: enrich missing vat_id from https://bizzy.ai/nl/nl/{kvk} (default: none)",
    )

    args = parser.parse_args()
    TELEMETRY_ENABLED = args.telemetry

    cc = args.country_code.strip().upper()
    if len(cc) != 2 or not cc.isalpha():
        ansi_error("--country-code must be two letters (ISO 3166-1 alpha-2).")
        sys.exit(2)

    if args.search_engine == "kvk" and cc != "NL":
        ansi_error(
            "--search-engine kvk requires --country-code NL (KVK is Netherlands-only)."
        )
        sys.exit(2)

    if args.vat_id_search_engine == "bizzy":
        if cc != "NL":
            ansi_error("--vat-id-search-engine bizzy requires --country-code NL.")
            sys.exit(2)
        if args.search_engine != "kvk":
            ansi_error("--vat-id-search-engine bizzy requires --search-engine kvk.")
            sys.exit(2)

    if not args.randomize_query and not args.query.strip():
        ansi_error(
            "Provide --query or use --randomize-query with a non-empty seed via --query."
        )
        sys.exit(2)

    words_path = args.randomize_words_file or os.path.join(_scripts_dir, "words.txt")
    if args.randomize_no_seed_csv:
        seed_csv_path = None
    elif args.randomize_seed_csv:
        seed_csv_path = args.randomize_seed_csv
    else:
        default_nl = os.path.join(_scripts_dir, "nl_companies.csv")
        seed_csv_path = default_nl if os.path.isfile(default_nl) else None

    northdata_cookie = os.getenv("NORTHDATA_COOKIE")
    if args.northdata_cookie_file:
        try:
            with open(args.northdata_cookie_file, encoding="utf-8") as cf:
                northdata_cookie = cf.read().strip()
        except OSError as e:
            ansi_error(f"Cannot read --northdata-cookie-file: {e}")
            sys.exit(2)

    insert = not args.no_insert
    hash_cache = load_hash_cache_from_csv(args.output)
    ee_registrikood_cache: set[str] | None = set() if cc == "EE" else None

    _rate_reset()
    connector = aiohttp.TCPConnector(limit=100, limit_per_host=100)
    async with aiohttp.ClientSession(connector=connector) as session:
        rate_task = asyncio.create_task(_rate_logger_loop())
        try:
            if args.randomize_query:
                await randomize_query_loop(
                    session,
                    args.search_engine,
                    cc,
                    args.query.strip() or "fi",
                    args.output,
                    args.athena_url,
                    args.athena_key,
                    args.athena_client,
                    insert,
                    args.max_workers,
                    args.min_query_length,
                    args.max_query_length,
                    args.kvk_api_url,
                    args.northdata_base_url,
                    northdata_cookie,
                    words_path,
                    seed_csv_path,
                    athena_user_id=args.athena_user_id,
                    athena_company_id=args.athena_company_id,
                    athena_organization_id=args.athena_organization_id,
                    vat_id_search_engine=args.vat_id_search_engine,
                )
            else:
                if args.search_engine == "northdata":
                    raw = await fetch_northdata_suggest(
                        session,
                        args.northdata_base_url,
                        args.query.strip(),
                        cc,
                        northdata_cookie,
                    )
                    rows = []
                    for hit in raw:
                        g = map_northdata_hit_to_global(hit, cc)
                        if g:
                            rows.append(g)
                else:
                    # fetch_all_kvk_companies already returns parsed nl_companies-shaped dicts
                    companies = await fetch_all_kvk_companies(
                        session, args.kvk_api_url, args.query.strip()
                    )
                    if args.vat_id_search_engine == "bizzy":
                        await enrich_nl_companies_vat_from_bizzy(session, companies)
                    rows = []
                    for company in companies:
                        g = map_kvk_company_to_global(company, cc)
                        if g:
                            rows.append(g)

                _rate_add_scraped(len(rows))
                ansi_info(f"Mapped {len(rows)} row(s).")

                await process_global_rows(
                    session,
                    rows,
                    args.athena_url,
                    args.athena_key,
                    args.athena_client,
                    insert,
                    args.max_workers,
                    hash_cache,
                    cc,
                    ee_registrikood_cache,
                    athena_user_id=args.athena_user_id,
                    athena_company_id=args.athena_company_id,
                    athena_organization_id=args.athena_organization_id,
                )

                if rows:
                    save_global_rows_csv(rows, args.output, append=args.append_csv)
                    for r in rows:
                        hk = r.get("hash")
                        if hk:
                            hash_cache.add(hk)
                        if ee_registrikood_cache is not None:
                            ee = map_global_row_to_ee_company(r)
                            if ee and ee.get("registrikood"):
                                ee_registrikood_cache.add(ee["registrikood"])
                else:
                    ansi_warning("No rows produced.")
        finally:
            rate_task.cancel()
            try:
                await rate_task
            except asyncio.CancelledError:
                pass


def main() -> None:
    try:
        asyncio.run(main_async())
    except KeyboardInterrupt:
        ansi_warning("Interrupted.")


if __name__ == "__main__":
    main()
