#!/usr/bin/env python3
"""
Country attribution pipeline — Step 1.

Reads a source CSV with columns ACCOUNT_ID, COUNTRY_ID, ATTRIBUTED_COUNTRY (and
optionally others), resolves each ACCOUNT_ID to its vendor_uuid via
art_relations.vendor, and writes a runner-ready CSV with columns
vendor_uuid, country_id for use with configs/country_attribution.json.

Usage:
    uv run python src/pipelines/country_attribution/1_prepare.py \\
        --input data/prepare/country_attribution_040926.csv

DB creds are read from .env (DB_URL). Requires `awsume prod`.
"""

import argparse
import csv
import os
import sys
from pathlib import Path
from urllib.parse import unquote

import pymysql  # type: ignore[import-untyped]
from dotenv import load_dotenv

ACCOUNT_ID_COL = 'ACCOUNT_ID'
COUNTRY_ID_COL = 'COUNTRY_ID'
BATCH_SIZE = 1000


def db_connect() -> pymysql.connections.Connection:
    load_dotenv()
    url = os.environ['DB_URL']
    rest = url.split('://', 1)[1]
    creds, hostdb = rest.split('@', 1)
    user, pw = creds.split(':', 1)
    host, db = hostdb.split('/', 1)
    return pymysql.connect(host=host, user=user, password=unquote(pw), database=db)


def load_source(csv_path: Path) -> list[tuple[int, int]]:
    """Return list of (account_id, country_id). Drops rows with missing values."""
    out: list[tuple[int, int]] = []
    with open(csv_path, newline='', encoding='utf-8-sig') as f:
        reader = csv.DictReader(f)
        if reader.fieldnames is None or ACCOUNT_ID_COL not in reader.fieldnames or COUNTRY_ID_COL not in reader.fieldnames:
            raise KeyError(f'Expected columns {ACCOUNT_ID_COL!r} and {COUNTRY_ID_COL!r} in {csv_path}; got {reader.fieldnames}')
        for row in reader:
            aid_raw = (row.get(ACCOUNT_ID_COL) or '').strip()
            cid_raw = (row.get(COUNTRY_ID_COL) or '').strip()
            if not aid_raw or not cid_raw:
                continue
            out.append((int(aid_raw), int(cid_raw)))
    return out


def lookup_uuids(account_ids: list[int]) -> dict[int, str]:
    conn = db_connect()
    try:
        cur = conn.cursor()
        result: dict[int, str] = {}
        unique_ids = list({aid for aid in account_ids})
        for i in range(0, len(unique_ids), BATCH_SIZE):
            batch = unique_ids[i : i + BATCH_SIZE]
            placeholders = ','.join(['%s'] * len(batch))
            cur.execute(
                f'SELECT vendor_id, vendor_uuid FROM vendor WHERE vendor_id IN ({placeholders})',
                batch,
            )
            for vid, vuuid in cur.fetchall():
                if vuuid:
                    result[vid] = vuuid
        cur.close()
        return result
    finally:
        conn.close()


def write_csv(path: Path, rows: list[tuple[str, int]]) -> None:
    with open(path, 'w', newline='', encoding='utf-8') as f:
        w = csv.writer(f)
        w.writerow(['vendor_uuid', 'country_id'])
        for vuuid, cid in rows:
            w.writerow([vuuid, cid])


def main() -> None:
    parser = argparse.ArgumentParser(description='Country attribution prep: source CSV → vendor_uuid/country_id CSV.')
    parser.add_argument('--input', required=True, help='Path to source CSV (ACCOUNT_ID, COUNTRY_ID, ATTRIBUTED_COUNTRY)')
    parser.add_argument(
        '--output',
        default='data/input/country_attribution.csv',
        help='Output CSV path (default: data/input/country_attribution.csv)',
    )
    args = parser.parse_args()

    src = Path(args.input)
    if not src.exists():
        print(f'Input not found: {src}', file=sys.stderr)
        sys.exit(1)

    pairs = load_source(src)
    print(f'Read {len(pairs)} (account_id, country_id) rows from {src}', file=sys.stderr)

    uuid_map = lookup_uuids([aid for aid, _ in pairs])
    resolved = [(uuid_map[aid], cid) for aid, cid in pairs if aid in uuid_map]
    missing = [aid for aid, _ in pairs if aid not in uuid_map]

    print(f'Resolved {len(resolved)}/{len(pairs)} vendor UUIDs', file=sys.stderr)
    if missing:
        print(
            f'WARNING: {len(missing)} account IDs had no vendor_uuid in art_relations.vendor',
            file=sys.stderr,
        )
        print(f'  sample: {missing[:10]}', file=sys.stderr)

    out = Path(args.output)
    out.parent.mkdir(parents=True, exist_ok=True)
    write_csv(out, resolved)
    print(f'Wrote {out} ({len(resolved)} rows)', file=sys.stderr)


if __name__ == '__main__':
    main()