#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You 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.
#
from pyspark.sql.connect.utils import check_dependencies

check_dependencies(__name__)

from typing import Dict
from typing import Optional, Union, List, overload, Tuple, cast, Callable
from typing import TYPE_CHECKING

from pyspark.sql.connect.plan import Read, DataSource, LogicalPlan, WriteOperation, WriteOperationV2
from pyspark.sql.types import StructType
from pyspark.sql.utils import to_str
from pyspark.sql.readwriter import (
    DataFrameWriter as PySparkDataFrameWriter,
    DataFrameReader as PySparkDataFrameReader,
    DataFrameWriterV2 as PySparkDataFrameWriterV2,
)
from pyspark.errors import PySparkAttributeError, PySparkTypeError, PySparkValueError
from pyspark.sql.connect.functions import builtin as F

if TYPE_CHECKING:
    from pyspark.sql.connect.dataframe import DataFrame
    from pyspark.sql.connect._typing import ColumnOrName, OptionalPrimitiveType
    from pyspark.sql.connect.session import SparkSession
    from pyspark.sql.metrics import ExecutionInfo

__all__ = ["DataFrameReader", "DataFrameWriter"]

PathOrPaths = Union[str, List[str]]
TupleOrListOfString = Union[List[str], Tuple[str, ...]]


class OptionUtils:
    def _set_opts(
        self,
        schema: Optional[Union[StructType, str]] = None,
        **options: "OptionalPrimitiveType",
    ) -> None:
        """
        Set named options (filter out those the value is None)
        """
        if schema is not None:
            self.schema(schema)  # type: ignore[attr-defined]
        for k, v in options.items():
            if v is not None:
                self.option(k, v)  # type: ignore[attr-defined]


