Autoparts-DB/scripts/assign_sat_keys.py

#!/usr/bin/env python3
"""Assign SAT ``clave_prod_serv`` / ``clave_unidad`` per product (CFDI 4.0).

Two modes:

  sale-items (default, writes): backfill the existing ``sale_items`` columns
      for a tenant.  Historically they are 100% NULL, so every CFDI fell back
      to the hard-coded default (25174800 / H87).  Consumers that read
      ``sale_items`` directly (``services/global_invoice.py``,
      ``blueprints/invoicing_bp.py``) keep the default themselves, so the only
      way to give them a real per-product key is to persist it here.

  inventory (report only): show the key distribution that *would* be assigned
      to ``inventory`` rows.  ``inventory`` has no clave columns today, so this
      mode never writes — see the module note in the summary for the proposed
      ``inventory.clave_prod_serv`` / ``inventory.clave_unidad varchar(10)``
      columns.

Resolution order per item: ``services.sat_keys`` keyword rules over the name,
then the canonical category (from ``product_categorizer`` or the tenant
``categories`` row), then optionally the ``ai_chat`` LLM fallback, then the SAT
default.  Every AI call degrades silently when the provider is down.

Usage (as root, from /home/Autopartes):
    python3 scripts/assign_sat_keys.py --tenant 45 --limit 200 --dry-run
    python3 scripts/assign_sat_keys.py --tenant 45 --limit 500 --use-llm
    python3 scripts/assign_sat_keys.py --tenant 45 --mode inventory --dry-run
    python3 scripts/assign_sat_keys.py --all --batch 500

Nothing is written unless ``--apply`` (mode sale-items); ``--dry-run`` is the
default to avoid surprises.
"""

import argparse
import json
import os
import sys
import time

sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..', 'pos'))

from tenant_db import get_tenant_conn, get_master_conn  # noqa: E402
from services import sat_keys  # noqa: E402
from services.product_categorizer import classify_name  # noqa: E402

DEFAULT_CHECKPOINT = os.path.join(
    os.path.dirname(os.path.abspath(__file__)), '..', 'backups',
    'sat_keys_checkpoint.json'
)


def list_tenants(only_id=None):
    conn = get_master_conn()
    cur = conn.cursor()
    try:
        if only_id:
            cur.execute(
                "SELECT id, name, db_name FROM tenants "
                "WHERE is_active = true AND id = %s", (only_id,))
        else:
            cur.execute(
                "SELECT id, name, db_name FROM tenants "
                "WHERE is_active = true ORDER BY id")
        return cur.fetchall()
    finally:
        cur.close()
        conn.close()


def _load_checkpoint(path):
    try:
        with open(path) as f:
            return json.load(f)
    except Exception:
        return {}


def _save_checkpoint(path, data):
    try:
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, 'w') as f:
            json.dump(data, f, indent=2)
    except Exception as e:
        print(f"  ! checkpoint no guardado: {e}")


def _resolve(prod_name, brand, category_name, part_number, use_llm):
    canonical = None
    try:
        canonical = classify_name(prod_name or part_number or '', brand)[0]
    except Exception:
        canonical = None
    return sat_keys.resolve_key(
        name=prod_name, brand=brand, category=category_name,
        canonical=canonical, part_number=part_number, use_llm=use_llm,
    )


def run_sale_items(tenant_id, name, conn, args, state):
    after_id = None if args.reset else state.get('after_id')
    t_key = str(tenant_id)
    totals = {'scanned': 0, 'updated': 0, 'by_source': {}}
    while True:
        cur = conn.cursor()
        where = ["(si.clave_prod_serv IS NULL OR si.clave_prod_serv = '')"]
        params = []
        if after_id is not None:
            where.append("si.id > %s")
            params.append(after_id)
        sql = f"""
            SELECT si.id, si.part_number, si.name, si.inventory_id,
                   i.name, i.brand, c.name
            FROM sale_items si
            LEFT JOIN inventory i ON i.id = si.inventory_id
            LEFT JOIN categories c ON c.id = i.category_id
            WHERE {' AND '.join(where)}
            ORDER BY si.id
            LIMIT %s
        """
        params.append(args.batch)
        cur.execute(sql, params)
        rows = cur.fetchall()
        cur.close()
        if not rows:
            break

        for (si_id, si_pn, si_name, inv_id, inv_name, inv_brand, cat_name) in rows:
            prod_name = inv_name or si_name
            clave, unidad, source = _resolve(
                prod_name, inv_brand, cat_name, si_pn, args.use_llm)
            totals['by_source'][source] = totals['by_source'].get(source, 0) + 1
            totals['scanned'] += 1
            after_id = si_id
            if args.apply:
                c2 = conn.cursor()
                try:
                    c2.execute(
                        "UPDATE sale_items SET clave_prod_serv = %s, "
                        "clave_unidad = %s WHERE id = %s",
                        (clave, unidad, si_id))
                    totals['updated'] += c2.rowcount
                    conn.commit()
                finally:
                    c2.close()
            if totals['scanned'] <= args.show:
                print(f"  {si_id:<8} {str(prod_name)[:40]:<40} -> {clave} / {unidad} ({source})")

        state[t_key] = {'after_id': after_id, 'updated': totals['updated'],
                        'updated_at': time.strftime('%Y-%m-%dT%H:%M:%S')}
        if args.apply and not args.dry_run:
            _save_checkpoint(args.checkpoint, state)
        if len(rows) < args.batch:
            break
        if args.limit and totals['scanned'] >= args.limit:
            break
    return totals


