#!/usr/bin/env python3
"""Enrich a CSV of companies through the TidyTools company enrichment API, one row at a time.

Docs: https://tools.yukai.uk/docs/company/enrich    Needs Python 3.8+, nothing else to install.

    export TIDYTOOLS_API_KEY=...        (PowerShell: $env:TIDYTOOLS_API_KEY = "...")
    python company-enrichment-csv.py --input companies.csv --output enriched.csv \\
        --name-column name --domain-column domain --country-column country --id-column id

A row with a domain sends only the domain; otherwise it sends the name and, when present, its two-letter country.
Every input row is written back unchanged and in order, followed by enrich_domain, enrich_name, enrich_primaryEmail,
enrich_primaryPhone, enrich_status and enrich_errorType.

enrich_status: success; error (the API or the network failed, reason in enrich_errorType, e.g. company_not_found);
skipped (the row has nothing valid to send, before or after a stop: empty_input or invalid_country); not_sent (the row
could be sent but the run had stopped; the reason is in enrich_errorType).

The run stops on unauthorized, insufficient_credits, account_suspended or Ctrl-C. The row that got such an answer is
error with that code; Ctrl-C during a call makes that row not_sent with interrupted. To finish a stopped run, send the
not_sent rows and the row with the stopping code again; skipped rows need their input fixed first.

Each sent row is one API call, charged as described in the docs and at https://tools.yukai.uk/pricing. Calls are not
retried, so nothing is sent twice. After a timeout or Ctrl-C during a call the outcome is unknown: check the balance
before sending that row again.

Exit code: 0 all rows handled, 1 stopped early, 2 nothing was done (bad arguments, file or key).
"""
import argparse
import csv
import json
import os
import re
import socket
import sys
import urllib.error
import urllib.request

API_URL = "https://tools.yukai.uk/api/v1/company/enrich"
KEY_ENV = "TIDYTOOLS_API_KEY"
USER_AGENT = "TidyTools-company-enrichment-csv/1.0 (+https://tools.yukai.uk/docs/company/enrich)"
TIMEOUT_SECONDS = 120
ENRICH_COLUMNS = ["enrich_domain", "enrich_name", "enrich_primaryEmail", "enrich_primaryPhone", "enrich_status", "enrich_errorType"]
# The next call would fail the same way: stop instead of walking through the rest of the file.
STOP_CODES = {"unauthorized", "insufficient_credits", "account_suspended"}
ERROR_CODE = re.compile(r"[A-Za-z0-9_.-]{1,64}")
COUNTRY = re.compile(r"[A-Za-z]{2}")
PHONE = re.compile(r"\+[\d\s().-]+")
FORMULA_START = re.compile(r"\s*[=+\-@]")


class _NoRedirect(urllib.request.HTTPRedirectHandler):
    """urllib would repeat the Authorization header to wherever a redirect points."""

    def redirect_request(self, *args, **kwargs):
        return None


_opener = urllib.request.build_opener(_NoRedirect)


def post_json(url, headers, body, timeout):
    """POST and return (status, headers with lowercase names, body bytes); any HTTP status is an answer, not an error."""
    request = urllib.request.Request(url, data=body, headers=headers, method="POST")
    try:
        with _opener.open(request, timeout=timeout) as response:
            return response.status, {k.lower(): v for k, v in response.headers.items()}, response.read()
    except urllib.error.HTTPError as e:
        with e:
            return e.code, {k.lower(): v for k, v in e.headers.items()}, e.read()


class Usage(Exception):
    """A problem to explain to the user before anything is sent."""


def read_csv(path):
    """Header and rows of a UTF-8 CSV, and whether it starts with a BOM (Excel's "CSV UTF-8" does)."""
    try:
        with open(path, "rb") as f:
            has_bom = f.read(3) == b"\xef\xbb\xbf"
        with open(path, encoding="utf-8-sig", newline="") as f:
            rows = list(csv.reader(f))
    except UnicodeDecodeError:
        raise Usage(f"{path} is not UTF-8: save it as \"CSV UTF-8\" and run again") from None
    except (OSError, csv.Error) as e:
        raise Usage(f"cannot read {path}: {e}") from None
    if not rows or not any(cell.strip() for cell in rows[0]):
        raise Usage(f"{path} has no header row")
    return rows[0], rows[1:], has_bom


def column_index(header, name, flag):
    if name is None:
        return None
    for i, cell in enumerate(header):
        if cell.strip() == name.strip():
            return i
    raise Usage(f"{flag}: no column named {name!r}; the file has: {', '.join(header)}")


def cell(row, index):
    return row[index].strip() if index is not None and index < len(row) else ""


def request_for(row, columns):
    """(JSON body, None) for a row to send, or (None, reason) for one to skip."""
    domain = cell(row, columns["domain"])
    if domain:
        return {"domain": domain}, None
    name = cell(row, columns["name"])
    if not name:
        return None, "empty_input"
    country = cell(row, columns["country"])
    if not country:
        return {"name": name}, None
    if not COUNTRY.fullmatch(country):
        return None, "invalid_country"
    return {"name": name, "country": country.upper()}, None


def safe_value(value):
    """An API value as a CSV cell. These come from third-party websites: a spreadsheet must not run them as a formula."""
    if isinstance(value, bool) or not isinstance(value, (str, int, float)):
        return ""
    text = str(value)
    if FORMULA_START.match(text) and not PHONE.fullmatch(text):
        return "'" + text
    return text


