#!/usr/bin/env python3
"""
CLI script to fetch and parse company data from the Dutch KVK API,
saving each company to a CSV file and inserting into the nl_companies table and global_companies_registry via the Athena Gateway.

Dependencies:
    - aiohttp
    - colorama

Install dependencies via:
    pip install aiohttp colorama
"""



import asyncio
import csv
import argparse
from colorama import Fore, Style, init as colorama_init
import inspect
import json
import hashlib
import random
import os
import subprocess
import sys
import time
from functools import wraps



colorama_init(autoreset=True)

# Global flag for telemetry
TELEMETRY_ENABLED = False

# Rate logging: scrapes and inserts per minute, every 15 sec
_rate_scraped = 0
_rate_inserted = 0
_rate_last_time = None
_rate_last_scraped = 0
_rate_last_inserted = 0
_rate_task = None


def _rate_add_scraped(n: int):
    global _rate_scraped
    _rate_scraped += n


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


def _rate_reset():
    """Reset rate counters at start of a run."""
    global _rate_scraped, _rate_inserted
    _rate_scraped = 0
    _rate_inserted = 0


async def _rate_logger_loop():
    """Log scrapes/min and inserts/min every 15 seconds."""
    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):
    print(
        f"{Fore.MAGENTA}[DEXTER]{Style.RESET_ALL} {Fore.CYAN}[INFO]{Style.RESET_ALL} {msg}"
    )


def ansi_success(msg):
    print(
        f"{Fore.MAGENTA}[DEXTER]{Style.RESET_ALL} {Fore.GREEN}[SUCCESS]{Style.RESET_ALL} {msg}"
    )


def ansi_warning(msg):
    print(
        f"{Fore.MAGENTA}[DEXTER]{Style.RESET_ALL} {Fore.YELLOW}[WARNING]{Style.RESET_ALL} {msg}"
    )


def ansi_error(msg):
    print(
        f"{Fore.MAGENTA}[DEXTER]{Style.RESET_ALL} {Fore.RED}[ERROR]{Style.RESET_ALL} {msg}"
    )


