from pyspark.sql import Row
from pyspark.sql.functions import *
from pyspark.sql.types import StructField
from easy_spark.easy_delta_helpers import EasyDeltaHelpers
from easy_spark.easy_sql_builder import EasySQLSelectBuilder
from easy_spark.lh_path import LHPath
from easy_spark.table_path import TablePath
from pyspark.sql import SparkSession
from easy_spark.wh_path import WHPath
from datetime import datetime
import numpy as np


class EasyDF:
    _spark: SparkSession = None

    # Constructor
    def __init__(self, df: DataFrame = None, spark: SparkSession = None):
        self._df: DataFrame = df
        if spark is not None:
            EasyDF._spark = spark

    @staticmethod
    def create(spark: SparkSession, df: DataFrame = None) -> 'EasyDF':
        return EasyDF(df, spark)

    @staticmethod
    def create_empty(spark: SparkSession, schema: StructType = None) -> 'EasyDF':
        return EasyDF(None, spark).empty(schema)

    @property
    def df(self) -> DataFrame:
        return self._df

    @df.setter
    def df(self, value: DataFrame):
        self._df = value

    @property
    def current_schema(self) -> StructType:
        return self._df.schema if self._df and self._df.schema else None

    def from_table_path(self, path: TablePath, df_format="delta") -> 'EasyDF':
        self._df = EasyDF._spark.read.format(df_format).load(path.path)
        return self

    def from_path(self, path: str, df_format="delta") -> 'EasyDF':
        self._df = EasyDF._spark.read.format(df_format).load(path)
        return self

    def from_lh_table(self, lh_name: str, table_name: str, limit: int = None) -> 'EasyDF':
        sql_builder = EasySQLSelectBuilder.new().select(["*"]).from_lh_table(lh_name, table_name).limit(limit)
        return self.from_sql(sql_builder.sql)

    def from_sql_builder(self, sql_builder: EasySQLSelectBuilder) -> 'EasyDF':
        return self.from_sql(sql_builder.sql)

    def from_lh_flat_path(self, path: tuple[str, str, str, str], df_format="delta") -> 'EasyDF':
        from easy_spark.path import Path
        return self.from_table_path(LHPath(Path(path[0], path[1]), path[2], path[3]), df_format)

    def from_wh_flat_path(self, path: tuple[str, str, str, str], df_format="delta") -> 'EasyDF':
        from easy_spark.path import Path
        return self.from_table_path(WHPath(Path(path[0], path[1]), path[2], path[3]), df_format)

    def from_dict(self, records: dict, schema: StructType = None) -> 'EasyDF':
        new_schema = schema if schema is not None else self.current_schema
        self._df = EasyDF._spark.createDataFrame([Row(**records)], new_schema)
        return self

    def from_list(self, records: list[dict], schema: StructType = None) -> 'EasyDF':
        new_schema = schema if schema is not None else self.current_schema
        self._df = EasyDF._spark.createDataFrame(records, new_schema)
        return self

    def from_tuple(self, records: list[tuple], schema: StructType = None) -> 'EasyDF':
        new_schema = schema if schema is not None else self.current_schema
        self._df = EasyDF._spark.createDataFrame(records, new_schema)
        return self

    def from_sql(self, sql: str) -> 'EasyDF':
        self._df = EasyDF._spark.sql(sql)
        return self

    def from_json(self, json: str) -> 'EasyDF':
        self._df = EasyDF._spark.read.json(json)
        return self

    def from_csv(self, csv_path: str, header=True, seperator=",", infer_schema=True) -> 'EasyDF':
        self._df = EasyDF._spark.read.option("header", header).csv(csv_path, sep=seperator, inferSchema=infer_schema)

        return self

    def from_lh_name_path(self, name: str, table_name: str, limt: int = None) -> 'EasyDF':
        if limt:
            self._df = EasyDF._spark.sql(f"SELECT * FROM {name}.{table_name} LIMIT {limt}")
        else:
            self._df = EasyDF._spark.sql(f"SELECT * FROM {name}.{table_name}")

        return self

    def from_instance_sql_builder(self):
        pass

    def empty(self, schema: StructType = None) -> 'EasyDF':
        new_schema = schema if schema is not None else self.current_schema
        self._df = EasyDF._spark.createDataFrame([], new_schema)

        return self

    def filter_by_filter(self, condition: str) -> 'EasyDF':
        if condition:
            self._df = self._df.filter(condition)
        return self

    def filter(self, keys: dict[str, any] = None):
        if keys:
            for key in keys:
                self._df = self._df[self._df[key] == keys[key]]

        return self

    def filter_using_filter(self, keys: dict[str, any] = None):
        if keys:
            conditions = EasyDeltaHelpers.build_condition(keys)
            self._df = self._df.filter(conditions)

        return self

    def where(self, condition: str) -> 'EasyDF':
        if condition:
            self._df = self._df.where(condition)
        return self

    def combine_from_df(self, df: DataFrame, type: str = 'unionByName',
                        allowMissingColumns: bool = True) -> 'EasyDF':
        self._df = EasyDeltaHelpers.combine_from_dfs([self._df, df], type, allowMissingColumns)
        return self

    def replace_column_empty_spaces(self, value="") -> 'EasyDF':
        for col_name in self._df.columns:
            self._df = self._df.withColumnRenamed(col_name, col_name.replace(" ", value))
        return self

    def drop_column_if_exists(self, column_name: str) -> 'EasyDF':
        if column_name in self._df.columns:
            self._df = self._df.drop(column_name)
        return self

    def add_column(self, column_name: str, value: any) -> 'EasyDF':
        self._df = self._df.withColumn(column_name, lit(value))
        return self

    def capitalize_column_between_spaces(self) -> 'EasyDF':
        for col_name in self._df.columns:
            self._df = self._df.withColumnRenamed(col_name, col_name.title())

        return self

    def replace_nan(self, value=None) -> 'EasyDF':
        self._df = self._df.replace({np.nan: value}).replace({"nan": value})

        return self

    def overwrite_types(self, struct_types: dict | list[StructField] = None) -> 'EasyDF':
        struct_types = EasyDeltaHelpers.build_struct_fields(struct_types)

        for struct_t in struct_types:
            filed_r = struct_t.name.replace("`", "")
            # Check if field is in df
            if filed_r not in self._df.columns:
                continue
            self._df = self._df.withColumn(filed_r, col(struct_t.name).cast(struct_t.dataType))

        return self

    def add_audit_fields(self) -> 'EasyDF':
        now_date = datetime.now()
        date_only = now_date.date()
        self._df = self._df.withColumn("CreatedDate", lit(now_date))
        self._df = self._df.withColumn("CreatedYear", lit(now_date.year))
        self._df = self._df.withColumn("CreatedMonth", lit(now_date.month))
        self._df = self._df.withColumn("CreatedDay", lit(now_date.day))
        self._df = self._df.withColumn("CreatedDateOnly", lit(date_only))

        return self

    def append_from_dict(self, record: dict) -> 'EasyDF':
        row = Row(**record)

        return self.append_from_row(row)

    def append_from_row(self, row: Row) -> 'EasyDF':
        df = EasyDF._spark.createDataFrame([row], self._df.schema)
        self._df = self.combine_from_df(df, type='union', allowMissingColumns=True)._df
        return self

    def save_from_table_path(self, path: TablePath, df_format="delta", mode="overwrite",
                             merge_option: str = "overwriteSchema") -> 'EasyDF':
        return self.save_from_path(path.path, df_format, mode, merge_option)

    def save_as_table(self, path: str, df_format="delta", mode="overwrite",
                      merge_option: str = "overwriteSchema") -> 'EasyDF':
        self._df.write.format(df_format).mode(mode).option(merge_option, "true").saveAsTable(path)
        return self

    def save_from_path(self, path: str, df_format="delta", mode="overwrite",
                       merge_option: str = "overwriteSchema", add_audit_fields=False) -> 'EasyDF':
        self._df.write.format(df_format).mode(mode).option(merge_option, "true").save(path)
        return self
