""" This module contains the core functions of the pytortilla package.
Functions:
    - create: Create a tortilla file.
    - load: Load the metadata of a tortilla file.
    - compile: Compile a subset of a tortilla file.    
"""

import concurrent.futures
import mmap
import pathlib
from typing import List, Union

import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import tqdm

import pytortilla.core
import pytortilla.utils
import pytortilla.utils_gdal


def create(
    files: Union[List[str], List[pathlib.Path]],
    output: Union[str, pathlib.Path],
    file_format: pytortilla.utils_gdal.GDAL_FILES,
    nworkers: int = 4,
    chunk_size: int = 1024 * 1024 * 100,
    quiet: bool = False,
) -> pathlib.Path:
    """Create a tortilla 🫓

    A tortilla is a new simple format for storing same format files
    optimized for very fast random access.

    Args:
        files (List[str]): The list of files to be included in the
            tortilla. All files must have the same format (same extension).
        file_format (GDAL_FILES): The format of the files. This format
            must be one of the GDAL formats. For example, "GTiff", "COG",
            "PNG", etc.
        nworkers (int, optional): The number of workers to use when writing
            the tortilla. Defaults to 4.
        chunk_size (int, optional): The size of the chunks to use when writing
            the tortilla. Defaults to 1024 * 1024 * 100 (100 MB).

    Returns:
        pathlib.Path: The tortilla file.
    """

    if file_format not in pytortilla.utils_gdal.GDAL_FILES.__args__:
        raise ValueError(
            f"Invalid file format: {file_format}. Must be one of {pytortilla.utils_gdal.GDAL_FILES.__args__}."
        )

    # Define the magic number
    MB: bytes = b"#y"

    # Get the total size of the files
    dict_bytes, bytes_counter = pytortilla.utils.process_files_concurrently(
        files=files, nworkers=nworkers
    )

    # Create the FOOTER metadata as pandas
    metadata = pd.DataFrame.from_dict(
        dict_bytes, orient="index", columns=["tortilla:offset", "tortilla:length"]
    ).reset_index(names="tortilla:id")

    # Create an in-memory Parquet file with BufferOutputStream
    with pa.BufferOutputStream() as sink:
        pq.write_table(
            pa.Table.from_pandas(metadata),
            sink,
            compression="zstd",  # Highly efficient codec
            compression_level=22,  # Maximum compression for Zstandard
            use_dictionary=False,  # Optimizes for repeated values
        )
        # return a blob of the in-memory Parquet file as bytes
        # Obtain the FOOTER metadata
        FOOTER: bytes = sink.getvalue().to_pybytes()

    # Define the FOOTER length and offset
    FL: bytes = len(FOOTER).to_bytes(8, "little")
    FO: bytes = (50 + bytes_counter).to_bytes(8, "little")

    # Data Format of items
    DF: int = file_format.encode().ljust(32)

    # Get the total size of the files
    total_size: int = 50 + bytes_counter + len(FOOTER)

    # If the folder does not exist, create it
    output = pathlib.Path(output)
    output.parent.mkdir(parents=True, exist_ok=True)

    # Create an empty file
    with open(output, "wb") as f:
        f.truncate(total_size)

    # Define the function to write into the main file
    def write_file(file, start):
        """read the file in chunks"""
        with open(file, "rb") as g:
            while True:
                chunk = g.read(chunk_size)
                if not chunk:
                    break
                mm[start : (start + len(chunk))] = chunk
                start += len(chunk)

    # Cook the tortilla 🫓
    with open(output, "r+b") as f:
        with mmap.mmap(f.fileno(), 0) as mm:
            # Write the magic number
            mm[:2] = MB

            # Write the FOOTER offset
            mm[2:10] = FO

            # Write the FOOTER length
            mm[10:18] = FL

            # Write the decoder
            mm[18:50] = DF

            # Write the DATA
            message = pytortilla.utils.tortilla_message()
            with concurrent.futures.ThreadPoolExecutor(
                max_workers=nworkers
            ) as executor:
                futures = []
                for file in files:
                    start, _ = dict_bytes[pathlib.Path(file).stem]
                    futures.append(executor.submit(write_file, file, start))

                # Wait for all futures to complete
                if not quiet:
                    list(
                        tqdm.tqdm(
                            concurrent.futures.as_completed(futures),
                            total=len(futures),
                            desc=message,
                            unit="file",
                        )
                    )
                else:
                    concurrent.futures.wait(futures)

            # Write the FOOTER
            mm[(50 + bytes_counter) :] = FOOTER

    return output


