# Copyright 2025 Snowflake Inc.
# SPDX-License-Identifier: Apache-2.0

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

# http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import json
import logging

from typing import Optional, Union

import pandera as pa

from hypothesis.strategies import DrawFn, SearchStrategy, composite

from snowflake.hypothesis_snowpark.constants import (
    CUSTOM_DATA_COLUMNS_KEY,
    CUSTOM_DATA_KEY,
    CUSTOM_DATA_NAME_KEY,
    CUSTOM_DATA_TYPE_KEY,
    PANDERA_SCHEMA_KEY,
    PYSPARK_ARRAY_TYPE,
    PYSPARK_BINARY_TYPE,
    PYSPARK_DATE_TYPE,
    PYSPARK_STRING_TYPE,
    PYSPARK_TO_SNOWPARK_SUPPORTED_TYPES,
)
from snowflake.hypothesis_snowpark.custom_strategies import update_pandas_df_strategy
from snowflake.hypothesis_snowpark.logging_utils import log
from snowflake.hypothesis_snowpark.strategies_utils import (
    apply_custom_null_values,
    generate_snowpark_dataframe,
    generate_snowpark_schema,
    load_json_schema,
    replace_surrogate_chars,
)
from snowflake.hypothesis_snowpark.telemetry.telemetry import report_telemetry
from snowflake.snowpark import DataFrame, Session


LOGGER = logging.getLogger(__name__)


@report_telemetry(params_list=["schema"])
@log
def dataframe_strategy(
    schema: Union[str, pa.DataFrameSchema], session: Session, size: Optional[int] = None
) -> SearchStrategy[DataFrame]:
    """Create a Hypothesis strategy for generating Snowpark DataFrames based on a given schema.

    Args:
        schema: A schema defining the columns, data types and checks that the generated DataFrame
                should satisfy. This can be a path to a JSON schema file generated by the
                :func:`snowflake.snowpark_checkpoints_collector.collect_dataframe_checkpoint`
                function when the collection mode is set to `SCHEMA`, or a Pandera DataFrameSchema
                object.
        session: The Snowpark session to use for creating the DataFrames.
        size: The number of rows to generate for each DataFrame. If not specified, the strategy will
              generate DataFrames of different sizes.

    Examples:
        Generate a Snowpark DataFrame from a JSON schema file:

        >>> from hypothesis import given
        >>> from snowflake.hypothesis_snowpark import dataframe_strategy
        >>> from snowflake.snowpark import DataFrame, Session
            <BLANKLINE>
        >>> @given(
        ...     df=dataframe_strategy(
        ...         schema="path/to/schema.json",
        ...         session=Session.builder.getOrCreate(),
        ...         size=10,
        ...     )
        ... )
        >>> def test_my_function(df: DataFrame):
        ...     ...

        Generate a Snowpark DataFrame from a Pandera DataFrameSchema object:

        >>> import pandera as pa
        >>> from hypothesis import given
        >>> from snowflake.hypothesis_snowpark import dataframe_strategy
        >>> from snowflake.snowpark import DataFrame, Session
            <BLANKLINE>
        >>> @given(
        ...    df=dataframe_strategy(
        ...        schema=pa.DataFrameSchema(
        ...            {
        ...                "A": pa.Column(pa.Int, checks=pa.Check.in_range(0, 10)),
        ...                "B": pa.Column(pa.Bool),
        ...            }
        ...        ),
        ...        session=Session.builder.getOrCreate(),
        ...        size=10,
        ...    )
        ... )
        >>> def test_my_function(df: DataFrame):
        ...     ...

        You can control aspects like the maximum number of test cases, the deadline for each test
        execution, verbosity levels and many others using the Hypothesis @settings decorator.

        >>> from datetime import timedelta
        >>> from hypothesis import given, settings
        >>> from snowflake.hypothesis_snowpark import dataframe_strategy
        >>> from snowflake.snowpark import DataFrame, Session
            <BLANKLINE>
        >>> @given(
        ...     df=dataframe_strategy(
        ...         schema="path/to/schema.json",
        ...         session=Session.builder.getOrCreate(),
        ...         size=10,
        ...     )
        ... )
        >>> @settings(
        ...     deadline=timedelta(milliseconds=800),
        ...     max_examples=25,
        ... )
        >>> def test_my_function(df: DataFrame):
        ...     ...

    Returns:
        A Hypothesis strategy that generates Snowpark DataFrames.

    """
    if not session:
        raise ValueError("Session cannot be None.")

    if isinstance(schema, pa.DataFrameSchema):
        return _dataframe_strategy_from_object_schema(schema, session, size)

    if isinstance(schema, str) and schema.endswith(".json"):
        return _dataframe_strategy_from_json_schema(schema, session, size)

    raise ValueError(
        "Schema must be a path to a JSON schema file or a Pandera DataFrameSchema object."
    )