def telemetry(func):
    """Decorator to measure and print execution time in milliseconds."""
    if inspect.iscoroutinefunction(func):
        @wraps(func)
        async def async_wrapper(*args, **kwargs):
            if TELEMETRY_ENABLED:
                start = time.perf_counter()
                result = await func(*args, **kwargs)
                end = time.perf_counter()
                elapsed_ms = (end - start) * 1000
                ansi_info(f"{func.__name__} took {elapsed_ms:.2f}ms")
                return result
            else:
                return await func(*args, **kwargs)
        return async_wrapper
    else:
        @wraps(func)
        def sync_wrapper(*args, **kwargs):
            if TELEMETRY_ENABLED:
                start = time.perf_counter()
                result = func(*args, **kwargs)
                end = time.perf_counter()
                elapsed_ms = (end - start) * 1000
                ansi_info(f"{func.__name__} took {elapsed_ms:.2f}ms")
                return result
            else:
                return func(*args, **kwargs)
        return sync_wrapper


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

    
@telemetry
def parse_nl_company_item(item):
    """
    Parse a single company item from the KVK API response and return a dict for nl_companies.
    """
    bezoek = item.get("bezoeklocatie", {}) or {}
    post = item.get("postlocatie", {}) or {}

    # Address fields
    street = bezoek.get("straat")
    house_number = bezoek.get("huisnummer")
    house_number = (
        int(house_number) if house_number and str(house_number).isdigit() else None
    )
    postal_code = bezoek.get("postcode")
    city = bezoek.get("plaats")
    area = bezoek.get("wijk") or None
    neighborhood = bezoek.get("buurt") or None
    province = bezoek.get("provincie") or None
    municipality = bezoek.get("gemeente") or None
    country = bezoek.get("land") or "Nederland"

    # Geolocation
    longitude = bezoek.get("longitude")
    latitude = bezoek.get("latitude")
    longitude = float(longitude) if longitude not in (None, "") else None
    latitude = float(latitude) if latitude not in (None, "") else None

    # Company identifiers
    kvk_number = item.get("kvkNummer")
    kvk_number = int(kvk_number) if kvk_number and str(kvk_number).isdigit() else None
    establishment_number = item.get("vestigingsnummer")
    establishment_number = (
        int(establishment_number)
        if establishment_number and str(establishment_number).isdigit()
        else None
    )
    vat_id = item.get("btwNummer") or None
    rsin = item.get("rsin") or None

    # Company status
    active = item.get("actief")
    if active is not None:
        active = bool(active)
    else:
        active = None

    # Rights form
    rights_form = item.get("rechtsvormOmschrijving") or None
    rights_form_code = item.get("rechtsvorm") or None

    # Activities
    activities = item.get("activiteiten", [])
    main_activity = None
    main_activity_description = None
    main_sbi_code = None
    standard_business_layout = None
    if activities:
        # Try to find hoofdactiviteit (main activity)
        main = None
        for act in activities:
            if act.get("hoofdactiviteit") is True:
                main = act
                break
        if not main and activities:
            main = activities[0]
        if main:
            main_activity = main.get("omschrijving") or None
            main_activity_description = main.get("omschrijving") or None
            main_sbi_code = main.get("code")
            main_sbi_code = (
                int(main_sbi_code)
                if main_sbi_code and str(main_sbi_code).isdigit()
                else None
            )
            standard_business_layout = main.get("sbiCode") or None
    activities_json = activities if activities else None

    # Establishment
    establishment = item.get("inschrijvingstype")
    establishment = (
        True
        if establishment and establishment.lower() == "hoofdvestiging"
        else False if establishment else None
    )

    # Compose company dict for nl_companies
    company = {
        # "id" is intentionally omitted to avoid violating NOT NULL constraint downstream
        "kvk_number": kvk_number,
        "name": item.get("naam"),
        "postal_code": postal_code,
        "active": active,
        "street": street,
        "area": area,
        "longitude": longitude,
        "latitude": latitude,
        "vat_id": vat_id,
        "rsin": rsin,
        "establishment_number": establishment_number,
        "rights_form": rights_form,
        "province": province,
        "municipality": municipality,
        "standard_business_layout": standard_business_layout,
        "rights_form_code": rights_form_code,
        "establishment": establishment,
        "house_number": house_number,
        "neighborhood": neighborhood,
        "country": country,
        "activities": json.dumps(activities_json) if activities_json else None,
        "city": city,
        "main_activity": main_activity,
        "main_activity_description": main_activity_description,
        "main_sbi_code": main_sbi_code,
    }
    # Remove any accidental "id" key if present
    if "id" in company:
        del company["id"]
    return company


