import argparse
import csv
import time
from typing import Literal

import daft


def string_data() -> str:
   return "hello world, this is a test, with commas, and \"quotes\", and unicode: 測試"


def write_csv(name: str, data_type: Literal["string", "number"], num_rows: int = 1**6, num_cols: int = 100):
  with open(name, 'w', newline='') as csvfile:
      writer = csv.writer(csvfile)
      writer.writerow([f"col_{i}" for i in range(num_cols)])
      match data_type:
          case "string":
              data = string_data()
          case "number":
              data = 1.2345
      writer.writerows([data for _ in range(num_cols)] for _ in range(num_rows))


def process_csv(name: str):
  df = daft.read_csv(name)
  batches = df.to_arrow_iter()
  start = time.perf_counter()
  next(batches)
  end = time.perf_counter()
  print(f"Read first batch in {end - start:.2f} seconds")
  count = len(list(batches))
  print(f"Read remaining {count} batches in {time.perf_counter() - end:.2f} seconds")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", help="Available commands", required=True)

    write_parser = subparsers.add_parser("write", help="Write a large CSV file")
    write_parser.add_argument("--rows", type=int, default=1**6, help="Number of rows to generate")
    write_parser.add_argument("--cols", type=int, default=100, help="Number of columns to generate")
    write_parser.add_argument("--type", type=str, choices=["string", "number"], default="string", help="Type of data to generate")
    write_parser.add_argument("filename", type=str, help="The name of the file to write")

    process_parser = subparsers.add_parser("process", help="Process a CSV file")
    process_parser.add_argument("filename", type=str, help="The name of the file to process")

    args = parser.parse_args()

    if args.command == "write":
        write_csv(args.filename, data_type=args.type, num_rows=args.rows, num_cols=args.cols)
    elif args.command == "process":
        process_csv(args.filename)

DAFT_PROGRESS_BAR=false uv run daft_read_csv.py write --rows 1000000 --cols 100 --type number numbers.data.csv

DAFT_PROGRESS_BAR=false uv run daft_read_csv.py write --rows 1000000 --cols 10 --type string strings.data.csv

DAFT_PROGRESS_BAR=false uv run daft_read_csv.py process numbers.data.csv
Read first batch in 0.04 seconds
Read remaining 4679 batches in 1.44 seconds

DAFT_PROGRESS_BAR=false uv run daft_read_csv.py process strings.data.csv
Read first batch in 1.20 seconds
Read remaining 5780 batches in 0.24 seconds

With very large files this causes the process to use a lot of memory before yielding the first chunk