from pyspark.sql import DataFrame
from delta import DeltaTable
from ._save import Save
from typing import Union, List, Any
from ._save_mode_type import SaveModeOptions
from pydantic import BaseModel, Field
from enum import Enum


class MatchOperator(Enum):
    ANY_NOT_EQUAL_EXCEPT = "any_not_equal_except"
    ALL_NOT_EQUAL_EXCEPT = "all_not_equal_except"
    ANY_EQUAL_EXCEPT = "any_equal_except"
    ALL_EQUAL_EXCEPT = "all_equal_except"


class Match(BaseModel):

    operator: MatchOperator = Field(...)
    columns: List[str] = Field(...)


class Merge(BaseModel):

    join: str = Field(...)
    update: Union[str, bool, Match] = Field(...)
    insert: Union[str, bool, Match] = Field(...)
    delete: Union[str, bool, Match] = Field(...)


class ErrorIfExistsSave(Save):
    def write(self):
        super().write()
        (
            self.dataset.dataframe.write.format(self.dataset.format.value)
            .options(**self.dataset.write.options)
            .mode(SaveModeOptions.ERROR_IF_EXISTS.value)
            .partitionBy(*self.dataset.partitioned_by)
            .save(self.dataset.path)
        )


class AppendSave(Save):
    def write(self):
        super().write()
        (
            self.dataset.dataframe.write.format(self.dataset.format.value)
            .options(**self.dataset.write.options)
            .mode(SaveModeOptions.APPEND.value)
            .partitionBy(*self.dataset.partitioned_by)
            .save(self.dataset.path)
        )


class OverwriteSchemaSave(Save):
    def write(self):
        super().write()
        options = self.dataset.write.options
        options[SaveModeOptions.OVERWRITE_SCHEMA.value] = True
        (
            self.dataset.dataframe.write.format(self.dataset.format.value)
            .options(**options)
            .mode(SaveModeOptions.OVERWRITE.value)
            .partitionBy(*self.dataset.partitioned_by)
            .save(self.dataset.path)
        )


class OverwriteSave(Save):
    def write(self):
        super().write()
        (
            self.dataset.dataframe.write.format(self.dataset.format.value)
            .options(**self.dataset.write.options)
            .mode(SaveModeOptions.OVERWRITE.value)
            .partitionBy(*self.dataset.partitioned_by)
            .save(self.dataset.path)
        )


class IgnoreSave(Save):
    def write(self):
        super().write()
        (
            self.dataset.dataframe.write.format(self.dataset.format.value)
            .options(**self.dataset.write.options)
            .mode(SaveModeOptions.IGNORE.value)
            .partitionBy(*self.dataset.partitioned_by)
            .save(self.dataset.path)
        )


class MergeSave(Save):
    def __init__(self, **data: Any) -> None:
        super().__init__(**data)

    join: str = Field(...)
    update: Union[str, bool, Match] = Field(default=False)
    insert: Union[str, bool, Match] = Field(default=False)
    delete: Union[str, bool, Match] = Field(default=False)

    def write(self):
        super().write()

        merge_update_match = self._derive_any_except_match(
            self.dataset.merge_update_match, self.dataset.dataframe
        )
        merge_insert_match = self._derive_any_except_match(
            self.dataset.merge_insert_match, self.dataset.dataframe
        )

        tbl = DeltaTable.forPath(self.dataset.context.spark, self.dataset.path)
        merger = tbl.alias("dst").merge(
            self.dataset.dataframe.alias("src"), self.dataset.merge_join
        )

        if self.dataset.merge_update:
            if merge_update_match:
                merger = merger.whenMatchedUpdateAll(merge_update_match)
            else:
                merger = merger.whenMatchedUpdateAll()

        if self.dataset.merge_insert:
            if merge_insert_match:
                merger = merger.whenNotMatchedInsertAll(merge_insert_match)
            else:
                merger = merger.whenNotMatchedInsertAll()

        if self.dataset.merge_delete:
            if self.dataset.merge_delete_match:
                merger = merger.whenMatchedDelete(self.dataset.merge_delete_match)

        merger.execute()

    def _derive_any_except_match(self, merge_match: Union[str, dict], df: DataFrame):

        if isinstance(merge_match, dict):
            any_except = merge_match.get("any_not_equal_except")
            derived_merge_match = [
                f"src.{c} != dst.{c}"
                for c in df.columns
                if c not in any_except and not c.startswith("_")
            ]
            derived_merge_match = " or ".join(derived_merge_match)
            return derived_merge_match
        else:
            return merge_match