@telemetry
def map_to_global_company(company):
    """
    Map a nl_companies dict to the global_companies_registry schema.
    """
    # Compose address string
    address_parts = []
    street = company.get("street") or None
    house_number = company.get("house_number") or None
    postal_code = company.get("postal_code") or None
    city = company.get("city") or None
    province = company.get("province") or None
    country = company.get("country") or "NL"

    if company.get("street"):
        address_parts.append(str(company.get("street")))
    if company.get("house_number"):
        address_parts.append(str(company.get("house_number")))
    if company.get("postal_code"):
        address_parts.append(str(company.get("postal_code")))
    if company.get("city"):
        address_parts.append(str(company.get("city")))
    if company.get("province"):
        address_parts.append(str(company.get("province")))
    if company.get("country"):
        address_parts.append(str(company.get("country")))
    address = ", ".join(address_parts) if address_parts else None

    # Compose hash (unique on name + kvk_number + country)
    hash_input = f"{company.get('name','')}|{company.get('kvk_number','')}|{company.get('country','')}"
    hash_val = hashlib.sha256(hash_input.encode("utf-8")).hexdigest()

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

    # Compose region/state
    region = company.get("province") or None
    state = company.get("province") or None

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

    # Compose start_date/establishment_date
    # KVK API does not provide explicit start date in this code, so leave as None

    # Compose language
    language = "nl"

    # Compose company_number
    company_number = (
        str(company.get("kvk_number"))
        if company.get("kvk_number") is not None
        else None
    )

    # Compose description
    description = (
        company.get("main_activity_description") or company.get("main_activity") or None
    )

    # Compose legal_persons (not available)
    legal_persons = None

    # Compose contact_data (not available)
    contact_data = None

    # Compose name_english (not available)
    name_english = None

    # Compose last_update (not available)
    last_update = None

    # Compose entity_type (not available)
    entity_type = None

    # Unix time in seconds (time of scrape/map)
    unix_timestamp = int(time.time())

    # Compose global_company_id (let DB default)

    global_company = {
        # "id" and "created_at" are omitted (auto-generated)
        "name": company.get("name"),
        "company_number": company_number,
        "rights_form": company.get("rights_form"),
        "country_code": "NL",
        "status": status,
        "last_update": last_update,
        "description": description,
        "region": region,
        "rights_form_code": company.get("rights_form_code"),
        "state": state,
        "name_english": name_english,
        "language": language,
        "hash": hash_val,
        "address": address,
        # "global_company_id": omitted (let DB default)
        "start_date": None,
        "legal_persons": legal_persons,
        "contact_data": contact_data,
        "activities": activities,
        "activity_code_table": activity_code_table,
        "establishment_date": None,
        "establishment_type": establishment_type,
        "entity_type": entity_type,
        "postal_code": postal_code,
        "house_number": house_number,
        "street": street,
        "province": province,
        "city": city,
        "country": country,
        "time": unix_timestamp,
    }
    # Remove accidental "id" or "created_at" if present
    for k in ["id", "created_at"]:
        if k in global_company:
            del global_company[k]
    return global_company


@telemetry
def save_to_csv(companies, filename="nl_companies.csv"):
    if not companies:
        ansi_warning("No companies to save.")
        return
    # Use all keys from the first company as fieldnames
    fieldnames = list(companies[0].keys())
    # Remove "id" from fieldnames if present
    if "id" in fieldnames:
        fieldnames.remove("id")
    # Open the file in append mode
    with open(filename, mode="a", newline="", encoding="utf-8") as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
        # Check if the file is empty to write the header
        csvfile.seek(0, 2)  # Move to the end of the file
        if csvfile.tell() == 0:
            writer.writeheader()
        for company in companies:
            # Remove "id" from each row if present
            row = dict(company)
            if "id" in row:
                del row["id"]
            writer.writerow(row)
    ansi_success(f"Saved {len(companies)} companies to '{filename}'.")


@telemetry
async def fetch_kvk_page(session, base_url, query, start=0, size=50):
    """
    Fetch a page of companies from the KVK API.
    """
    if start >= 5000:
        ansi_warning(f"Start value {start} is too high, skipping fetch.")
        return None

    url = f"{base_url}?q={query}&language=nl&site=kvk2014&size={size}&start={start}&inschrijvingsstatus=ingeschreven"
    headers = {"profileId": "5C10A89D-635E-49CC-94B8-042DD533B64A"}
    try:
        async with session.get(
            url, headers=headers, timeout=aiohttp.ClientTimeout(total=30)
        ) as resp:
            if resp.status == 200:
                return await resp.json()
            else:
                ansi_warning(f"Failed to fetch {url}: {resp.status}")
                print(await resp.text())
                return None
    except Exception as e:
        ansi_error(f"Exception fetching {url}: {e}")
        return None