class DataFrameReader(OptionUtils):
    # TODO(SPARK-40539) Achieve parity with PySpark.

    def __init__(self, client: "SparkSession"):
        self._client = client
        self._format: Optional[str] = None
        self._schema = ""
        self._options: Dict[str, str] = {}

    def format(self, source: str) -> "DataFrameReader":
        self._format = source
        return self

    format.__doc__ = PySparkDataFrameReader.format.__doc__

    def schema(self, schema: Union[StructType, str]) -> "DataFrameReader":
        if isinstance(schema, StructType):
            self._schema = schema.json()
        elif isinstance(schema, str):
            self._schema = schema
        else:
            raise PySparkTypeError(
                errorClass="NOT_STR_OR_STRUCT",
                messageParameters={
                    "arg_name": "schema",
                    "arg_type": type(schema).__name__,
                },
            )
        return self

    schema.__doc__ = PySparkDataFrameReader.schema.__doc__

    def option(self, key: str, value: "OptionalPrimitiveType") -> "DataFrameReader":
        self._options[key] = cast(str, to_str(value))
        return self

    option.__doc__ = PySparkDataFrameReader.option.__doc__

    def options(self, **options: "OptionalPrimitiveType") -> "DataFrameReader":
        for k in options:
            self.option(k, to_str(options[k]))
        return self

    options.__doc__ = PySparkDataFrameReader.options.__doc__

    def load(
        self,
        path: Optional[PathOrPaths] = None,
        format: Optional[str] = None,
        schema: Optional[Union[StructType, str]] = None,
        **options: "OptionalPrimitiveType",
    ) -> "DataFrame":
        if format is not None:
            self.format(format)
        if schema is not None:
            self.schema(schema)
        self.options(**options)

        paths = path
        if isinstance(path, str):
            paths = [path]

        plan = DataSource(
            format=self._format,
            schema=self._schema,
            options=self._options,
            paths=paths,  # type: ignore[arg-type]
        )
        return self._df(plan)

    load.__doc__ = PySparkDataFrameReader.load.__doc__

    def _df(self, plan: LogicalPlan) -> "DataFrame":
        from pyspark.sql.connect.dataframe import DataFrame

        return DataFrame(plan, self._client)

    def table(self, tableName: str) -> "DataFrame":
        return self._df(Read(tableName, self._options))

    table.__doc__ = PySparkDataFrameReader.table.__doc__

    def json(
        self,
        path: PathOrPaths,
        schema: Optional[Union[StructType, str]] = None,
        primitivesAsString: Optional[Union[bool, str]] = None,
        prefersDecimal: Optional[Union[bool, str]] = None,
        allowComments: Optional[Union[bool, str]] = None,
        allowUnquotedFieldNames: Optional[Union[bool, str]] = None,
        allowSingleQuotes: Optional[Union[bool, str]] = None,
        allowNumericLeadingZero: Optional[Union[bool, str]] = None,
        allowBackslashEscapingAnyCharacter: Optional[Union[bool, str]] = None,
        mode: Optional[str] = None,
        columnNameOfCorruptRecord: Optional[str] = None,
        dateFormat: Optional[str] = None,
        timestampFormat: Optional[str] = None,
        multiLine: Optional[Union[bool, str]] = None,
        allowUnquotedControlChars: Optional[Union[bool, str]] = None,
        lineSep: Optional[str] = None,
        samplingRatio: Optional[Union[float, str]] = None,
        dropFieldIfAllNull: Optional[Union[bool, str]] = None,
        encoding: Optional[str] = None,
        locale: Optional[str] = None,
        pathGlobFilter: Optional[Union[bool, str]] = None,
        recursiveFileLookup: Optional[Union[bool, str]] = None,
        modifiedBefore: Optional[Union[bool, str]] = None,
        modifiedAfter: Optional[Union[bool, str]] = None,
        allowNonNumericNumbers: Optional[Union[bool, str]] = None,
        useUnsafeRow: Optional[Union[bool, str]] = None,
    ) -> "DataFrame":
        self._set_opts(
            primitivesAsString=primitivesAsString,
            prefersDecimal=prefersDecimal,
            allowComments=allowComments,
            allowUnquotedFieldNames=allowUnquotedFieldNames,
            allowSingleQuotes=allowSingleQuotes,
            allowNumericLeadingZero=allowNumericLeadingZero,
            allowBackslashEscapingAnyCharacter=allowBackslashEscapingAnyCharacter,
            mode=mode,
            columnNameOfCorruptRecord=columnNameOfCorruptRecord,
            dateFormat=dateFormat,
            timestampFormat=timestampFormat,
            multiLine=multiLine,
            allowUnquotedControlChars=allowUnquotedControlChars,
            lineSep=lineSep,
            samplingRatio=samplingRatio,
            dropFieldIfAllNull=dropFieldIfAllNull,
            encoding=encoding,
            locale=locale,
            pathGlobFilter=pathGlobFilter,
            recursiveFileLookup=recursiveFileLookup,
            modifiedBefore=modifiedBefore,
            modifiedAfter=modifiedAfter,
            allowNonNumericNumbers=allowNonNumericNumbers,
            useUnsafeRow=useUnsafeRow,
        )
        if isinstance(path, str):
            path = [path]
        return self.load(path=path, format="json", schema=schema)

    json.__doc__ = PySparkDataFrameReader.json.__doc__

    def parquet(self, *paths: str, **options: "OptionalPrimitiveType") -> "DataFrame":
        mergeSchema = options.get("mergeSchema", None)
        pathGlobFilter = options.get("pathGlobFilter", None)
        modifiedBefore = options.get("modifiedBefore", None)
        modifiedAfter = options.get("modifiedAfter", None)
        recursiveFileLookup = options.get("recursiveFileLookup", None)
        datetimeRebaseMode = options.get("datetimeRebaseMode", None)
        int96RebaseMode = options.get("int96RebaseMode", None)
        self._set_opts(
            mergeSchema=mergeSchema,
            pathGlobFilter=pathGlobFilter,
            recursiveFileLookup=recursiveFileLookup,
            modifiedBefore=modifiedBefore,
            modifiedAfter=modifiedAfter,
            datetimeRebaseMode=datetimeRebaseMode,
            int96RebaseMode=int96RebaseMode,
        )

        return self.load(path=list(paths), format="parquet")

    parquet.__doc__ = PySparkDataFrameReader.parquet.__doc__

    def text(
        self,
        paths: PathOrPaths,
        wholetext: bool = False,
        lineSep: Optional[str] = None,
        pathGlobFilter: Optional[Union[bool, str]] = None,
        recursiveFileLookup: Optional[Union[bool, str]] = None,
        modifiedBefore: Optional[Union[bool, str]] = None,
        modifiedAfter: Optional[Union[bool, str]] = None,
    ) -> "DataFrame":
        self._set_opts(
            wholetext=wholetext,
            lineSep=lineSep,
            pathGlobFilter=pathGlobFilter,
            recursiveFileLookup=recursiveFileLookup,
            modifiedBefore=modifiedBefore,
            modifiedAfter=modifiedAfter,
        )

        if isinstance(paths, str):
            paths = [paths]
        return self.load(path=paths, format="text")

    text.__doc__ = PySparkDataFrameReader.text.__doc__

    def csv(
        self,
        path: PathOrPaths,
        schema: Optional[Union[StructType, str]] = None,
        sep: Optional[str] = None,
        encoding: Optional[str] = None,
        quote: Optional[str] = None,
        escape: Optional[str] = None,
        comment: Optional[str] = None,
        header: Optional[Union[bool, str]] = None,
        inferSchema: Optional[Union[bool, str]] = None,
        ignoreLeadingWhiteSpace: Optional[Union[bool, str]] = None,
        ignoreTrailingWhiteSpace: Optional[Union[bool, str]] = None,
        nullValue: Optional[str] = None,
        nanValue: Optional[str] = None,
        positiveInf: Optional[str] = None,
        negativeInf: Optional[str] = None,
        dateFormat: Optional[str] = None,
        timestampFormat: Optional[str] = None,
        maxColumns: Optional[Union[int, str]] = None,
        maxCharsPerColumn: Optional[Union[int, str]] = None,
        maxMalformedLogPerPartition: Optional[Union[int, str]] = None,
        mode: Optional[str] = None,
        columnNameOfCorruptRecord: Optional[str] = None,
        multiLine: Optional[Union[bool, str]] = None,
        charToEscapeQuoteEscaping: Optional[str] = None,
        samplingRatio: Optional[Union[float, str]] = None,
        enforceSchema: Optional[Union[bool, str]] = None,
        emptyValue: Optional[str] = None,
        locale: Optional[str] = None,
        lineSep: Optional[str] = None,
        pathGlobFilter: Optional[Union[bool, str]] = None,
        recursiveFileLookup: Optional[Union[bool, str]] = None,
        modifiedBefore: Optional[Union[bool, str]] = None,
        modifiedAfter: Optional[Union[bool, str]] = None,
        unescapedQuoteHandling: Optional[str] = None,
    ) -> "DataFrame":
        self._set_opts(
            sep=sep,
            encoding=encoding,
            quote=quote,
            escape=escape,
            comment=comment,
            header=header,
            inferSchema=inferSchema,
            ignoreLeadingWhiteSpace=ignoreLeadingWhiteSpace,
            ignoreTrailingWhiteSpace=ignoreTrailingWhiteSpace,
            nullValue=nullValue,
            nanValue=nanValue,
            positiveInf=positiveInf,
            negativeInf=negativeInf,
            dateFormat=dateFormat,
            timestampFormat=timestampFormat,
            maxColumns=maxColumns,
            maxCharsPerColumn=maxCharsPerColumn,
            maxMalformedLogPerPartition=maxMalformedLogPerPartition,
            mode=mode,
            columnNameOfCorruptRecord=columnNameOfCorruptRecord,
            multiLine=multiLine,
            charToEscapeQuoteEscaping=charToEscapeQuoteEscaping,
            samplingRatio=samplingRatio,
            enforceSchema=enforceSchema,
            emptyValue=emptyValue,
            locale=locale,
            lineSep=lineSep,
            pathGlobFilter=pathGlobFilter,
            recursiveFileLookup=recursiveFileLookup,
            modifiedBefore=modifiedBefore,
            modifiedAfter=modifiedAfter,
            unescapedQuoteHandling=unescapedQuoteHandling,
        )
        if isinstance(path, str):
            path = [path]
        return self.load(path=path, format="csv", schema=schema)

    csv.__doc__ = PySparkDataFrameReader.csv.__doc__

    def xml(
        self,
        path: PathOrPaths,
        rowTag: Optional[str] = None,
        schema: Optional[Union[StructType, str]] = None,
        excludeAttribute: Optional[Union[bool, str]] = None,
        attributePrefix: Optional[str] = None,
        valueTag: Optional[str] = None,
        ignoreSurroundingSpaces: Optional[Union[bool, str]] = None,
        rowValidationXSDPath: Optional[str] = None,
        ignoreNamespace: Optional[Union[bool, str]] = None,
        wildcardColName: Optional[str] = None,
        encoding: Optional[str] = None,
        inferSchema: Optional[Union[bool, str]] = None,
        nullValue: Optional[str] = None,
        dateFormat: Optional[str] = None,
        timestampFormat: Optional[str] = None,
        mode: Optional[str] = None,
        columnNameOfCorruptRecord: Optional[str] = None,
        multiLine: Optional[Union[bool, str]] = None,
        samplingRatio: Optional[Union[float, str]] = None,
        locale: Optional[str] = None,
    ) -> "DataFrame":
        self._set_opts(
            rowTag=rowTag,
            schema=schema,
            excludeAttribute=excludeAttribute,
            attributePrefix=attributePrefix,
            valueTag=valueTag,
            ignoreSurroundingSpaces=ignoreSurroundingSpaces,
            rowValidationXSDPath=rowValidationXSDPath,
            ignoreNamespace=ignoreNamespace,
            wildcardColName=wildcardColName,
            encoding=encoding,
            inferSchema=inferSchema,
            nullValue=nullValue,
            dateFormat=dateFormat,
            timestampFormat=timestampFormat,
            mode=mode,
            columnNameOfCorruptRecord=columnNameOfCorruptRecord,
            multiLine=multiLine,
            samplingRatio=samplingRatio,
            locale=locale,
        )
        if isinstance(path, str):
            path = [path]
        return self.load(path=path, format="xml", schema=schema)

    xml.__doc__ = PySparkDataFrameReader.xml.__doc__

    def orc(
        self,
        path: PathOrPaths,
        mergeSchema: Optional[bool] = None,
        pathGlobFilter: Optional[Union[bool, str]] = None,
        recursiveFileLookup: Optional[Union[bool, str]] = None,
        modifiedBefore: Optional[Union[bool, str]] = None,
        modifiedAfter: Optional[Union[bool, str]] = None,
    ) -> "DataFrame":
        self._set_opts(
            mergeSchema=mergeSchema,
            pathGlobFilter=pathGlobFilter,
            modifiedBefore=modifiedBefore,
            modifiedAfter=modifiedAfter,
            recursiveFileLookup=recursiveFileLookup,
        )
        if isinstance(path, str):
            path = [path]
        return self.load(path=path, format="orc")

    orc.__doc__ = PySparkDataFrameReader.orc.__doc__

    @overload
    def jdbc(
        self, url: str, table: str, *, properties: Optional[Dict[str, str]] = None
    ) -> "DataFrame":
        ...

    @overload
    def jdbc(
        self,
        url: str,
        table: str,
        column: str,
        lowerBound: Union[int, str],
        upperBound: Union[int, str],
        numPartitions: int,
        *,
        properties: Optional[Dict[str, str]] = None,
    ) -> "DataFrame":
        ...

    @overload
    def jdbc(
        self,
        url: str,
        table: str,
        *,
        predicates: List[str],
        properties: Optional[Dict[str, str]] = None,
    ) -> "DataFrame":
        ...

    def jdbc(
        self,
        url: str,
        table: str,
        column: Optional[str] = None,
        lowerBound: Optional[Union[int, str]] = None,
        upperBound: Optional[Union[int, str]] = None,
        numPartitions: Optional[int] = None,
        predicates: Optional[List[str]] = None,
        properties: Optional[Dict[str, str]] = None,
    ) -> "DataFrame":
        if properties is None:
            properties = dict()

        self.format("jdbc")

        if column is not None:
            assert lowerBound is not None, "lowerBound can not be None when ``column`` is specified"
            assert upperBound is not None, "upperBound can not be None when ``column`` is specified"
            assert (
                numPartitions is not None
            ), "numPartitions can not be None when ``column`` is specified"
            self.options(
                partitionColumn=column,
                lowerBound=lowerBound,
                upperBound=upperBound,
                numPartitions=numPartitions,
            )
            self.options(**properties)
            self.options(url=url, dbtable=table)
            return self.load()
        else:
            self.options(**properties)
            self.options(url=url, dbtable=table)
            if predicates is not None:
                plan = DataSource(
                    format=self._format,
                    schema=self._schema,
                    options=self._options,
                    predicates=predicates,
                )
                return self._df(plan)
            else:
                return self.load()

    jdbc.__doc__ = PySparkDataFrameReader.jdbc.__doc__

    @property
    def _jreader(self) -> None:
        raise PySparkAttributeError(
            errorClass="JVM_ATTRIBUTE_NOT_SUPPORTED", messageParameters={"attr_name": "_jreader"}
        )


