from pathlib import Path
from polyfactory.factories.pydantic_factory import ModelFactory
from polyfactory import Use

from benchmark.pydantic_benchmark.models import (
    ColumnsInput,
    PlanogramInput,
    PlanogramsBulkInputPayload,
)


COLUMNS_COUNT = 1000
PLANOGRAMS_COUNT = 1000


class ColumnsInputFactory(ModelFactory):
    __model__ = ColumnsInput


class PlanogramInputFactory(ModelFactory):
    __model__ = PlanogramInput
    columns = Use(ColumnsInputFactory.batch, size=COLUMNS_COUNT)


class PlanogramsBulkInputPayloadFactory(ModelFactory):
    __model__ = PlanogramsBulkInputPayload

    planograms = Use(PlanogramInputFactory.batch, size=PLANOGRAMS_COUNT)


def create_test_file(filename: Path):
    bulk = PlanogramsBulkInputPayloadFactory.build()

    out_data = bulk.model_dump_json(by_alias=True)

    with open(filename, "w") as f:
        f.write(out_data)

    print(f"{len(out_data)} bytes of test data written to {filename}.")