def _dataframe_strategy_from_object_schema(
    schema: pa.DataFrameSchema, session: Session, size: Optional[int] = None
) -> SearchStrategy[DataFrame]:
    """Create a Hypothesis strategy for generating Snowpark DataFrames based on a Pandera DataFrameSchema object.

    Args:
        schema: The Pandera DataFrameSchema object.
        session: The Snowpark session to use for creating the DataFrames.
        size: The number of rows to generate.

    Returns:
        A Hypothesis strategy that generates Snowpark DataFrames.

    """
    LOGGER.info(
        "Starting to create a Hypothesis strategy from the given Pandera DataFrameSchema object"
    )
    original_schema = copy.deepcopy(schema)
    str_columns = []

    for column in schema.columns.values():
        if isinstance(column.dtype, pa.engines.pandas_engine.Date):
            # Data generation for date type is currently unsupported by Pandera.
            # As a workaround, we can change the data type to pa.DateTime to
            # avoid an exception and let Snowpark handle the conversion to Date.
            LOGGER.warning(
                "Data generation for Date type is not supported; converting column '%s' to Timestamp",
                column.name,
            )
            column.dtype = pa.DateTime
        elif isinstance(column.dtype, pa.String):
            str_columns.append(column.name)

    @composite
    @log
    def _dataframe_strategy(draw: DrawFn) -> DataFrame:
        pandas_strategy = schema.strategy(size=size)
        pandas_df = draw(pandas_strategy)
        LOGGER.info("Generated a Pandas DataFrame with shape: %s", pandas_df.shape)
        pandas_df = replace_surrogate_chars(pandas_df, str_columns)
        snowpark_schema = generate_snowpark_schema(original_schema)
        snowpark_df = generate_snowpark_dataframe(pandas_df, snowpark_schema, session)
        return snowpark_df

    LOGGER.info(
        "Successfully created a Hypothesis strategy from the given Pandera DataFrameSchema object"
    )
    return _dataframe_strategy()


def _dataframe_strategy_from_json_schema(
    schema: str, session: Session, size: Optional[int] = None
) -> SearchStrategy[DataFrame]:
    """Create a Hypothesis strategy for generating Snowpark DataFrames based on a JSON schema.

    Args:
        schema: The path to the JSON schema file.
        session: The Snowpark session to use for creating the DataFrames.
        size: The number of rows to generate.

    Returns:
        A Hypothesis strategy that generates Snowpark DataFrames.

    """
    LOGGER.info(
        "Starting to create a Hypothesis strategy from the given JSON schema file"
    )
    json_schema_dict = load_json_schema(schema)
    pandera_schema = json_schema_dict.get(PANDERA_SCHEMA_KEY)
    custom_data = json_schema_dict.get(CUSTOM_DATA_KEY)

    if not (pandera_schema and custom_data):
        raise ValueError(
            f"Invalid JSON schema. The JSON schema must contain '{PANDERA_SCHEMA_KEY}' and '{CUSTOM_DATA_KEY}' keys."
        )

    custom_data_columns = custom_data.get(CUSTOM_DATA_COLUMNS_KEY, [])
    not_supported_columns, columns_with_custom_strategy, str_columns = [], [], []

    for column in custom_data_columns:
        dtype = column.get(CUSTOM_DATA_TYPE_KEY)
        if dtype not in PYSPARK_TO_SNOWPARK_SUPPORTED_TYPES:
            not_supported_columns.append(column)
        elif dtype in (PYSPARK_ARRAY_TYPE, PYSPARK_BINARY_TYPE, PYSPARK_DATE_TYPE):
            columns_with_custom_strategy.append(column)
        elif dtype == PYSPARK_STRING_TYPE:
            str_columns.append(column.get(CUSTOM_DATA_NAME_KEY))

    if not_supported_columns:
        raise ValueError(
            f"The following data types are not supported by the Snowpark DataFrame strategy: "
            f"{[column.get(CUSTOM_DATA_TYPE_KEY) for column in not_supported_columns]}"
        )

    LOGGER.info("Creating a Pandera DataFrameSchema object from the given schema")
    df_schema = pa.DataFrameSchema.from_json(json.dumps(pandera_schema))
    LOGGER.debug("Loaded Pandera DataFrameSchema: %s", repr(df_schema))
    df_schema = _process_dataframe_schema(df_schema)

    @composite
    @log
    def _dataframe_strategy(draw: DrawFn) -> DataFrame:
        pandas_strategy = df_schema.strategy(size=size)
        pandas_df = draw(pandas_strategy)
        LOGGER.info("Generated a Pandas DataFrame with shape: %s", pandas_df.shape)
        pandas_df = draw(
            update_pandas_df_strategy(pandas_df, columns_with_custom_strategy)
        )
        pandas_df = apply_custom_null_values(pandas_df, custom_data)
        pandas_df = replace_surrogate_chars(pandas_df, str_columns)
        snowpark_schema = generate_snowpark_schema(df_schema, custom_data)
        snowpark_df = generate_snowpark_dataframe(pandas_df, snowpark_schema, session)
        return snowpark_df

    LOGGER.info(
        "Successfully created a Hypothesis strategy from the given JSON schema file"
    )
    return _dataframe_strategy()


def _process_dataframe_schema(df_schema: pa.DataFrameSchema) -> pa.DataFrameSchema:
    df_schema_copy = copy.copy(df_schema)

    for column_obj in df_schema_copy.columns.values():
        if isinstance(column_obj.dtype, pa.engines.pandas_engine.Date):
            # Data generation for date type is currently unsupported by Pandera. As a workaround,
            # we can change the data type to any supported type to avoid an exception and manually
            # generate the dates.
            LOGGER.warning(
                "Data generation for Date type is not supported; converting column '%s' to Timestamp",
                column_obj.name,
            )
            column_obj.dtype = pa.DateTime

    return df_schema_copy