def run_inventory(tenant_id, conn, args):
    """Report-only: distribution of keys the rules would assign to inventory."""
    cur = conn.cursor()
    cur.execute("""
        SELECT i.id, i.part_number, i.name, i.brand, c.name
        FROM inventory i
        LEFT JOIN categories c ON c.id = i.category_id
        WHERE i.is_active = true
        ORDER BY i.id
        LIMIT %s
    """, (args.batch,))
    rows = cur.fetchall()
    cur.close()
    by_source, by_key, sample = {}, {}, []
    for inv_id, pn, nm, brand, cat in rows:
        clave, unidad, source = _resolve(nm, brand, cat, pn, args.use_llm)
        by_source[source] = by_source.get(source, 0) + 1
        by_key[clave] = by_key.get(clave, 0) + 1
        if len(sample) < args.show:
            sample.append((inv_id, nm, clave, unidad, source))
    print(f"  inventory scanned={len(rows)} by_source={by_source}")
    print(f"  keys: {by_key}")
    for inv_id, nm, clave, unidad, source in sample:
        print(f"  {inv_id:<8} {str(nm)[:40]:<40} -> {clave} / {unidad} ({source})")
    return {'scanned': len(rows), 'by_source': by_source, 'by_key': by_key}


def main():
    p = argparse.ArgumentParser(description='Assign SAT product/unit keys.')
    p.add_argument('--tenant', type=int)
    p.add_argument('--all', action='store_true')
    p.add_argument('--mode', choices=['sale-items', 'inventory'], default='sale-items')
    p.add_argument('--limit', type=int, default=0, help='max items per tenant (0=all)')
    p.add_argument('--batch', type=int, default=500)
    p.add_argument('--apply', action='store_true', help='persist (sale-items only)')
    p.add_argument('--dry-run', action='store_true', default=True)
    p.add_argument('--use-llm', action='store_true', help='enable ai_chat fallback')
    p.add_argument('--reset', action='store_true', help='ignore checkpoint')
    p.add_argument('--show', type=int, default=8)
    p.add_argument('--checkpoint', default=DEFAULT_CHECKPOINT)
    args = p.parse_args()

    # --apply is the only thing that turns off dry-run; keep it explicit.
    if args.apply:
        args.dry_run = False
    if not args.all and not args.tenant:
        p.error('--tenant ID o --all requerido')

    state = _load_checkpoint(args.checkpoint)
    tenants = list_tenants(None if args.all else args.tenant)
    if not tenants:
        print('sin tenants para procesar')
        return 1

    print(f"mode={args.mode} apply={not args.dry_run} use_llm={args.use_llm}")
    grand = {}
    for tid, name, db_name in tenants:
        print(f"\n=== tenant {tid} — {name} ({db_name}) ===")
        conn = get_tenant_conn(tid)
        try:
            if args.mode == 'inventory':
                run_inventory(tid, conn, args)
            else:
                totals = run_sale_items(tid, name, conn, args, state)
                print(f"  total: scanned={totals['scanned']} updated={totals['updated']} "
                      f"by_source={totals['by_source']}")
                for k, v in totals['by_source'].items():
                    grand[k] = grand.get(k, 0) + v
        finally:
            conn.close()
    if args.mode == 'sale-items':
        print(f"\nTOTAL by_source={grand}")
    return 0


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