def load(file: Union[str, pathlib.Path]) -> pd.DataFrame:
    """Load the metadata of a tortilla file.

    Args:
        file (Union[str, pathlib.Path]): A tortilla file.

    Returns:
        pd.DataFrame: The metadata of the tortilla file.
    """
    if isinstance(file, str):
        if pytortilla.utils.is_valid_url(file):
            metadata: pd.DataFrame = pytortilla.core.read_tortilla_metadata_online(file)
            partial_file: pd.Series = metadata.apply(
                lambda row: f"/vsisubfile/{row['tortilla:offset']}_{row['tortilla:length']},/vsicurl/{file}",
                axis=1,
            )
        else:
            metadata: pd.DataFrame = pytortilla.core.read_tortilla_metadata_local(file)
            partial_file: pd.Series = metadata.apply(
                lambda row: f"/vsisubfile/{row['tortilla:offset']}_{row['tortilla:length']},{file}",
                axis=1,
            )

    elif isinstance(file, pathlib.Path):
        metadata: pd.DataFrame = pytortilla.core.read_tortilla_metadata_local(file)
        partial_file: pd.Series = metadata.apply(
            lambda row: f"/vsisubfile/{row['tortilla:offset']}_{row['tortilla:length']},{file}",
            axis=1,
        )
    else:
        raise ValueError("Invalid file type. Must be a string or pathlib.Path.")

    metadata["tortilla:subfile"] = partial_file

    return metadata


def compile(
    dataset: pd.DataFrame,
    output: Union[str, pathlib.Path],
    chunk_size: int = 1024 * 1024 * 100,
    nworkers: int = 4,
    force: bool = False,
    quiet: bool = False,
) -> pd.DataFrame:
    """Prepare a subset of a Tortilla file and write it to a new local file.

    Args:
        dataframe (pd.DataFrame): A subset of the metadata of a Tortilla file.
        output_folder (Union[str, pathlib.Path]): The folder where the Tortilla file
            will be saved. If the folder does not exist, it will be created.
        chunk_size (int, optional): The size of the chunks to use when writing
            the tortilla. Defaults to 1024 * 1024 * 100 (100 MB).
        nworkers (int, optional): The number of workers to use when writing
            the tortilla. Defaults to 4. Only used when the Tortilla file is
            local.
        force (bool, optional): If True, the function overwrites the file if it
            already exists. By default, it is False.
        quiet (bool, optional): If True, the function does not print any
            message. By default, it is False.
    Returns:
        pd.DataFrame: The metadata of the new Tortilla file.
    """

    # If the folder does not exist, create it
    output = pathlib.Path(output)
    output.parent.mkdir(parents=True, exist_ok=True)

    # Check if the file already exists
    if output.exists() and force:
        output.unlink()

    # Remove the index from the previous dataset
    dataset = dataset.copy()
    dataset.sort_values("tortilla:offset", inplace=True)
    dataset.reset_index(drop=True, inplace=True)

    # Compile your tortilla
    if dataset["tortilla:mode"].iloc[0] == "local":
        pytortilla.core.compile_local(dataset, output, chunk_size, nworkers, quiet)

    if dataset["tortilla:mode"].iloc[0] == "online":
        pytortilla.core.compile_online(dataset, output, chunk_size, quiet)

    return load(output)
