"""Extract data from OneStream Data Adapters and load it to Redshift Spectrum."""

from typing import Any, Literal

from prefect import flow
from prefect.logging import get_run_logger

from viadot.orchestration.prefect.tasks import (
    create_batch_list_of_custom_subst_vars,
    df_to_redshift_spectrum,
    onestream_to_df,
)
from viadot.orchestration.prefect.utils import with_flow_timeout_param


@flow(
    name="extract--onestream_data_adapters--redshift_spectrum",
    description="Extract data from OneStream Data Adapters and load it to Redshift Spectrum.",
    retries=1,
    retry_delay_seconds=60,
)
@with_flow_timeout_param()
def onestream_data_adapters_to_redshift_spectrum(  # noqa: PLR0913
    base_url: str,
    application: str,
    adapter_name: str,
    to_path: str,
    schema_name: str,
    table: str,
    workspace_name: str = "MainWorkspace",
    adapter_response_key: str = "Results",
    custom_subst_vars: dict[str, list[Any]] | None = None,
    batch_by_subst_vars: bool = False,
    params: dict[str, str] | None = None,
    extension: str = ".parquet",
    if_exists: Literal["overwrite", "append"] = "overwrite",
    if_empty: Literal["warn", "skip", "fail"] = "fail",
    partition_cols: list[str] | None = None,
    index: bool = False,
    compression: Literal["snappy", "gzip", "zstd"] | None = None,
    sep: str = ",",
    aws_config_key: str | None = None,
    credentials_secret: str | None = None,
    onestream_credentials_secret: str | None = None,
    onestream_config_key: str | None = None,
) -> None:
    """Extract data from OneStream Data Adapter and load it into AWS Redshift Spectrum.

    API documentation:
    https://documentation.onestream.com/1384528/Content/REST%20API/OneStream%20WebAPI%20Endpoints.html#:~:text=Data%20Provider%20GetAdoDataSetForAdapter%20endpoint

    This flow retrieves data from a OneStream Data Adapter using provided parameters
    and uploads it to AWS Redshift Spectrum.

    Flow can be executed with the custom substitution variables mapped to its values.
    These variable-value combination is used to filter data inside the data
    adapters table allowing to return only the datasets where rows have
    the value for a given column values set.The column - variable mapping is defined in
    the OneStream Data Adapter configuration.The filtering is done for each
    variable-values combination, hence the flow might generate a couple of API calls
    to get data for all combinations.

    When custom_subst_vars are provided and batch_by_subst_vars is True, the ingestion
    process will be split into batches. Each batch represents one combination of the
    custom_subst_vars - variable:[value]. When batch_by_subst_vars is False, all
    substitution variable combinations are processed together into a single data frame.
    Warning! Processing custom substitution vars without batching might lead to
    out-of-memory errors when each variable-value combination data set size is too big.

    Args:
        base_url (str): OneStream base server URL.
        application (str): OneStream application name.
        adapter_name (str): Data Adapter name to query.
        to_path (str): Path to an S3 folder where the table will be located.
        schema_name (str): AWS Glue catalog database name.
        table (str): AWS Glue catalog table name.
        workspace_name (str): OneStream workspace name. Defaults to "MainWorkspace".
        adapter_response_key (str): Key in the JSON response that contains
            the adapter's returned data. Defaults to "Results".
        custom_subst_vars (dict[str, list[Any]], optional): A dictionary mapping
            substitution variable names to lists of possible values.Substition variables
            used as substitution variables that are mapped to the column names in the
            Data Adapter configuration.They provides a kind of SQL WHERE clause for rows
            filtering where variable refers to a name of a set containings a list
            of possible values.
            For example: "prm_activity_region":["MK", "LG"]
                        refers to column UD3 that has custom variable mapping of
                        prm_activity_region.Data will be extracted only for rows where
                        the values (MK,LG) are present.
            Values can be of any type that can be converted to strings.Defaults to None.
        batch_by_subst_vars (bool): Whether to process data in batches based on
            substitution variable combinations. When True and custom_subst_vars is
            provided, each combination of substitution variables and its values will be
            processed as a separate batch, uploading individual parquet files in S3.
            When False,all substitution variable combinations are collected in a loop
            and are uploaded together in a single operation. Defaults to False.
        params (dict[str, str], optional): Additional API parameters. Defaults to None.
        extension (str): Required file type. Accepted formats: 'csv', 'parquet'.
            Defaults to ".parquet".
        if_exists (Literal["overwrite", "append"], optional): Whether to 'overwrite'
            or 'append' to existing table. Defaults to "overwrite".
        if_empty (Literal["warn", "skip", "fail"], optional): What to do if the
            API returns no data. Defaults to "fail".
        partition_cols (list[str], optional): Columns used to create partitions.
            Only applies when dataset=True. Defaults to None.
        index (bool): Write row names (index). Defaults to False.
        compression (Literal["snappy", "gzip", "zstd"], optional): Compression style
            (None, snappy, gzip, zstd). Defaults to None.
        sep (str): Field delimiter for the output file. Defaults to ','.
        aws_config_key (str, optional): Key in viadot config for AWS credentials.
            Defaults to None.
        credentials_secret (str, optional): Name of Prefect secret block with AWS
            credentials. Defaults to None.
        onestream_credentials_secret (str, optional): Name of secret storing OneStream
            credentials. Defaults to None.
        onestream_config_key (str, optional)): Key in viadot config for OneStream
            credentials.Defaults to None.
    """
    logger = get_run_logger()

    logger.info(f"Starting OneStream Data Adapter extraction: {adapter_name}")

    if batch_by_subst_vars and not custom_subst_vars:
        msg = (
            "Invalid parameter combination: batch_by_subst_vars=True requires "
            "custom_subst_vars to be provided. Either set batch_by_subst_vars=False "
            "or provide a custom_subst_vars dictionary."
        )
        logger.error(msg)
        raise ValueError(msg)
    if custom_subst_vars and batch_by_subst_vars:
        # Process data in batches - each substitution variable combination separately
        custom_subst_vars_batch_list = create_batch_list_of_custom_subst_vars(
            custom_subst_vars
        )
        logger.info(
            f"Processing {len(custom_subst_vars_batch_list)} batches based on substitution variable combinations."
        )

        for i, custom_subst_var in enumerate(custom_subst_vars_batch_list, 1):
            logger.info(
                f"Processing batch {i}/{len(custom_subst_vars_batch_list)}: {custom_subst_var}."
            )
            df = onestream_to_df(
                base_url=base_url,
                application=application,
                api="data_adapter",
                adapter_name=adapter_name,
                workspace_name=workspace_name,
                adapter_response_key=adapter_response_key,
                custom_subst_vars=custom_subst_var,
                if_empty=if_empty,
                params=params,
                credentials_secret=onestream_credentials_secret,
                config_key=onestream_config_key,
            )
            df_to_redshift_spectrum(
                df=df,
                to_path=to_path,
                schema_name=schema_name,
                table=table,
                extension=extension,
                if_exists=if_exists,
                partition_cols=partition_cols,
                index=index,
                compression=compression,
                sep=sep,
                config_key=aws_config_key,
                credentials_secret=credentials_secret,
            )
            if_exists = "append"  # Changed to "append" to add batches to the table.
            logger.info(
                f"Batch {i}/{len(custom_subst_vars_batch_list)} completed successfully."
            )

        logger.info("All batches processed successfully.")
    else:
        # Process all data together - either no custom_subst_vars or batching disabled
        if custom_subst_vars:
            logger.info(
                "Processing all substitution variable combinations as one data frame."
            )
        else:
            logger.info("Processing data without substitution variables.")
        df = onestream_to_df(
            base_url=base_url,
            application=application,
            api="data_adapter",
            adapter_name=adapter_name,
            workspace_name=workspace_name,
            adapter_response_key=adapter_response_key,
            custom_subst_vars=custom_subst_vars,
            params=params,
            credentials_secret=onestream_credentials_secret,
            config_key=onestream_config_key,
        )
        df_to_redshift_spectrum(
            df=df,
            to_path=to_path,
            schema_name=schema_name,
            table=table,
            extension=extension,
            if_exists=if_exists,
            partition_cols=partition_cols,
            index=index,
            compression=compression,
            sep=sep,
            config_key=aws_config_key,
            credentials_secret=credentials_secret,
        )
        logger.info("Data processing completed successfully.")