@telemetry
async def fetch_all_kvk_companies(session, base_url, query):
    """
    Fetch all companies for the given query from the KVK API, handling pagination.
    Uses concurrent fetching to speed up pagination.
    """
    all_companies = []
    size = 50
    
    # Fetch first page to get total hits
    first_data = await fetch_kvk_page(session, base_url, query, start=0, size=size)
    if not first_data or "data" not in first_data or "items" not in first_data["data"]:
        return all_companies
    
    items = first_data["data"]["items"]
    for item in items:
        company = parse_nl_company_item(item)
        all_companies.append(company)
    
    total_hits = first_data["data"].get("numberOfHits", 0)
    remaining_pages = []
    start = size
    while start < total_hits and start < 5000:
        remaining_pages.append(start)
        start += size
    ansi_info(f"KVK query '{query}': {total_hits} hits, fetching {1 + len(remaining_pages)} pages")

    # Fetch all remaining pages concurrently
    if remaining_pages:
        fetch_tasks = [
            fetch_kvk_page(session, base_url, query, start=start, size=size)
            for start in remaining_pages
        ]
        results = await asyncio.gather(*fetch_tasks)
        
        # Process results
        for data in results:
            if data and "data" in data and "items" in data["data"]:
                items = data["data"]["items"]
                for item in items:
                    company = parse_nl_company_item(item)
                    all_companies.append(company)
    
    return all_companies


@telemetry
def load_companies_from_csv(filename):
    companies = []
    if not os.path.exists(filename):
        return companies
    with open(filename, newline="", encoding="utf-8") as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            companies.append(row)
    return companies


@telemetry
def get_existing_kvk_numbers_from_csv(filename):
    """
    Return a set of kvk_numbers (as int) from the CSV file.
    """
    kvk_numbers, _ = get_existing_caches_from_csv(filename)
    return kvk_numbers


@telemetry
def get_existing_company_numbers_from_csv(filename):
    """
    Return a set of company_number (as str) from the CSV file.
    """
    _, company_numbers = get_existing_caches_from_csv(filename)
    return company_numbers


@telemetry
def get_existing_caches_from_csv(filename):
    """
    One-pass read: return (kvk_numbers set, company_numbers set) from the CSV file.
    Use this when both caches are needed to avoid reading the file twice.
    """
    kvk_numbers = set()
    company_numbers = set()
    if not os.path.exists(filename):
        return kvk_numbers, company_numbers
    with open(filename, newline="", encoding="utf-8") as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            kvk = row.get("kvk_number")
            if kvk is not None and str(kvk).isdigit():
                kvk_numbers.add(int(kvk))
                company_numbers.add(str(kvk))
    return kvk_numbers, company_numbers


def build_athena_headers(
    athena_key,
    athena_client,
    user_id=None,
    company_id=None,
    organization_id=None,
    strip_nulls=None,
):
    headers = {
        "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


async def athena_insert(session, athena_url, headers, table_name, payload):
    endpoint = f"{athena_url.rstrip('/')}/gateway/insert"
    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)


def _athena_headers_or_build(athena_key, athena_client, user_id, company_id, organization_id, prebuilt=None):
    """Use prebuilt headers if provided, else build new ones."""
    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_nl_company_via_athena(
    session,
    company,
    athena_url,
    athena_key,
    athena_client,
    kvk_cache,
    user_id=None,
    company_id=None,
    organization_id=None,
    headers=None,
):
    """
    Insert a company dict into the nl_companies table via Athena Gateway.
    If the company already exists (duplicate kvk_number + establishment_number), ignore the error.
    Only insert if kvk_number is present and not in kvk_cache.
    """
    kvk_number = company.get("kvk_number")
    name = company.get("name")
    if not kvk_number or kvk_number == "" or not name or name == "":
        return

    if kvk_number in kvk_cache:
        return

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

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

    status, detail = await athena_insert(
        session, athena_url, h, "nl_companies", payload
    )
    if status == "ok":
        _rate_add_inserted(1)
        ansi_success(
            f"Inserted: ({company.get('kvk_number')}) {company.get('name')} [nl_companies]"
        )
        kvk_cache.add(kvk_number)
    elif status == "duplicate":
        kvk_cache.add(kvk_number)
    elif status == "error":
        ansi_error(
            f"Failed to insert: {company.get('name')} ({company.get('kvk_number')}): {detail} [nl_companies]"
        )
    else:
        ansi_error(
            f"Exception during insert: {company.get('name')} ({company.get('kvk_number')}): {detail} [nl_companies]"
        )


