scripts/backfill_pivot_months.py

#!/usr/bin/env python3
"""Backfill one-shot: meses fechados -> pivot_transporte_{MM} (+ demais bases pivot).

Exemplos:
  ./scripts/backfill_pivot_months.py --year 2026 --months 1-5
  ./scripts/backfill_pivot_months.py --year 2026 --auto   # jan..mes_anterior
  ./scripts/backfill_pivot_months.py --upload-dataset     # apos backfill, push pro HF dataset
"""
from __future__ import annotations

import argparse
import logging
import os
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from dashboard_apn.backfill import backfill_pivot_months, missing_pivot_months  # noqa: E402

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("backfill_pivot_months")


def load_env_local(path: Path) -> None:
    if not path.exists():
        return
    for line in path.read_text().splitlines():
        line = line.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        k, v = line.split("=", 1)
        os.environ.setdefault(k.strip(), v.strip().strip('"').strip("'"))


def parse_months(spec: str) -> list[int]:
    out: list[int] = []
    for part in spec.split(","):
        part = part.strip()
        if not part:
            continue
        if "-" in part:
            a, b = part.split("-", 1)
            out.extend(range(int(a), int(b) + 1))
        else:
            out.append(int(part))
    return sorted(set(m for m in out if 1 <= m <= 12))


def upload_dataset(db_path: Path) -> None:
    token = os.environ.get("HF_TOKEN", "")
    repo = os.environ.get("DATASET_REPO", "gabrielsapucaia02/cco-apoena-data")
    if not token:
        log.error("HF_TOKEN ausente — pulando upload do dataset")
        return
    from huggingface_hub import HfApi

    api = HfApi(token=token)
    api.upload_file(
        path_or_fileobj=str(db_path),
        path_in_repo="f2m_local.duckdb",
        repo_id=repo,
        repo_type="dataset",
        commit_message="backfill pivot months",
    )
    log.info("Dataset upload OK: %s", repo)


def main() -> int:
    load_env_local(ROOT / ".env.local")

    p = argparse.ArgumentParser(description="Backfill pivot_* meses fechados F2M")
    p.add_argument("--db", default=os.environ.get("F2M_DB_PATH", str(ROOT / "f2m_local.duckdb")))
    p.add_argument("--year", type=int, default=2026)
    p.add_argument("--months", help="Ex: 1-5 ou 1,2,3")
    p.add_argument("--auto", action="store_true", help="Todos os meses fechados do ano")
    p.add_argument("--force", action="store_true", help="Re-baixa mesmo se pivot ja existir")
    p.add_argument("--upload-dataset", action="store_true", help="Push f2m_local.duckdb pro HF dataset")
    args = p.parse_args()

    db_path = Path(args.db)
    if not db_path.exists():
        log.error("DB nao encontrado: %s", db_path)
        return 1

    email = os.environ.get("F2M_EMAIL", "")
    password = os.environ.get("F2M_PASSWORD", "")
    if not email or not password:
        log.error("F2M_EMAIL/F2M_PASSWORD ausentes (.env.local ou env)")
        return 1

    if args.auto:
        months = missing_pivot_months(str(db_path), args.year)
    elif args.months:
        months = parse_months(args.months)
    else:
        months = list(range(1, 6))  # default jan-maio

    log.info("Backfill %d meses em %s: %s", len(months), db_path.name, months)
    missing = missing_pivot_months(str(db_path), args.year)
    log.info("Pivots ausentes antes: %s", missing)

    result = backfill_pivot_months(
        str(db_path),
        email,
        password,
        args.year,
        months,
        skip_existing=not args.force,
    )

    for mes, tbl, n in result["written"]:
        log.info("  OK mes=%02d %s (%d linhas transporte)", mes, tbl, n)
    for mes in result["skipped"]:
        log.info("  SKIP mes=%02d (ja existia)", mes)
    for mes, err in result["failed"]:
        log.error("  FAIL mes=%02d: %s", mes, err)

    log.info(
        "Concluido: %d gravados, %d pulados, %d falhas | ausentes depois: %s",
        len(result["written"]),
        len(result["skipped"]),
        len(result["failed"]),
        result["missing_after"],
    )

    if args.upload_dataset and not result["failed"]:
        upload_dataset(db_path)

    return 1 if result["failed"] else 0


if __name__ == "__main__":
    raise SystemExit(main())