DataFrameReader.__doc__ = PySparkDataFrameReader.__doc__


class DataFrameWriter(OptionUtils):
    def __init__(
        self,
        plan: "LogicalPlan",
        session: "SparkSession",
        callback: Optional[Callable[["ExecutionInfo"], None]] = None,
    ):
        self._df: "LogicalPlan" = plan
        self._spark: "SparkSession" = session
        self._write: "WriteOperation" = WriteOperation(self._df)

        self._callback = callback if callback is not None else lambda _: None

    def mode(self, saveMode: Optional[str]) -> "DataFrameWriter":
        # At the JVM side, the default value of mode is already set to "error".
        # So, if the given saveMode is None, we will not call JVM-side's mode method.
        if saveMode is not None:
            self._write.mode = saveMode
        return self

    mode.__doc__ = PySparkDataFrameWriter.mode.__doc__

    def format(self, source: str) -> "DataFrameWriter":
        self._write.source = source
        return self

    format.__doc__ = PySparkDataFrameWriter.format.__doc__

    def option(self, key: str, value: "OptionalPrimitiveType") -> "DataFrameWriter":
        self._write.options[key] = to_str(value)
        return self

    option.__doc__ = PySparkDataFrameWriter.option.__doc__

    def options(self, **options: "OptionalPrimitiveType") -> "DataFrameWriter":
        for k in options:
            self._write.options[k] = to_str(options[k])
        return self

    options.__doc__ = PySparkDataFrameWriter.options.__doc__

    @overload
    def partitionBy(self, *cols: str) -> "DataFrameWriter":
        ...

    @overload
    def partitionBy(self, *cols: List[str]) -> "DataFrameWriter":
        ...

    def partitionBy(self, *cols: Union[str, List[str]]) -> "DataFrameWriter":
        if len(cols) == 1 and isinstance(cols[0], (list, tuple)):
            cols = cols[0]  # type: ignore[assignment]

        self._write.partitioning_cols = cast(List[str], cols)
        return self

    partitionBy.__doc__ = PySparkDataFrameWriter.partitionBy.__doc__

    @overload
    def bucketBy(self, numBuckets: int, col: str, *cols: str) -> "DataFrameWriter":
        ...

    @overload
    def bucketBy(self, numBuckets: int, col: TupleOrListOfString) -> "DataFrameWriter":
        ...

    def bucketBy(
        self, numBuckets: int, col: Union[str, TupleOrListOfString], *cols: Optional[str]
    ) -> "DataFrameWriter":
        if not isinstance(numBuckets, int):
            raise PySparkValueError(
                errorClass="CANNOT_SET_TOGETHER",
                messageParameters={
                    "arg_list": f"`col` of type {type(col).__name__} and `cols`",
                },
            )

        if isinstance(col, (list, tuple)):
            if cols:
                raise PySparkValueError(
                    errorClass="NOT_INT",
                    messageParameters={
                        "arg_list": "numBuckets",
                    },
                )

            col, cols = col[0], col[1:]  # type: ignore[assignment]

        for c in cols:
            if not isinstance(c, str):
                raise PySparkTypeError(
                    errorClass="NOT_LIST_OF_STR",
                    messageParameters={
                        "arg_name": "cols",
                        "arg_type": type(c).__name__,
                    },
                )
        if not isinstance(col, str):
            raise PySparkTypeError(
                errorClass="NOT_LIST_OF_STR",
                messageParameters={
                    "arg_name": "col",
                    "arg_type": type(col).__name__,
                },
            )

        self._write.num_buckets = numBuckets
        self._write.bucket_cols = cast(List[str], [col, *cols])
        return self

    bucketBy.__doc__ = PySparkDataFrameWriter.bucketBy.__doc__

    @overload
    def sortBy(self, col: str, *cols: str) -> "DataFrameWriter":
        ...

    @overload
    def sortBy(self, col: TupleOrListOfString) -> "DataFrameWriter":
        ...

    def sortBy(
        self, col: Union[str, TupleOrListOfString], *cols: Optional[str]
    ) -> "DataFrameWriter":
        if isinstance(col, (list, tuple)):
            if cols:
                raise PySparkValueError(
                    errorClass="CANNOT_SET_TOGETHER",
                    messageParameters={
                        "arg_list": f"`col` of type {type(col).__name__} and `cols`",
                    },
                )

            col, cols = col[0], col[1:]  # type: ignore[assignment]

        for c in cols:
            if not isinstance(c, str):
                raise PySparkTypeError(
                    errorClass="NOT_LIST_OF_STR",
                    messageParameters={
                        "arg_name": "cols",
                        "arg_type": type(c).__name__,
                    },
                )
        if not isinstance(col, str):
            raise PySparkTypeError(
                errorClass="NOT_LIST_OF_STR",
                messageParameters={
                    "arg_name": "col",
                    "arg_type": type(col).__name__,
                },
            )

        self._write.sort_cols = cast(List[str], [col, *cols])
        return self

    sortBy.__doc__ = PySparkDataFrameWriter.sortBy.__doc__

    @overload
    def clusterBy(self, *cols: str) -> "DataFrameWriter":
        ...

    @overload
    def clusterBy(self, *cols: List[str]) -> "DataFrameWriter":
        ...

    def clusterBy(self, *cols: Union[str, List[str]]) -> "DataFrameWriter":
        if len(cols) == 1 and isinstance(cols[0], (list, tuple)):
            cols = cols[0]  # type: ignore[assignment]
        assert len(cols) > 0, "clusterBy needs one or more clustering columns."
        self._write.clustering_cols = cast(List[str], cols)
        return self

    clusterBy.__doc__ = PySparkDataFrameWriter.clusterBy.__doc__

    def save(
        self,
        path: Optional[str] = None,
        format: Optional[str] = None,
        mode: Optional[str] = None,
        partitionBy: Optional[Union[str, List[str]]] = None,
        **options: "OptionalPrimitiveType",
    ) -> None:
        self.mode(mode).options(**options)
        if partitionBy is not None:
            self.partitionBy(partitionBy)
        if format is not None:
            self.format(format)
        self._write.path = path
        _, _, ei = self._spark.client.execute_command(
            self._write.command(self._spark.client), self._write.observations
        )
        self._callback(ei)

    save.__doc__ = PySparkDataFrameWriter.save.__doc__

    def insertInto(self, tableName: str, overwrite: Optional[bool] = None) -> None:
        if overwrite is not None:
            self.mode("overwrite" if overwrite else "append")
        self._write.table_name = tableName
        self._write.table_save_method = "insert_into"
        _, _, ei = self._spark.client.execute_command(
            self._write.command(self._spark.client), self._write.observations
        )
        self._callback(ei)

    insertInto.__doc__ = PySparkDataFrameWriter.insertInto.__doc__

    def saveAsTable(
        self,
        name: str,
        format: Optional[str] = None,
        mode: Optional[str] = None,
        partitionBy: Optional[Union[str, List[str]]] = None,
        **options: "OptionalPrimitiveType",
    ) -> None:
        self.mode(mode).options(**options)
        if partitionBy is not None:
            self.partitionBy(partitionBy)
        if format is not None:
            self.format(format)
        self._write.table_name = name
        self._write.table_save_method = "save_as_table"
        _, _, ei = self._spark.client.execute_command(
            self._write.command(self._spark.client), self._write.observations
        )
        self._callback(ei)

    saveAsTable.__doc__ = PySparkDataFrameWriter.saveAsTable.__doc__

    def json(
        self,
        path: str,
        mode: Optional[str] = None,
        compression: Optional[str] = None,
        dateFormat: Optional[str] = None,
        timestampFormat: Optional[str] = None,
        lineSep: Optional[str] = None,
        encoding: Optional[str] = None,
        ignoreNullFields: Optional[Union[bool, str]] = None,
    ) -> None:
        self.mode(mode)
        self._set_opts(
            compression=compression,
            dateFormat=dateFormat,
            timestampFormat=timestampFormat,
            lineSep=lineSep,
            encoding=encoding,
            ignoreNullFields=ignoreNullFields,
        )
        self.format("json").save(path)

    json.__doc__ = PySparkDataFrameWriter.json.__doc__

    def parquet(
        self,
        path: str,
        mode: Optional[str] = None,
        partitionBy: Optional[Union[str, List[str]]] = None,
        compression: Optional[str] = None,
    ) -> None:
        self.mode(mode)
        if partitionBy is not None:
            self.partitionBy(partitionBy)
        self._set_opts(compression=compression)
        self.format("parquet").save(path)

    parquet.__doc__ = PySparkDataFrameWriter.parquet.__doc__

    def text(
        self, path: str, compression: Optional[str] = None, lineSep: Optional[str] = None
    ) -> None:
        self._set_opts(compression=compression, lineSep=lineSep)
        self.format("text").save(path)

    text.__doc__ = PySparkDataFrameWriter.text.__doc__

    def csv(
        self,
        path: str,
        mode: Optional[str] = None,
        compression: Optional[str] = None,
        sep: Optional[str] = None,
        quote: Optional[str] = None,
        escape: Optional[str] = None,
        header: Optional[Union[bool, str]] = None,
        nullValue: Optional[str] = None,
        escapeQuotes: Optional[Union[bool, str]] = None,
        quoteAll: Optional[Union[bool, str]] = None,
        dateFormat: Optional[str] = None,
        timestampFormat: Optional[str] = None,
        ignoreLeadingWhiteSpace: Optional[Union[bool, str]] = None,
        ignoreTrailingWhiteSpace: Optional[Union[bool, str]] = None,
        charToEscapeQuoteEscaping: Optional[str] = None,
        encoding: Optional[str] = None,
        emptyValue: Optional[str] = None,
        lineSep: Optional[str] = None,
    ) -> None:
        self.mode(mode)
        self._set_opts(
            compression=compression,
            sep=sep,
            quote=quote,
            escape=escape,
            header=header,
            nullValue=nullValue,
            escapeQuotes=escapeQuotes,
            quoteAll=quoteAll,
            dateFormat=dateFormat,
            timestampFormat=timestampFormat,
            ignoreLeadingWhiteSpace=ignoreLeadingWhiteSpace,
            ignoreTrailingWhiteSpace=ignoreTrailingWhiteSpace,
            charToEscapeQuoteEscaping=charToEscapeQuoteEscaping,
            encoding=encoding,
            emptyValue=emptyValue,
            lineSep=lineSep,
        )
        self.format("csv").save(path)

    csv.__doc__ = PySparkDataFrameWriter.csv.__doc__

    def xml(
        self,
        path: str,
        rowTag: Optional[str] = None,
        mode: Optional[str] = None,
        attributePrefix: Optional[str] = None,
        valueTag: Optional[str] = None,
        rootTag: Optional[str] = None,
        declaration: Optional[bool] = None,
        arrayElementName: Optional[str] = None,
        nullValue: Optional[str] = None,
        dateFormat: Optional[str] = None,
        timestampFormat: Optional[str] = None,
        compression: Optional[str] = None,
        encoding: Optional[str] = None,
        validateName: Optional[bool] = None,
    ) -> None:
        self.mode(mode)
        self._set_opts(
            rowTag=rowTag,
            attributePrefix=attributePrefix,
            valueTag=valueTag,
            rootTag=rootTag,
            declaration=declaration,
            arrayElementName=arrayElementName,
            nullValue=nullValue,
            dateFormat=dateFormat,
            timestampFormat=timestampFormat,
            compression=compression,
            encoding=encoding,
            validateName=validateName,
        )
        self.format("xml").save(path)

    xml.__doc__ = PySparkDataFrameWriter.xml.__doc__

    def orc(
        self,
        path: str,
        mode: Optional[str] = None,
        partitionBy: Optional[Union[str, List[str]]] = None,
        compression: Optional[str] = None,
    ) -> None:
        self.mode(mode)
        if partitionBy is not None:
            self.partitionBy(partitionBy)
        self._set_opts(compression=compression)
        self.format("orc").save(path)

    orc.__doc__ = PySparkDataFrameWriter.orc.__doc__

    def jdbc(
        self,
        url: str,
        table: str,
        mode: Optional[str] = None,
        properties: Optional[Dict[str, str]] = None,
    ) -> None:
        if properties is None:
            properties = dict()

        self.format("jdbc").mode(mode).options(**properties).options(url=url, dbtable=table).save()

    jdbc.__doc__ = PySparkDataFrameWriter.jdbc.__doc__


