# 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 json
import logging
import os
import random
import re

from contextlib import contextmanager
from typing import Optional, Union

import numpy as np
import pandas as pd
import pandera as pa

from snowflake.hypothesis_snowpark.constants import (
    CUSTOM_DATA_COLUMNS_KEY,
    CUSTOM_DATA_NAME_KEY,
    CUSTOM_DATA_ROWS_COUNT,
    CUSTOM_DATA_ROWS_NULL_COUNT_KEY,
    CUSTOM_DATA_TYPE_KEY,
)
from snowflake.hypothesis_snowpark.datatype_mapper import (
    pandera_dtype_to_snowpark_dtype,
    pyspark_dtype_to_snowpark_dtype,
)
from snowflake.snowpark import DataFrame, Session
from snowflake.snowpark.types import StructField, StructType


LOGGER = logging.getLogger(__name__)


def load_json_schema(json_schema: str) -> dict:
    """Load the JSON schema from a file.

    Args:
        json_schema: The path to the JSON file.

    Returns:
        The JSON schema as a dictionary.

    """
    if not os.path.isfile(json_schema):
        raise ValueError(f"Invalid JSON schema path: {json_schema}")

    try:
        with open(json_schema, encoding="utf-8") as file:
            return json.load(file)
    except (OSError, json.JSONDecodeError) as e:
        raise ValueError(f"Error reading JSON schema file: {e}") from None


def generate_snowpark_dataframe(
    pandas_df: pd.DataFrame,
    snowpark_schema: StructType,
    session: Session,
) -> DataFrame:
    """Generate a Snowpark DataFrame from a Pandas DataFrame.

    Args:
        pandas_df: The Pandas DataFrame to convert.
        snowpark_schema: The Snowpark schema to use for creating the DataFrame.
        session: The Snowpark session to use for creating the DataFrame.

    Returns:
        A Snowpark DataFrame.

    """
    pandas_df = pandas_df.replace({np.nan: None})

    # Snowpark ignores the schema argument if the data argument is a pandas dataframe. We need to convert the pandas
    # dataframe into a list of tuples to be able to specify a schema.
    data = list(pandas_df.itertuples(index=False, name=None))

    # In some scenarios, the create_dataframe function performs some random operations (e.g., generating temporary table
    # names or temporary stage names). Due to Hypothesis controls randomness during its tests, we need to temporarily
    # set a random seed to ensure that the names generated by Snowpark are unique to avoid conflicts.
    with temporary_random_seed():
        LOGGER.info(
            "Creating Snowpark DataFrame with schema=%s and size=%s",
            snowpark_schema,
            len(data),
        )
        snowpark_df = session.create_dataframe(data=data, schema=snowpark_schema)
    return snowpark_df


def apply_custom_null_values(
    pandas_df: pd.DataFrame, custom_data: dict
) -> pd.DataFrame:
    """Apply null values to a Pandas DataFrame based on the custom data.

    This function modifies the DataFrame in place and returns a new DataFrame with null values applied.

    Args:
        pandas_df: The Pandas DataFrame to apply null values to.
        custom_data: The custom data associated with the Pandera schema.

    Returns:
        A Pandas DataFrame with null values applied.

    """
    null_proportions = {}
    for col in custom_data.get(CUSTOM_DATA_COLUMNS_KEY, []):
        rows_count = col.get(CUSTOM_DATA_ROWS_COUNT, 0)
        if rows_count > 0:
            null_proportions[col[CUSTOM_DATA_NAME_KEY]] = (
                col[CUSTOM_DATA_ROWS_NULL_COUNT_KEY] / rows_count
            )

    df_copy = pandas_df.copy()
    total_rows = len(df_copy)

    for column, target_null_proportion in null_proportions.items():
        if column not in df_copy.columns:
            LOGGER.warning(
                "Column '%s' not found in DataFrame during null value correction. Skipping it.",
                column,
            )
            continue

        required_nulls = round(total_rows * target_null_proportion)
        current_nulls = df_copy[column].isnull().sum()

        if current_nulls < required_nulls:
            LOGGER.debug(
                "Column '%s' has %s null values but %s are required.",
                column,
                current_nulls,
                required_nulls,
            )
            additional_nulls = required_nulls - current_nulls
            non_null_positions = np.where(df_copy[column].notnull())[0]
            selected_positions = np.random.choice(
                non_null_positions, size=additional_nulls, replace=False
            )
            df_copy[column] = df_copy[column].astype(object)
            df_copy.iloc[selected_positions, df_copy.columns.get_loc(column)] = None
            LOGGER.info(
                "Added %s null values to column '%s'",
                additional_nulls,
                column,
            )

    return df_copy.replace({np.nan: None})


