import argparse
from pathlib import Path

import rich

from openai import OpenAI

# From the repository root, generate fixtures and clean up when finished:
#   fixture_dir=$(bash examples/generate_file.sh 1048576)
#   uv run --locked --all-extras python examples/uploads.py "$fixture_dir/big_test_file.txt"
#   rm -r -- "$fixture_dir"
# Omit the byte count for 1 GiB. Add "memory" to demonstrate an in-memory upload.


def from_disk(client: OpenAI, file: Path) -> None:
    print("uploading file from disk")

    upload = client.uploads.upload_file_chunked(
        file=file,
        mime_type="txt",
        purpose="batch",
    )
    rich.print(upload)


def from_in_memory(client: OpenAI, file: Path) -> None:
    print("uploading file from memory")

    # read the data into memory ourselves to simulate
    # it coming from somewhere else
    data = file.read_bytes()
    filename = "my_file.txt"

    upload = client.uploads.upload_file_chunked(
        file=data,
        filename=filename,
        bytes=len(data),
        mime_type="txt",
        purpose="batch",
    )
    rich.print(upload)


def main() -> None:
    parser = argparse.ArgumentParser(description="Upload a file generated by examples/generate_file.sh")
    parser.add_argument("file", type=Path)
    parser.add_argument("mode", nargs="?", choices=("disk", "memory"), default="disk")
    args = parser.parse_args()

    with OpenAI() as client:
        if args.mode == "memory":
            from_in_memory(client, args.file)
        else:
            from_disk(client, args.file)


if __name__ == "__main__":
    main()
