#!/usr/bin/env python3
"""
AWAL terminations pipeline — Step 1.

Reads 'AWAL Core - Terminations (revised).csv' and resolves each Account ID to
its vendor_uuid via art_relations.vendor. Writes two CSVs with identical content
but distinct filenames so the downstream bulk_runner output files
(`<stem>_responses.json`) for the two endpoints do not collide:

  data/input/awal_terminate_staff.csv    → used with configs/awal_terminate_staff.json
  data/input/awal_terminate_delete.csv   → used with configs/awal_terminate_delete.json

Usage:
    uv run python src/pipelines/awal_terminations/1_prepare.py \\
        --input "data/prepare/AWAL Core - Terminations (revised).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_HEADER = 'ABACUS Account Account 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_account_ids(csv_path: Path) -> list[int]:
    with open(csv_path, encoding='utf-8-sig', newline='') as f:
        reader = csv.DictReader(f)
        if reader.fieldnames is None or ACCOUNT_ID_HEADER not in reader.fieldnames:
            raise KeyError(f'Column {ACCOUNT_ID_HEADER!r} not found; header={reader.fieldnames}')
        ids: list[int] = []
        for row in reader:
            v = (row.get(ACCOUNT_ID_HEADER) or '').strip()
            if not v:
                continue
            ids.append(int(v))
    return ids


def lookup_uuids(account_ids: list[int]) -> dict[int, str]:
    conn = db_connect()
    try:
        cur = conn.cursor()
        result: dict[int, str] = {}
        for i in range(0, len(account_ids), BATCH_SIZE):
            batch = account_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, uuids: list[str]) -> None:
    with open(path, 'w', newline='', encoding='utf-8') as f:
        w = csv.writer(f)
        w.writerow(['vendor_uuid'])
        for u in uuids:
            w.writerow([u])


def main() -> None:
    parser = argparse.ArgumentParser(description='AWAL terminations prep: CSV → vendor_uuid CSVs.')
    parser.add_argument('--input', required=True, help='Path to AWAL terminations CSV')
    parser.add_argument(
        '--output-dir',
        default='data/input',
        help='Directory for output CSVs (default: data/input)',
    )
    args = parser.parse_args()

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

    account_ids = load_account_ids(src)
    print(f'Read {len(account_ids)} account IDs from {src.name}', file=sys.stderr)

    uuid_map = lookup_uuids(account_ids)
    resolved = [uuid_map[aid] for aid in account_ids if aid in uuid_map]
    missing = [aid for aid in account_ids if aid not in uuid_map]

    print(f'Resolved {len(resolved)}/{len(account_ids)} 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_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    staff_csv = out_dir / 'awal_terminate_staff.csv'
    delete_csv = out_dir / 'awal_terminate_delete.csv'
    write_csv(staff_csv, resolved)
    write_csv(delete_csv, resolved)
    print(f'Wrote {staff_csv} and {delete_csv} ({len(resolved)} rows each)', file=sys.stderr)


if __name__ == '__main__':
    main()