#!/usr/bin/env python3
import json
import sys
import time
from urllib import request, error


def build_chunk(file_url: str, chunk_index: int, chunk_total: int, total_lines: int):
    return {
        "fileUrl": file_url,
        "originalFilename": "demo.xlsx",
        "relativePath": None,
        "mainSheetName": "Sheet1",
        "chunkIndex": chunk_index,
        "chunkTotal": chunk_total,
        "totalLines": total_lines,
        "keptRows": [f"brand-{chunk_index}"],
        "invalidBrands": [],
        "queryFailedBrands": []
    }


def post_json(url: str, payload: dict):
    data = json.dumps(payload).encode("utf-8")
    req = request.Request(url, data=data, headers={"Content-Type": "application/json"}, method="POST")
    with request.urlopen(req, timeout=60) as resp:
        return resp.status, resp.read().decode("utf-8", errors="ignore")


def main():
    if len(sys.argv) < 4:
        print("usage: python brand_chunk_load_test.py <base_url> <task_id> <chunk_total> [file_url]")
        sys.exit(1)
    base_url = sys.argv[1].rstrip("/")
    task_id = int(sys.argv[2])
    chunk_total = int(sys.argv[3])
    file_url = sys.argv[4] if len(sys.argv) > 4 else "https://example.com/demo.xlsx"

    submit_url = f"{base_url}/api/brand/tasks/{task_id}/result"
    started = time.time()
    for i in range(1, chunk_total + 1):
        payload = {
            "strategy": "Terms",
            "files": [build_chunk(file_url, i, chunk_total, chunk_total)]
        }
        status, body = post_json(submit_url, payload)
        if i == 1 or i == chunk_total or i % 100 == 0:
            print(f"chunk {i}/{chunk_total} status={status} elapsed={time.time() - started:.2f}s")
        if status < 200 or status >= 300:
            print(body)
            sys.exit(2)
    print(f"done {chunk_total} chunks in {time.time() - started:.2f}s")


if __name__ == "__main__":
    main()