#
# 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.
#

import os
import shutil
import tempfile
import time

from pyspark.sql import Row
from pyspark.sql.functions import lit
from pyspark.sql.types import StructType, StructField, IntegerType, StringType, TimestampType
from pyspark.testing.sqlutils import ReusedSQLTestCase
from pyspark.errors import PySparkValueError


class StreamingTestsMixin:
    def test_streaming_query_functions_basic(self):
        df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
        query = (
            df.writeStream.format("memory")
            .queryName("test_streaming_query_functions_basic")
            .start()
        )
        try:
            self.assertEqual(query.name, "test_streaming_query_functions_basic")
            self.assertTrue(isinstance(query.id, str))
            self.assertTrue(isinstance(query.runId, str))
            self.assertTrue(query.isActive)
            self.assertEqual(query.exception(), None)
            self.assertFalse(query.awaitTermination(1))
            query.processAllAvailable()
            lastProgress = query.lastProgress
            recentProgress = query.recentProgress
            self.assertEqual(lastProgress["name"], query.name)
            self.assertEqual(lastProgress["id"], query.id)
            self.assertTrue(any(p == lastProgress for p in recentProgress))
            query.explain()

        except Exception as e:
            self.fail(
                "Streaming query functions sanity check shouldn't throw any error. "
                "Error message: " + str(e)
            )

        finally:
            query.stop()

    def test_streaming_progress(self):
        """
        Should be able to access fields using attributes in lastProgress / recentProgress
        e.g. q.lastProgress.id
        """
        df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
        query = df.writeStream.format("noop").start()
        try:
            query.processAllAvailable()
            lastProgress = query.lastProgress
            recentProgress = query.recentProgress
            self.assertEqual(lastProgress["name"], query.name)
            # Return str when accessed using dict get.
            self.assertEqual(lastProgress["id"], query.id)
            # SPARK-48567 Use attribute to access fields in q.lastProgress
            self.assertEqual(lastProgress.name, query.name)
            # Return uuid when accessed using attribute.
            self.assertEqual(str(lastProgress.id), query.id)
            self.assertTrue(any(p == lastProgress for p in recentProgress))
            self.assertTrue(lastProgress.numInputRows > 0)
            # Also access source / sink progress with attributes
            self.assertTrue(len(lastProgress.sources) > 0)
            self.assertTrue(lastProgress.sources[0].numInputRows > 0)
            self.assertTrue(lastProgress["sources"][0]["numInputRows"] > 0)
            self.assertTrue(lastProgress.sink.numOutputRows > 0)
            self.assertTrue(lastProgress["sink"]["numOutputRows"] > 0)
            # In Python, for historical reasons, changing field value
            # in StreamingQueryProgress is allowed.
            new_name = "myNewQuery"
            lastProgress["name"] = new_name
            self.assertEqual(lastProgress.name, new_name)

        except Exception as e:
            self.fail(
                "Streaming query functions sanity check shouldn't throw any error. "
                "Error message: " + str(e)
            )
        finally:
            query.stop()

    def test_streaming_query_name_edge_case(self):
        # Query name should be None when not specified
        q1 = self.spark.readStream.format("rate").load().writeStream.format("noop").start()
        self.assertEqual(q1.name, None)

        # Cannot set query name to be an empty string
        error_thrown = False
        try:
            (
                self.spark.readStream.format("rate")
                .load()
                .writeStream.format("noop")
                .queryName("")
                .start()
            )
        except PySparkValueError:
            error_thrown = True

        self.assertTrue(error_thrown)

    def test_stream_trigger(self):
        df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")

        # Should take at least one arg
        try:
            df.writeStream.trigger()
        except ValueError:
            pass

        # Should not take multiple args
        try:
            df.writeStream.trigger(once=True, processingTime="5 seconds")
        except ValueError:
            pass

        # Should not take multiple args
        try:
            df.writeStream.trigger(processingTime="5 seconds", continuous="1 second")
        except ValueError:
            pass

        # Should take only keyword args
        try:
            df.writeStream.trigger("5 seconds")
            self.fail("Should have thrown an exception")
        except TypeError:
            pass

    def test_stream_read_options(self):
        schema = StructType([StructField("data", StringType(), False)])
        df = (
            self.spark.readStream.format("text")
            .option("path", "python/test_support/sql/streaming")
            .schema(schema)
            .load()
        )
        self.assertTrue(df.isStreaming)
        self.assertEqual(df.schema.simpleString(), "struct<data:string>")

    def test_stream_read_options_overwrite(self):
        bad_schema = StructType([StructField("test", IntegerType(), False)])
        schema = StructType([StructField("data", StringType(), False)])
        # SPARK-32516 disables the overwrite behavior by default.
        with self.sql_conf({"spark.sql.legacy.pathOptionBehavior.enabled": True}):
            df = (
                self.spark.readStream.format("csv")
                .option("path", "python/test_support/sql/fake")
                .schema(bad_schema)
                .load(path="python/test_support/sql/streaming", schema=schema, format="text")
            )
            self.assertTrue(df.isStreaming)
            self.assertEqual(df.schema.simpleString(), "struct<data:string>")

    def test_stream_save_options(self):
        df = (
            self.spark.readStream.format("text")
            .load("python/test_support/sql/streaming")
            .withColumn("id", lit(1))
        )
        for q in self.spark.streams.active:
            q.stop()
        tmpPath = tempfile.mkdtemp()
        shutil.rmtree(tmpPath)
        self.assertTrue(df.isStreaming)
        out = os.path.join(tmpPath, "out")
        chk = os.path.join(tmpPath, "chk")
        q = (
            df.writeStream.option("checkpointLocation", chk)
            .queryName("this_query")
            .format("parquet")
            .partitionBy("id")
            .outputMode("append")
            .option("path", out)
            .start()
        )
        try:
            self.assertEqual(q.name, "this_query")
            self.assertTrue(q.isActive)
            q.processAllAvailable()
            output_files = []
            for _, _, files in os.walk(out):
                output_files.extend([f for f in files if not f.startswith(".")])
            self.assertTrue(len(output_files) > 0)
            self.assertTrue(len(os.listdir(chk)) > 0)
        finally:
            q.stop()
            shutil.rmtree(tmpPath)

    def test_stream_save_options_overwrite(self):
        df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
        for q in self.spark.streams.active:
            q.stop()
        tmpPath = tempfile.mkdtemp()
        shutil.rmtree(tmpPath)
        self.assertTrue(df.isStreaming)
        out = os.path.join(tmpPath, "out")
        chk = os.path.join(tmpPath, "chk")
        fake1 = os.path.join(tmpPath, "fake1")
        fake2 = os.path.join(tmpPath, "fake2")
        # SPARK-32516 disables the overwrite behavior by default.
        with self.sql_conf({"spark.sql.legacy.pathOptionBehavior.enabled": True}):
            q = (
                df.writeStream.option("checkpointLocation", fake1)
                .format("memory")
                .option("path", fake2)
                .queryName("fake_query")
                .outputMode("append")
                .start(path=out, format="parquet", queryName="this_query", checkpointLocation=chk)
            )

        try:
            self.assertEqual(q.name, "this_query")
            self.assertTrue(q.isActive)
            q.processAllAvailable()
            output_files = []
            for _, _, files in os.walk(out):
                output_files.extend([f for f in files if not f.startswith(".")])
            self.assertTrue(len(output_files) > 0)
            self.assertTrue(len(os.listdir(chk)) > 0)
            self.assertFalse(os.path.isdir(fake1))  # should not have been created
            self.assertFalse(os.path.isdir(fake2))  # should not have been created
        finally:
            q.stop()
            shutil.rmtree(tmpPath)

    def test_stream_status_and_progress(self):
        df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
        for q in self.spark.streams.active:
            q.stop()
        tmpPath = tempfile.mkdtemp()
        shutil.rmtree(tmpPath)
        self.assertTrue(df.isStreaming)
        out = os.path.join(tmpPath, "out")
        chk = os.path.join(tmpPath, "chk")

        def func(x):
            time.sleep(1)
            return x

        from pyspark.sql.functions import col, udf

        sleep_udf = udf(func)

        # Use "sleep_udf" to delay the progress update so that we can test `lastProgress` when there
        # were no updates.
        q = df.select(sleep_udf(col("value")).alias("value")).writeStream.start(
            path=out, format="parquet", queryName="this_query", checkpointLocation=chk
        )
        try:
            # "lastProgress" will return None in most cases. However, as it may be flaky when
            # Jenkins is very slow, we don't assert it. If there is something wrong, "lastProgress"
            # may throw error with a high chance and make this test flaky, so we should still be
            # able to detect broken codes.
            q.lastProgress

            q.processAllAvailable()
            lastProgress = q.lastProgress
            recentProgress = q.recentProgress
            status = q.status
            self.assertEqual(lastProgress["name"], q.name)
            self.assertEqual(lastProgress["id"], q.id)
            self.assertTrue(any(p == lastProgress for p in recentProgress))
            self.assertTrue(
                "message" in status and "isDataAvailable" in status and "isTriggerActive" in status
            )
        finally:
            q.stop()
            shutil.rmtree(tmpPath)

    def test_stream_await_termination(self):
        df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
        for q in self.spark.streams.active:
            q.stop()
        tmpPath = tempfile.mkdtemp()
        shutil.rmtree(tmpPath)
        self.assertTrue(df.isStreaming)
        out = os.path.join(tmpPath, "out")
        chk = os.path.join(tmpPath, "chk")
        q = df.writeStream.start(
            path=out, format="parquet", queryName="this_query", checkpointLocation=chk
        )
        try:
            self.assertTrue(q.isActive)
            try:
                q.awaitTermination("hello")
                self.fail("Expected a value exception")
            except ValueError:
                pass
            now = time.time()
            # test should take at least 2 seconds
            res = q.awaitTermination(2.6)
            duration = time.time() - now
            self.assertTrue(duration >= 2)
            self.assertFalse(res)

            q.processAllAvailable()
            q.stop()
            # Sanity check when no parameter is set
            q.awaitTermination()
            self.assertFalse(q.isActive)
        finally:
            q.stop()
            shutil.rmtree(tmpPath)

    def test_stream_exception(self):
        with self.sql_conf({"spark.sql.execution.pyspark.udf.simplifiedTraceback.enabled": True}):
            sdf = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
            sq = sdf.writeStream.format("memory").queryName("query_explain").start()
            try:
                sq.processAllAvailable()
                self.assertEqual(sq.exception(), None)
            finally:
                sq.stop()

            from pyspark.sql.functions import col, udf
            from pyspark.errors import StreamingQueryException

            bad_udf = udf(lambda x: 1 / 0)
            sq = (
                sdf.select(bad_udf(col("value")))
                .writeStream.format("memory")
                .queryName("this_query")
                .start()
            )
            try:
                # Process some data to fail the query
                sq.processAllAvailable()
                self.fail("bad udf should fail the query")
            except StreamingQueryException as e:
                # This is expected
                self._assert_exception_tree_contains_msg(e, "ZeroDivisionError")
            finally:
                exception = sq.exception()
                sq.stop()
            self.assertIsInstance(exception, StreamingQueryException)
            self._assert_exception_tree_contains_msg(exception, "ZeroDivisionError")

    def test_query_manager_no_recreation(self):
        # SPARK-46873: There should not be a new StreamingQueryManager created every time
        # spark.streams is called.
        for i in range(5):
            self.assertTrue(self.spark.streams == self.spark.streams)

    def test_query_manager_get(self):
        df = self.spark.readStream.format("rate").load()
        for q in self.spark.streams.active:
            q.stop()
        q = df.writeStream.format("noop").start()

        self.assertTrue(q.isActive)
        self.assertTrue(q.id == self.spark.streams.get(q.id).id)

        q.stop()

        self.assertIsNone(self.spark.streams.get(q.id))

    def test_query_manager_await_termination(self):
        df = self.spark.readStream.format("text").load("python/test_support/sql/streaming")
        for q in self.spark.streams.active:
            q.stop()
        tmpPath = tempfile.mkdtemp()
        shutil.rmtree(tmpPath)
        self.assertTrue(df.isStreaming)
        out = os.path.join(tmpPath, "out")
        chk = os.path.join(tmpPath, "chk")
        q = df.writeStream.start(
            path=out, format="parquet", queryName="this_query", checkpointLocation=chk
        )
        try:
            self.assertTrue(q.isActive)
            try:
                self.spark.streams.awaitAnyTermination("hello")
                self.fail("Expected a value exception")
            except ValueError:
                pass
            now = time.time()
            # test should take at least 2 seconds
            res = self.spark.streams.awaitAnyTermination(2.6)
            duration = time.time() - now
            self.assertTrue(duration >= 2)
            self.assertFalse(res)
        finally:
            q.processAllAvailable()
            q.stop()
            shutil.rmtree(tmpPath)

    def test_streaming_read_from_table(self):
        with self.table("input_table", "this_query"):
            self.spark.sql("CREATE TABLE input_table (value string) USING parquet")
            self.spark.sql("INSERT INTO input_table VALUES ('aaa'), ('bbb'), ('ccc')")
            df = self.spark.readStream.table("input_table")
            self.assertTrue(df.isStreaming)
            q = df.writeStream.format("memory").queryName("this_query").start()
            q.processAllAvailable()
            q.stop()
            result = self.spark.sql("SELECT * FROM this_query ORDER BY value").collect()
            self.assertEqual(
                set([Row(value="aaa"), Row(value="bbb"), Row(value="ccc")]), set(result)
            )

    def test_streaming_write_to_table(self):
        with self.table("output_table"), tempfile.TemporaryDirectory(prefix="to_table") as tmpdir:
            df = self.spark.readStream.format("rate").option("rowsPerSecond", 10).load()
            q = df.writeStream.toTable("output_table", format="parquet", checkpointLocation=tmpdir)
            self.assertTrue(q.isActive)
            time.sleep(10)
            q.stop()
            result = self.spark.sql("SELECT value FROM output_table").collect()
            self.assertTrue(len(result) > 0)

    def test_streaming_write_to_table_cluster_by(self):
        with self.table("output_table"), tempfile.TemporaryDirectory(prefix="to_table") as tmpdir:
            df = self.spark.readStream.format("rate").option("rowsPerSecond", 10).load()
            q = df.writeStream.clusterBy("value").toTable(
                "output_table", format="parquet", checkpointLocation=tmpdir
            )
            self.assertTrue(q.isActive)
            time.sleep(10)
            q.stop()
            result = self.spark.sql("DESCRIBE output_table").collect()
            self.assertEqual(
                set(
                    [
                        Row(col_name="timestamp", data_type="timestamp", comment=None),
                        Row(col_name="value", data_type="bigint", comment=None),
                        Row(col_name="# Clustering Information", data_type="", comment=""),
                        Row(col_name="# col_name", data_type="data_type", comment="comment"),
                        Row(col_name="value", data_type="bigint", comment=None),
                    ]
                ),
                set(result),
            )
            result = self.spark.sql("SELECT value FROM output_table").collect()
            self.assertTrue(len(result) > 0)

    def test_streaming_with_temporary_view(self):
        """
        This verifies createOrReplaceTempView() works with a streaming dataframe. An SQL
        SELECT query on such a table results in a streaming dataframe and the streaming query works
        as expected.
        """
        with self.table("input_table", "this_query"):
            self.spark.sql("CREATE TABLE input_table (value string) USING parquet")
            self.spark.sql("INSERT INTO input_table VALUES ('a'), ('b'), ('c')")
            df = self.spark.readStream.table("input_table")
            self.assertTrue(df.isStreaming)
            # Create a temp view
            df.createOrReplaceTempView("test_view")
            # Create a select query
            view_df = self.spark.sql("SELECT CONCAT('view_', value) as vv from test_view")
            self.assertTrue(view_df.isStreaming)
            q = view_df.writeStream.format("memory").queryName("this_query").start()
            q.processAllAvailable()
            q.stop()
            result = self.spark.sql("SELECT * FROM this_query ORDER BY vv").collect()
            self.assertEqual(
                set([Row(value="view_a"), Row(value="view_b"), Row(value="view_c")]), set(result)
            )

    def test_streaming_drop_duplicate_within_watermark(self):
        """
        This verifies dropDuplicatesWithinWatermark works with a streaming dataframe.
        """
        user_schema = StructType().add("time", TimestampType()).add("id", "integer")
        df = (
            self.spark.readStream.option("sep", ";")
            .schema(user_schema)
            .csv("python/test_support/sql/streaming/time")
        )
        q1 = (
            df.withWatermark("time", "2 seconds")
            .dropDuplicatesWithinWatermark(["id"])
            .writeStream.outputMode("update")
            .format("memory")
            .queryName("test_streaming_drop_duplicates_within_wm")
            .start()
        )
        self.assertTrue(q1.isActive)
        q1.processAllAvailable()
        q1.stop()
        result = self.spark.sql("SELECT * FROM test_streaming_drop_duplicates_within_wm").collect()
        self.assertTrue(len(result) >= 6 and len(result) <= 9)


class StreamingTests(StreamingTestsMixin, ReusedSQLTestCase):
    def _assert_exception_tree_contains_msg(self, exception, msg):
        e = exception
        contains = msg in e._desc
        while e._cause is not None and not contains:
            e = e._cause
            contains = msg in e._desc
        self.assertTrue(contains, "Exception tree doesn't contain the expected message: %s" % msg)


if __name__ == "__main__":
    import unittest
    from pyspark.sql.tests.streaming.test_streaming import *  # noqa: F401

    try:
        import xmlrunner

        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
    except ImportError:
        testRunner = None
    unittest.main(testRunner=testRunner, verbosity=2)