def call_api(payload, key, http):
    """(enrich columns, credits charged, balance). Never put the exception text in the result: it can quote the headers."""
    headers = {"Authorization": f"Bearer {key}", "Content-Type": "application/json", "User-Agent": USER_AGENT}
    body = json.dumps(payload).encode("utf-8")
    try:
        status, response_headers, raw = http(API_URL, headers, body, TIMEOUT_SECONDS)
    except Exception as e:
        reason = getattr(e, "reason", e)
        code = "timeout" if isinstance(reason, (socket.timeout, TimeoutError)) else "network_error"
        return result("error", code), 0.0, None
    charged = to_number(response_headers.get("x-credits-charged"))
    balance = response_headers.get("x-credits-balance")
    try:
        data = json.loads(raw.decode("utf-8"))
    except ValueError:
        data = None
    if not isinstance(data, dict):
        return result("error", "invalid_response" if 200 <= status < 300 else f"http_{status}"), charged, balance
    if 200 <= status < 300 and data.get("success") is not False:
        return result("success", "", data), charged, balance
    error = data.get("error")
    code = error.get("code") if isinstance(error, dict) else None
    if not (isinstance(code, str) and ERROR_CODE.fullmatch(code)):
        code = f"http_{status}"
    return result("error", code), charged, balance


def result(status, error_type, data=None):
    data = data or {}
    return [safe_value(data.get(f)) for f in ("domain", "name", "primaryEmail", "primaryPhone")] + [status, error_type]


def to_number(value):
    try:
        return float(value)
    except (TypeError, ValueError):
        return 0.0


def process(rows, width, columns, key, http, write, log):
    """Send and write the rows in order. Returns (counts by status, credits charged, last balance, the reason for stopping)."""
    counts = {"success": 0, "error": 0, "skipped": 0, "not_sent": 0}
    charged, balance, stopped = 0.0, None, None
    for number, row in enumerate(rows, 1):
        label = cell(row, columns["id"]) or f"row {number}"
        payload, skip = request_for(row, columns)
        if skip:
            out = result("skipped", skip)
        elif stopped:
            out = result("not_sent", stopped)
        else:
            try:
                out, used, last = call_api(payload, key, http)
                charged += used
                balance = last if last is not None else balance
                if out[-1] in STOP_CODES:
                    stopped = out[-1]
            except KeyboardInterrupt:
                stopped = "interrupted"
                out = result("not_sent", stopped)
        counts[out[-2]] += 1
        write(row + [""] * (width - len(row)) + out)
        log(f"[{number}/{len(rows)}] {label}: {out[-2]}{' ' + out[-1] if out[-1] else ''}")
    return counts, charged, balance, stopped


def main(argv=None, environ=os.environ, http=post_json, log=None):
    log = log or (lambda line: print(line, file=sys.stderr, flush=True))
    parser = argparse.ArgumentParser(description="Enrich a CSV of companies through the TidyTools company enrichment API, one row at a time.")
    parser.add_argument("--input", required=True, help="CSV to read (UTF-8, first row is the header)")
    parser.add_argument("--output", required=True, help="CSV to write: the input columns followed by the enrich_ columns")
    parser.add_argument("--domain-column", help="column with the company domain or website; preferred when a row has one")
    parser.add_argument("--name-column", help="column with the company name")
    parser.add_argument("--country-column", help="column with a two-letter country code, sent with the name")
    parser.add_argument("--id-column", help="column that identifies the row in the progress log")
    parser.add_argument("--overwrite", action="store_true", help="replace --output if it exists (it may hold results you paid for)")
    args = parser.parse_args(argv)
    try:
        if not args.domain_column and not args.name_column:
            raise Usage("give --name-column, --domain-column or both")
        key = environ.get(KEY_ENV, "").strip()
        if not key:
            raise Usage(f"set the environment variable {KEY_ENV} to your API key (from https://tools.yukai.uk/pricing)")
        if not re.fullmatch(r"[\x21-\x7e]+", key):
            raise Usage(f"the value of {KEY_ENV} has characters an API key cannot have")
        if os.path.exists(args.output):
            if os.path.exists(args.input) and os.path.samefile(args.input, args.output):
                raise Usage("--output must be a different file from --input")
            if not args.overwrite:
                raise Usage(f"{args.output} exists; choose another name or add --overwrite")
        header, rows, has_bom = read_csv(args.input)
        used = [c for c in ENRICH_COLUMNS if c in (h.strip() for h in header)]
        if used:
            raise Usage(f"{args.input} already has {', '.join(used)}: is it the output of an earlier run?")
        columns = {
            "domain": column_index(header, args.domain_column, "--domain-column"),
            "name": column_index(header, args.name_column, "--name-column"),
            "country": column_index(header, args.country_column, "--country-column"),
            "id": column_index(header, args.id_column, "--id-column"),
        }
        width = max([len(header)] + [len(row) for row in rows])
        try:
            out = open(args.output, "w", encoding="utf-8-sig" if has_bom else "utf-8", newline="")
        except OSError as e:
            raise Usage(f"cannot write {args.output}: {e}") from None
    except Usage as e:
        log(f"error: {e}")
        return 2
    with out:
        writer = csv.writer(out)
        writer.writerow(header + [""] * (width - len(header)) + ENRICH_COLUMNS)

        def write(row):
            writer.writerow(row)
            out.flush()  # a stopped run keeps every result already paid for

        counts, charged, balance, stopped = process(rows, width, columns, key, http, write, log)
    if stopped:
        log(f"stopped early: {stopped}")
    summary = ", ".join(f"{n} {status}" for status, n in counts.items())
    log(f"{summary}; {charged:g} credits charged{f', balance {balance}' if balance is not None else ''}; wrote {args.output}")
    return 1 if stopped else 0


if __name__ == "__main__":
    sys.exit(main())