from zenml import step, pipeline
import pandas as pd
import numpy as np
from datetime import datetime, timedelta


@step(enable_cache=False)
def save_df(num_rows: int = 10000) -> pd.DataFrame:
    """Reads the data from the source."""
    data = {
        "ID": range(1, num_rows + 1),
        "Value": np.random.rand(num_rows),  # Random float values
        # Create dates cycling through a year to avoid 'date value out of range' error
        "Date": [datetime.now() - timedelta(days=x % 365) for x in range(num_rows)],
    }
    return pd.DataFrame(data)


@step(enable_cache=False)
def read_df(df: pd.DataFrame) -> pd.DataFrame:
    """Reads the data from the source."""
    return df


@pipeline(enable_cache=False)
def my_pipeline():
    # df = save_df()  # 10k rows
    # df = save_df(num_rows=1000000)  # 1M rows
    df = save_df(num_rows=28000000)  # 28M rows
    reloaded_df = read_df(df=df)


my_pipeline()

def save_chunk(self, df, start, chunk_size, path, compression):
        """Save a chunk of the dataframe to a parquet file."""
        chunk = df.iloc[start : start + chunk_size]
        with fileio.open(path, mode="wb") as f:
            chunk.to_parquet(f, compression=compression)

    def save(self, df: Union[pd.DataFrame, pd.Series]) -> None:
        """Writes a pandas dataframe or series to the specified filename in parallel chunks."""
        if isinstance(df, pd.Series):
            df = df.to_frame(name="series")

        if self.pyarrow_exists:
            # Write the length of the dataframe to a file for later use
            io_utils.write_file_contents_as_string(
                self.metadata_path, str(len(df))
            )

            # Define the maximum number of threads
            # Fallback to 10 if cpu_count is None
            max_workers = os.cpu_count() or 10

            # Prepare the thread pool for parallel writes
            with concurrent.futures.ThreadPoolExecutor(
                max_workers=max_workers
            ) as executor:
                futures = []
                for i in range(0, len(df), CHUNK_SIZE):
                    chunk_path = f"{self.parquet_path}_{i//CHUNK_SIZE}"
                    # Submit the chunk-saving task to the executor
                    futures.append(
                        executor.submit(
                            self.save_chunk,
                            df,
                            i,
                            CHUNK_SIZE,
                            chunk_path,
                            COMPRESSION_TYPE,
                        )
                    )

                # Wait for all futures to complete
                for future in concurrent.futures.as_completed(futures):
                    # Check for exceptions and handle them if necessary
                    future.result()
        else:
            with fileio.open(self.csv_path, mode="wb") as f:
                df.to_csv(f, index=True)