@telemetry
async def insert_global_company_via_athena(
    session,
    global_company,
    athena_url,
    athena_key,
    athena_client,
    company_number_cache,
    user_id=None,
    company_id=None,
    organization_id=None,
    headers=None,
):
    """
    Insert a company dict into the global_companies_registry table via Athena Gateway.
    If the company already exists (duplicate hash), ignore the error.
    Only insert if company_number is present and not in company_number_cache.
    """
    company_number = global_company.get("company_number")
    if not company_number:
        return

    if company_number in company_number_cache:
        return

    payload = {k: v for k, v in global_company.items() if v is not None}

    rights_form = global_company.get("rights_form")
    entity_type = None
    if rights_form == "Eenmanszaak":
        entity_type = "EMZ"
    elif rights_form == "Vennootschap Onder Firma":
        entity_type = "VOF"
    elif rights_form == "Maatschap":
        entity_type = "MS"
    elif rights_form == "Stichting":
        entity_type = "ST"
    elif rights_form == "Besloten Vennootschap":
        entity_type = "BV"
    payload["entity_type"] = entity_type

    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(
        session, athena_url, h, "global_companies_registry", payload
    )
    if status == "ok":
        _rate_add_inserted(1)
        ansi_success(
            f"Inserted: ({global_company.get('company_number')}) {global_company.get('name')}  [global_companies_registry]"
        )
        company_number_cache.add(company_number)
    elif status == "duplicate":
        company_number_cache.add(company_number)
    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 process_nl_companies(
    session,
    companies,
    athena_url,
    athena_key,
    athena_client,
    insert=True,
    max_workers=25,
    kvk_cache=None,
    company_number_cache=None,
    athena_user_id=None,
    athena_company_id=None,
    athena_organization_id=None,
):
    if kvk_cache is None:
        kvk_cache = set()
    if company_number_cache is None:
        company_number_cache = set()

    if insert and companies:
        from asyncio import Semaphore, create_task, gather

        # Build headers once for all inserts
        headers = build_athena_headers(
            athena_key, athena_client,
            user_id=athena_user_id,
            company_id=athena_company_id,
            organization_id=athena_organization_id,
        )
        semaphore = Semaphore(max_workers)

        async def worker(company):
            async with semaphore:
                try:
                    global_company = map_to_global_company(company)
                    # Run both inserts in parallel for this company
                    await gather(
                        insert_nl_company_via_athena(
                            session,
                            company,
                            athena_url,
                            athena_key,
                            athena_client,
                            kvk_cache,
                            user_id=athena_user_id,
                            company_id=athena_company_id,
                            organization_id=athena_organization_id,
                            headers=headers,
                        ),
                        insert_global_company_via_athena(
                            session,
                            global_company,
                            athena_url,
                            athena_key,
                            athena_client,
                            company_number_cache,
                            user_id=athena_user_id,
                            company_id=athena_company_id,
                            organization_id=athena_organization_id,
                            headers=headers,
                        ),
                    )
                except Exception as e:
                    ansi_error(f"Exception during insert: {e}")

        # Launch workers for each company, up to max_workers at a time
        tasks = [create_task(worker(company)) for company in companies]
        await gather(*tasks)
    elif insert:
        ansi_warning("No companies to insert.")
    return companies


