#
# Copyright (c) 2012-2024 Snowflake Computing Inc. All rights reserved.
#

import json
import os
import random

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


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():
        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:
            continue

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

        if 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

    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)