class DataFrameWriterV2(OptionUtils):
    def __init__(
        self,
        plan: "LogicalPlan",
        session: "SparkSession",
        table: str,
        callback: Optional[Callable[["ExecutionInfo"], None]] = None,
    ):
        self._df: "LogicalPlan" = plan
        self._spark: "SparkSession" = session
        self._table_name: str = table
        self._write: "WriteOperationV2" = WriteOperationV2(self._df, self._table_name)
        self._callback = callback if callback is not None else lambda _: None

    def using(self, provider: str) -> "DataFrameWriterV2":
        self._write.provider = provider
        return self

    using.__doc__ = PySparkDataFrameWriterV2.using.__doc__

    def option(self, key: str, value: "OptionalPrimitiveType") -> "DataFrameWriterV2":
        self._write.options[key] = to_str(value)
        return self

    option.__doc__ = PySparkDataFrameWriterV2.option.__doc__

    def options(self, **options: "OptionalPrimitiveType") -> "DataFrameWriterV2":
        for k in options:
            self._write.options[k] = to_str(options[k])
        return self

    options.__doc__ = PySparkDataFrameWriterV2.options.__doc__

    def tableProperty(self, property: str, value: str) -> "DataFrameWriterV2":
        self._write.table_properties[property] = value
        return self

    tableProperty.__doc__ = PySparkDataFrameWriterV2.tableProperty.__doc__

    def partitionedBy(self, col: "ColumnOrName", *cols: "ColumnOrName") -> "DataFrameWriterV2":
        self._write.partitioning_columns = [F._to_col(c) for c in [col] + list(cols)]
        return self

    partitionedBy.__doc__ = PySparkDataFrameWriterV2.partitionedBy.__doc__

    def clusterBy(self, col: str, *cols: str) -> "DataFrameWriterV2":
        self._write.clustering_columns = [col]
        self._write.clustering_columns.extend(cols)
        return self

    clusterBy.__doc__ = PySparkDataFrameWriterV2.clusterBy.__doc__

    def create(self) -> None:
        self._write.mode = "create"
        _, _, ei = self._spark.client.execute_command(
            self._write.command(self._spark.client), self._write.observations
        )
        self._callback(ei)

    create.__doc__ = PySparkDataFrameWriterV2.create.__doc__

    def replace(self) -> None:
        self._write.mode = "replace"
        _, _, ei = self._spark.client.execute_command(
            self._write.command(self._spark.client), self._write.observations
        )
        self._callback(ei)

    replace.__doc__ = PySparkDataFrameWriterV2.replace.__doc__

    def createOrReplace(self) -> None:
        self._write.mode = "create_or_replace"
        _, _, ei = self._spark.client.execute_command(
            self._write.command(self._spark.client), self._write.observations
        )
        self._callback(ei)

    createOrReplace.__doc__ = PySparkDataFrameWriterV2.createOrReplace.__doc__

    def append(self) -> None:
        self._write.mode = "append"
        _, _, ei = self._spark.client.execute_command(
            self._write.command(self._spark.client), self._write.observations
        )
        self._callback(ei)

    append.__doc__ = PySparkDataFrameWriterV2.append.__doc__

    def overwrite(self, condition: "ColumnOrName") -> None:
        self._write.mode = "overwrite"
        self._write.overwrite_condition = F._to_col(condition)
        _, _, ei = self._spark.client.execute_command(
            self._write.command(self._spark.client), self._write.observations
        )
        self._callback(ei)

    overwrite.__doc__ = PySparkDataFrameWriterV2.overwrite.__doc__

    def overwritePartitions(self) -> None:
        self._write.mode = "overwrite_partitions"
        _, _, ei = self._spark.client.execute_command(
            self._write.command(self._spark.client), self._write.observations
        )
        self._callback(ei)

    overwritePartitions.__doc__ = PySparkDataFrameWriterV2.overwritePartitions.__doc__


def _test() -> None:
    import sys
    import os
    import doctest
    from pyspark.sql import SparkSession as PySparkSession
    import pyspark.sql.connect.readwriter

    globs = pyspark.sql.connect.readwriter.__dict__.copy()

    globs["spark"] = (
        PySparkSession.builder.appName("sql.connect.readwriter tests")
        .remote(os.environ.get("SPARK_CONNECT_TESTING_REMOTE", "local[4]"))
        .getOrCreate()
    )

    (failure_count, test_count) = doctest.testmod(
        pyspark.sql.connect.readwriter,
        globs=globs,
        optionflags=doctest.ELLIPSIS
        | doctest.NORMALIZE_WHITESPACE
        | doctest.IGNORE_EXCEPTION_DETAIL,
    )

    globs["spark"].stop()

    if failure_count:
        sys.exit(-1)


if __name__ == "__main__":
    _test()