@telemetry
async def randomize_query_loop(
    session,
    base_url,
    output_file,
    athena_url,
    athena_key,
    athena_client,
    insert,
    max_workers,
    max_query_length,
    csv_file="nl_companies.csv",
    athena_user_id=None,
    athena_company_id=None,
    athena_organization_id=None,
):
    """
    Infinite recursive random query loop:
    - Pick a random company from csv_file (or output_file)
    - Use the first X chars of its name as the next query
    - Fetch, process, save, repeat
    """
    seen_queries = set()
    # Build caches from CSV (one-pass read)
    kvk_cache, company_number_cache = get_existing_caches_from_csv(csv_file)
    while True:
        # Load companies from CSV
        companies_csv = load_companies_from_csv(csv_file)
        if not companies_csv:
            ansi_warning(
                f"No companies found in '{csv_file}'. Using default query 'fi'."
            )
            query = "fi"
        else:
            company = random.choice(companies_csv)
            name = company.get("name", "")
            if not name or len(name.strip()) == 0:
                query = "fi"
            else:
                query = name[:max_query_length].strip()
                if not query:
                    query = "fi"
        if query in seen_queries:
            ansi_info(f"Already tried query '{query}', picking another.")
            continue
        seen_queries.add(query)
        ansi_info(f"Randomized query: '{query}'")
        companies = await fetch_all_kvk_companies(session, base_url, query)
        _rate_add_scraped(len(companies))
        ansi_info(
            f"Fetched {len(companies)} companies from KVK API for query '{query}'."
        )
        await process_nl_companies(
            session,
            companies,
            athena_url,
            athena_key,
            athena_client,
            insert=insert,
            max_workers=max_workers,
            kvk_cache=kvk_cache,
            company_number_cache=company_number_cache,
            athena_user_id=athena_user_id,
            athena_company_id=athena_company_id,
            athena_organization_id=athena_organization_id,
        )
        if companies:
            save_to_csv(companies, filename=output_file)
            # Update caches with new companies
            for company in companies:
                kvk = company.get("kvk_number")
                if kvk is not None:
                    kvk_cache.add(int(kvk))
                company_number = (
                    str(company.get("kvk_number"))
                    if company.get("kvk_number") is not None
                    else None
                )
                if company_number:
                    company_number_cache.add(company_number)
        else:
            ansi_warning("No companies to save.")
        ansi_success(
            f"Completed fetching and saving companies to '{output_file}'. Total companies: {len(companies)}"
        )
        # Now, loop again, using the new CSV as the source for random company
        # (output_file is always written, so use that)
        csv_file = output_file


