from __future__ import annotations

import argparse
import json
import time
from typing import Optional, Sequence


def run_v2_backfill_batches(
    *,
    apply: bool,
    batch_size: int,
    start_offset: int = 0,
    end_offset: Optional[int] = None,
    reset_first: bool = False,
) -> dict:
    from sqlalchemy import text

    from db.session import get_engine
    from app.jobs.v2_backfill_foundation import run_v2_backfill

    engine = get_engine()
    with engine.connect() as conn:
        total_products = int(conn.execute(text("select count(*) from master_product")).scalar() or 0)

    effective_end = total_products if end_offset is None else min(int(end_offset), total_products)
    offset = max(0, int(start_offset))
    batch_size = max(1, int(batch_size))

    summary = {
        "status": "ok",
        "applied": bool(apply),
        "batch_size": batch_size,
        "start_offset": offset,
        "end_offset": effective_end,
        "total_products": total_products,
        "batches": [],
    }

    batch_index = 0
    while offset < effective_end:
        limit = min(batch_size, effective_end - offset)
        started = time.time()
        result = run_v2_backfill(
            apply=apply,
            sku_offset=offset,
            sku_limit=limit,
            reset=bool(reset_first and batch_index == 0),
        )
        elapsed = round(time.time() - started, 2)
        batch_summary = {
            "offset": offset,
            "limit": limit,
            "elapsed_seconds": elapsed,
            **result,
        }
        summary["batches"].append(batch_summary)
        print(json.dumps(batch_summary, sort_keys=True), flush=True)
        offset += limit
        batch_index += 1

    return summary


def main(argv: Optional[Sequence[str]] = None) -> int:
    parser = argparse.ArgumentParser(description="Run v2 foundation backfill in sequential batches")
    parser.add_argument("--apply", action="store_true", help="Commit writes to v2 tables. Default is dry-run.")
    parser.add_argument("--batch-size", type=int, default=500, help="Number of products per batch")
    parser.add_argument("--start-offset", type=int, default=0, help="Initial product offset")
    parser.add_argument("--end-offset", type=int, help="Optional exclusive end offset")
    parser.add_argument("--reset-first", action="store_true", help="Clear v2 foundation tables before the first batch")
    args = parser.parse_args(list(argv) if argv is not None else None)

    result = run_v2_backfill_batches(
        apply=args.apply,
        batch_size=args.batch_size,
        start_offset=args.start_offset,
        end_offset=args.end_offset,
        reset_first=args.reset_first,
    )
    print(json.dumps(result, indent=2, sort_keys=True))
    return 0


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