import logging
import multiprocessing as mp
import threading
from datetime import datetime
from pathlib import Path
from typing import Any

import numpy as np
import numpy.typing as npt
import typer

from . import globals

app = typer.Typer()
logger = logging.getLogger(__name__)


def write_to_file(path: Path, queue: "mp.Queue[Any]") -> None:
    import h5py

    with h5py.File(path, "w") as file:
        logger.info("Starting writer")
        while True:
            data = queue.get()
            if data is None:
                break
            logger.info("Writing data to file")
            file.create_dataset(datetime.now().isoformat(), data=data)
        logger.info("Writer stopped")


@app.command()
def capture(output_path: Path) -> None:
    """Capture CSI data to a file"""
    logger.info(f"Capturing data to {output_path}")

    queue: "mp.Queue[None | npt.NDArray[np.complex128]]" = mp.Queue()
    writer = threading.Thread(
        target=write_to_file,
        args=(
            output_path,
            queue,
        ),
    )
    writer.start()

    def callback(antenna_data: npt.NDArray[np.complex128]) -> None:
        logger.info(f"Got final CSI data with shape {antenna_data.shape}")
        queue.put(antenna_data)

    globals.csi_producer(csi_callback=callback)
    queue.put(None)