@contextmanager
def temporary_random_seed(seed: Union[int, float, str, bytes, bytearray, None] = None):
    """Set a temporary random seed within a context.

    This context manager saves the current state of the random number generator,
    sets a new seed, and then restores the original state when exiting the context.

    Args:
        seed: The seed to set for the random number generator. If not provided, seeds from current time or from an
              operating system specific randomness source if available.

    Yields:
        None: The context manager does not return anything.

    """
    current_state = random.getstate()
    random.seed(seed)
    try:
        yield
    finally:
        random.setstate(current_state)


def generate_snowpark_schema(
    pandera_schema: pa.DataFrameSchema, custom_data: Optional[dict] = None
) -> StructType:
    """Generate a Snowpark schema from a Pandera schema.

    Args:
        pandera_schema: The Pandera schema to convert.
        custom_data: The custom data associated with the Pandera schema.

    Returns:
        A Snowpark struct type.

    """
    struct_fields = []

    if custom_data is None:
        for column_name, column_data in pandera_schema.columns.items():
            snowpark_dtype = pandera_dtype_to_snowpark_dtype(column_data.dtype)
            struct_field = StructField(
                column_identifier=column_name,
                datatype=snowpark_dtype,
                nullable=column_data.nullable,
            )
            struct_fields.append(struct_field)
    else:
        for column_name, column_data in pandera_schema.columns.items():
            custom_data_column = next(
                (
                    column
                    for column in custom_data.get(CUSTOM_DATA_COLUMNS_KEY, [])
                    if column.get(CUSTOM_DATA_NAME_KEY) == column_name
                ),
                None,
            )

            if custom_data_column is None:
                raise ValueError(f"Column '{column_name}' is missing from custom_data")

            pyspark_type = custom_data_column.get(CUSTOM_DATA_TYPE_KEY)

            if pyspark_type is None:
                raise ValueError(
                    f"Type for column '{column_name}' is missing from custom_data"
                )

            snowpark_type = pyspark_dtype_to_snowpark_dtype(pyspark_type)

            struct_field = StructField(
                column_identifier=column_name,
                datatype=snowpark_type,
                nullable=column_data.nullable,
            )
            struct_fields.append(struct_field)

    return StructType(struct_fields)


def replace_surrogate_chars(df: pd.DataFrame, columns: list[str]) -> pd.DataFrame:
    """Replace surrogate characters in the specified columns of a Pandas DataFrame with a space.

    Args:
        df: The Pandas DataFrame to process.
        columns: The list of columns to process.

    Returns:
        A new DataFrame with surrogate characters replaced.

    """
    surrogate_pattern = re.compile(r"[\uD800-\uDFFF]")
    replacement = " "
    df_copy = df.copy()

    for column in columns:
        if column not in df_copy.columns:
            LOGGER.warning(
                "Column '%s' not found in DataFrame during surrogate replacement. Skipping it.",
                column,
            )
            continue

        LOGGER.info("Replacing surrogate characters in column '%s'", column)
        df_copy[column] = df_copy[column].apply(
            lambda value: re.sub(surrogate_pattern, replacement, value)
            if isinstance(value, str)
            else value
        )

    return df_copy