async def main_async():
    global TELEMETRY_ENABLED
    
    parser = argparse.ArgumentParser(
        description="Fetch and parse Dutch company data from the KVK API."
    )
    parser.add_argument(
        "--output", type=str, default="nl_companies.csv", help="Output CSV filename"
    )
    parser.add_argument(
        "--query",
        type=str,
        default="secra",
        help="Search query for KVK API (e.g. company name fragment)",
    )
    parser.add_argument(
        "--no-insert",
        action="store_true",
        help="Do not insert via Athena Gateway, just save CSV",
    )
    parser.add_argument(
        "--max-workers",
        type=int,
        default=25,
        help="Maximum number of concurrent insert workers (default: 25)",
    )
    parser.add_argument(
        "--kvk-api-url",
        type=str,
        default="https://web-api.kvk.nl/zoeken/v3/search",
        help="KVK API base URL",
    )
    parser.add_argument(
        "--athena-url",
        type=str,
        default="https://mirror3.athena-db.com",
        help="Athena Gateway base URL",
    )
    parser.add_argument(
        "--athena-key",
        type=str,
        default=os.getenv("ATHENA_API_KEY", ""),
        help="Athena Gateway API key (also sent as x-api-key); defaults to ATHENA_API_KEY env",
    )
    parser.add_argument(
        "--athena-client",
        type=str,
        default="railway_xbp",
        help="X-Athena-Client value (e.g. athena, xbp, railway_direct)",
    )
    parser.add_argument(
        "--athena-user-id",
        type=str,
        default=os.getenv("ATHENA_USER_ID", None),
        help="Optional X-User-Id header for Athena inserts (falls back to API key if server requires)",
    )
    parser.add_argument(
        "--athena-company-id",
        type=str,
        default=os.getenv("ATHENA_COMPANY_ID", None),
        help="Optional X-Company-Id header for Athena inserts",
    )
    parser.add_argument(
        "--athena-organization-id",
        type=str,
        default=os.getenv("ATHENA_ORGANIZATION_ID", None),
        help="Optional X-Organization-Id header for Athena inserts",
    )
    parser.add_argument(
        "--randomize-query",
        action="store_true",
        help="If set, recursively randomize the query using company names from CSV",
    )
    parser.add_argument(
        "--max-query-length",
        type=int,
        default=2,
        help="Maximum number of characters to use from company name for random query (default: 2)",
    )
    parser.add_argument(
        "--telemetry",
        action="store_true",
        help="Enable telemetry to print execution time for each function",
    )
    args = parser.parse_args()

    # Set global telemetry flag
    TELEMETRY_ENABLED = args.telemetry

    output_file = args.output
    query = args.query
    insert = not args.no_insert
    athena_url = args.athena_url
    athena_key = args.athena_key
    athena_client = args.athena_client
    athena_user_id = args.athena_user_id
    athena_company_id = args.athena_company_id
    athena_organization_id = args.athena_organization_id
    max_workers = args.max_workers
    kvk_api_url = args.kvk_api_url
    randomize_query = args.randomize_query
    max_query_length = args.max_query_length

    # Build caches from CSV (one-pass read)
    kvk_cache, company_number_cache = get_existing_caches_from_csv(output_file)

    _rate_reset()
    try:
        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 randomize_query:
                    await randomize_query_loop(
                        session,
                        kvk_api_url,
                        output_file,
                        athena_url,
                        athena_key,
                        athena_client,
                        insert,
                        max_workers,
                        max_query_length,
                        csv_file=output_file,
                        athena_user_id=athena_user_id,
                        athena_company_id=athena_company_id,
                        athena_organization_id=athena_organization_id,
                    )
                else:
                    companies = await fetch_all_kvk_companies(session, kvk_api_url, query)
                    _rate_add_scraped(len(companies))
                    ansi_info(f"Fetched {len(companies)} companies from KVK API.")

                    await process_nl_companies(
                        session,
                        companies,
                        athena_url,
                        athena_key,
                        athena_client,
                        insert=insert,
                        max_workers=max_workers,
                        kvk_cache=kvk_cache,
                        company_number_cache=company_number_cache,
                        athena_user_id=athena_user_id,
                        athena_company_id=athena_company_id,
                        athena_organization_id=athena_organization_id,
                    )

                    if companies:
                        save_to_csv(companies, filename=output_file)
                        # Update caches with new companies
                        for company in companies:
                            kvk = company.get("kvk_number")
                            if kvk is not None:
                                kvk_cache.add(int(kvk))
                            company_number = (
                                str(company.get("kvk_number"))
                                if company.get("kvk_number") is not None
                                else None
                            )
                            if company_number:
                                company_number_cache.add(company_number)
                    else:
                        ansi_warning("No companies to save.")
                    ansi_success(
                        f"Completed fetching and saving companies to '{output_file}'. Total companies: {len(companies)}"
                    )
            finally:
                rate_task.cancel()
                try:
                    await rate_task
                except asyncio.CancelledError:
                    pass
    except KeyboardInterrupt:
        ansi_warning("Interrupted by user (Ctrl+C). Exiting gracefully...")
        return


def main():
    try:
        asyncio.run(main_async())
    except KeyboardInterrupt:
        ansi_warning("Interrupted by user (Ctrl+C). Exiting...")


if __name__ == "__main__":
    main()
