from __future__ import annotations

import doctest
import os
import time
from typing import TYPE_CHECKING

import pytest
from _pytest.doctest import DoctestItem
from pyspark.sql import SparkSession

from pysail.tests.spark.utils import SAIL_ONLY, is_jvm_spark


def pytest_configure(config):
    # Register custom markers.
    # Note: pytest-bdd converts @sail-only tag to "sail-only" marker (preserves hyphen)
    config.addinivalue_line(
        "markers",
        "sail-only: mark test as Sail-only (skipped when running against Spark JVM)",
    )
    # Load all pytest-bdd step modules.
    config.pluginmanager.import_plugin("pysail.tests.spark.steps.file_tree")
    config.pluginmanager.import_plugin("pysail.tests.spark.steps.sql")
    config.pluginmanager.import_plugin("pysail.tests.spark.steps.plan")
    config.pluginmanager.import_plugin("pysail.tests.spark.steps.delta_log")
    config.pluginmanager.import_plugin("pysail.tests.spark.steps.dataframe")


if TYPE_CHECKING:
    import pyspark.sql.connect.session


@pytest.fixture(scope="session")
def remote():
    """Creates a Spark Connect server if there is not one already running
    whose address is set in the `SPARK_REMOTE` environment variable.

    :yields: The remote address of the Spark Connect server to connect to.
    """
    if r := os.environ.get("SPARK_REMOTE"):
        yield r
    else:
        from pysail.spark import SparkConnectServer

        server = SparkConnectServer("127.0.0.1", 0)
        server.start(background=True)
        _, port = server.listening_address
        yield f"sc://localhost:{port}"
        server.stop()


@pytest.fixture(scope="module")
def spark(remote):
    """Create and configure a Spark Session to be used in the tests.
    After the tests are finished, the Spark Session is stopped.

    :param remote: The remote address of the Spark Connect server to connect to.
    :yields: A Spark Session configured for the tests.
    """
    spark = SparkSession.builder.remote(remote).getOrCreate()
    configure_spark_session(spark)
    patch_spark_connect_session(spark)
    yield spark
    spark.stop()


@pytest.fixture
def spark_session_factory(remote):
    """Factory for creating independent SparkSessions.

    Each call to the factory creates a new SparkSession with a unique session ID,
    allowing tests to verify session isolation behavior.

    :param remote: The remote address of the Spark Connect server.
    :yields: A factory function that creates new SparkSessions.
    """
    import contextlib
    import uuid

    sessions = []

    def create_session():
        # Use a unique app name to ensure we get a fresh session
        # The session ID is generated internally by Spark Connect
        unique_app = f"test_session_{uuid.uuid4().hex[:8]}"
        session = (
            SparkSession.builder.appName(unique_app).remote(remote).create()
        )  # Use create() instead of getOrCreate() to force new session
        configure_spark_session(session)
        patch_spark_connect_session(session)
        sessions.append(session)
        return session

    yield create_session

    # Cleanup all created sessions
    for session in sessions:
        # Best-effort cleanup: ignore errors while stopping Spark sessions during test teardown.
        with contextlib.suppress(Exception):
            session.stop()


def configure_spark_session(session):
    # Set the Spark session time zone to UTC by default.
    # Some test data (e.g. TPC-DS data) may generate timestamps that is invalid
    # in some local time zones. This would result in `pytz.exceptions.NonExistentTimeError`
    # when converting such timestamps from the local time zone to UTC.
    session.conf.set("spark.sql.session.timeZone", "UTC")
    # Enable Arrow to avoid data type errors when creating Spark DataFrame from Pandas.
    session.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")


def patch_spark_connect_session(session: pyspark.sql.connect.session.SparkSession):
    """
    Patch the Spark Connect session to avoid deadlock when closing the session.
    """
    f = session._client.close  # noqa: SLF001

    def close():
        if session._client._closed:  # noqa: SLF001
            return
        return f()

    session._client.close = close  # noqa: SLF001


@pytest.fixture(scope="module", autouse=True)
def spark_doctest(doctest_namespace, spark):
    # The Spark session is scoped to each module, so that the registered
    # temporary views and UDFs do not interfere with each other.
    doctest_namespace["spark"] = spark


@pytest.fixture
def session_timezone(spark, request):
    tz = spark.conf.get("spark.sql.session.timeZone")
    spark.conf.set("spark.sql.session.timeZone", request.param)
    yield
    spark.conf.set("spark.sql.session.timeZone", tz)


@pytest.fixture
def local_timezone(request):
    tz = os.environ.get("TZ")
    os.environ["TZ"] = request.param
    time.tzset()
    yield
    if tz is None:
        os.environ.pop("TZ")
    else:
        os.environ["TZ"] = tz
    time.tzset()


def pytest_collection_modifyitems(session, config, items):  # noqa: ARG001
    if is_jvm_spark():
        skip_sail_only = pytest.mark.skip(reason="Sail-only feature, not supported by Spark")
        for item in items:
            if isinstance(item, DoctestItem):
                for example in item.dtest.examples:
                    if example.options.get(SAIL_ONLY):
                        example.options[doctest.SKIP] = True
            # Skip pytest-bdd scenarios with @sail-only tag
            # Note: pytest-bdd preserves the hyphen in marker names
            elif item.get_closest_marker("sail-only"):
                item.add_marker(skip_sail_only)
