"""Tests for partition transforms using PySpark's DataFrameWriterV2 API with Glue catalog.

NOTE: These tests are currently skipped because moto does not support the AWS Glue
OpenTableFormatInput API which is required for creating Iceberg tables. The Rust
integration tests for Glue partition transforms are also skipped for the same reason.

Tracking issue: https://github.com/getmoto/moto/issues/9705

To run these tests, you need a real AWS Glue endpoint or a mock that supports Iceberg.
"""

import platform

import pytest
from pyiceberg.catalog.glue import GlueCatalog
from pyiceberg.partitioning import (
    BucketTransform,
    DayTransform,
    HourTransform,
    MonthTransform,
    YearTransform,
)

from pysail.testing.spark.utils.common import is_jvm_spark

try:
    from pyspark.sql.functions import partitioning

    HAS_PARTITIONING = True
except ImportError:
    HAS_PARTITIONING = False
    partitioning = None

# Skip reason for Glue Iceberg tests
MOTO_ICEBERG_SKIP = "moto does not support Glue OpenTableFormatInput API for Iceberg tables"


@pytest.mark.skip(reason=MOTO_ICEBERG_SKIP)
@pytest.mark.skipif(platform.system() == "Windows", reason="may not work on Windows")
@pytest.mark.skipif(is_jvm_spark(), reason="Spark does not handle v1 and v2 tables properly")
@pytest.mark.skipif(not HAS_PARTITIONING, reason="partitioning module not available in this Spark version")
@pytest.mark.parametrize(
    ("transform_name", "transform_func", "expected_transform_type"),
    [
        ("years", lambda p: p.years("dt"), YearTransform),
        ("months", lambda p: p.months("dt"), MonthTransform),
        ("days", lambda p: p.days("dt"), DayTransform),
        ("hours", lambda p: p.hours("ts"), HourTransform),
    ],
    ids=["years", "months", "days", "hours"],
)
def test_glue_partition_transform_time_based(
    glue_spark, moto_endpoint, transform_name, transform_func, expected_transform_type
):
    """Test time-based partition transforms via PySpark V2 API with Glue catalog."""
    table_name = f"t_{transform_name}"

    # Create dataframe with timestamp data
    df = glue_spark.createDataFrame(
        [(1, "2023-03-15 10:30:00"), (2, "2024-06-20 14:45:00")],
        schema="id INT, ts STRING",
    )
    df = df.selectExpr("id", "to_date(ts) as dt", "to_timestamp(ts) as ts")

    # Create table with partition transform
    # Note: Glue catalog doesn't support OPTIONS, so we don't use .option("location", ...)
    df.writeTo(f"test_db.{table_name}").using("iceberg").partitionedBy(transform_func(partitioning)).create()

    try:
        # Verify with pyiceberg
        catalog = GlueCatalog(
            "glue",
            **{
                "glue.region": "us-east-1",
                "glue.endpoint": moto_endpoint,
                "glue.access-key-id": "testing",
                "glue.secret-access-key": "testing",
            },
        )
        table = catalog.load_table("test_db." + table_name)
        spec = table.spec()

        assert len(spec.fields) == 1, f"Expected 1 partition field, got {len(spec.fields)}"
        assert isinstance(spec.fields[0].transform, expected_transform_type), (
            f"Expected {expected_transform_type.__name__}, got {type(spec.fields[0].transform).__name__}"
        )

        # Verify data can be read back via table name
        result = glue_spark.table(f"test_db.{table_name}").select("id").orderBy("id").collect()
        assert [r[0] for r in result] == [1, 2]
    finally:
        glue_spark.sql(f"DROP TABLE IF EXISTS test_db.{table_name}")


@pytest.mark.skip(reason=MOTO_ICEBERG_SKIP)
@pytest.mark.skipif(platform.system() == "Windows", reason="may not work on Windows")
@pytest.mark.skipif(is_jvm_spark(), reason="Spark does not handle v1 and v2 tables properly")
@pytest.mark.skipif(not HAS_PARTITIONING, reason="partitioning module not available in this Spark version")
def test_glue_partition_transform_bucket(glue_spark, moto_endpoint):
    """Test bucket partition transform via PySpark V2 API with Glue catalog."""
    table_name = "t_bucket"
    num_buckets = 10

    # Create dataframe
    df = glue_spark.createDataFrame([(1, "a"), (2, "b"), (3, "c"), (4, "d")], schema="id INT, name STRING")

    # Create table with bucket partition
    # Note: Glue catalog doesn't support OPTIONS, so we don't specify a location
    df.writeTo(f"test_db.{table_name}").using("iceberg").partitionedBy(partitioning.bucket(num_buckets, "id")).create()

    try:
        # Verify with pyiceberg
        catalog = GlueCatalog(
            "glue",
            **{
                "glue.region": "us-east-1",
                "glue.endpoint": moto_endpoint,
                "glue.access-key-id": "testing",
                "glue.secret-access-key": "testing",
            },
        )
        table = catalog.load_table("test_db." + table_name)
        spec = table.spec()

        assert len(spec.fields) == 1, f"Expected 1 partition field, got {len(spec.fields)}"
        assert isinstance(spec.fields[0].transform, BucketTransform), (
            f"Expected BucketTransform, got {type(spec.fields[0].transform).__name__}"
        )
        assert spec.fields[0].transform.num_buckets == num_buckets

        # Verify data can be read back via table name
        result = glue_spark.table(f"test_db.{table_name}").select("id").orderBy("id").collect()
        assert [r[0] for r in result] == [1, 2, 3, 4]
    finally:
        glue_spark.sql(f"DROP TABLE IF EXISTS test_db.{table_name}")
