from django.db.models import Count
from django.utils import timezone

from creative_module.models import AdCreativeData
from creative_module.types import UpdateStatus
from data_pipeline_module.services.ad_creative_sync import (
    build_ad_creative_row,
    push_ad_creatives_to_clickhouse,
)


ACCOUNT_ID = "0a61564f-6d44-4f04-946a-30575a6c2590"

# True pushes only rows not already marked as pushed.
# False re-pushes every AdCreativeData row for the account.
ONLY_CH_PUSHED_AT_NULL = True

# True matches the standard safe reconciliation path.
# False pushes every selected row regardless of sync_status.
ONLY_SUCCESS_ROWS = False

BATCH_SIZE = 500
DRY_RUN = False


qs = (
    AdCreativeData.objects.filter(account_id=ACCOUNT_ID)
    .select_related("account", "creative")
    .order_by("ad_id")
)

if ONLY_CH_PUSHED_AT_NULL:
    qs = qs.filter(ch_pushed_at__isnull=True)

if ONLY_SUCCESS_ROWS:
    qs = qs.filter(sync_status=UpdateStatus.SUCCESS.value)


print("Target rows by status:")
for row in qs.values("account__ad_account_name", "account_id", "sync_status").annotate(count=Count("id")).order_by("sync_status"):
    print(row)

total = qs.count()
print(
    {
        "account_id": ACCOUNT_ID,
        "selected": total,
        "only_ch_pushed_at_null": ONLY_CH_PUSHED_AT_NULL,
        "only_success_rows": ONLY_SUCCESS_ROWS,
        "dry_run": DRY_RUN,
    }
)

if DRY_RUN or total == 0:
    raise SystemExit

batch = []
batch_ids = []
total_built = 0
total_inserted = 0
total_marked = 0
total_skipped = 0
total_failed = 0


def flush_batch():
    global total_inserted, total_marked, total_failed
    if not batch:
        return

    inserted = push_ad_creatives_to_clickhouse(batch)
    total_inserted += inserted

    if inserted:
        marked = AdCreativeData.objects.filter(id__in=batch_ids).update(ch_pushed_at=timezone.now())
        total_marked += marked
        print(f"Pushed batch rows={len(batch)} inserted={inserted} marked={marked}")
    else:
        total_failed += len(batch)
        print(f"WARNING: ClickHouse push returned 0 for batch size={len(batch)}; not marking pushed")

    batch.clear()
    batch_ids.clear()


for ad in qs.iterator(chunk_size=BATCH_SIZE):
    row = build_ad_creative_row(ad)
    if row is None:
        total_skipped += 1
        continue

    batch.append(row)
    batch_ids.append(ad.id)
    total_built += 1

    if len(batch) >= BATCH_SIZE:
        flush_batch()

flush_batch()

print(
    {
        "selected": total,
        "rows_built": total_built,
        "inserted": total_inserted,
        "marked_ch_pushed_at": total_marked,
        "skipped_no_row": total_skipped,
        "failed_not_marked": total_failed,
    }
)

print("Remaining ch_pushed_at NULL by status:")
remaining = AdCreativeData.objects.filter(account_id=ACCOUNT_ID, ch_pushed_at__isnull=True)
for row in remaining.values("account__ad_account_name", "account_id", "sync_status").annotate(count=Count("id")).order_by("sync_status"):
    print(row)