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

import pyarrow as pa
import pytest
from datafusion import (
    ExecutionPlan,
    LogicalPlan,
    Metric,
    MetricsSet,
    SessionContext,
    col,
    udf,
)


# Note: We must use CSV because memory tables are currently not supported for
# conversion to/from protobuf.
@pytest.fixture
def df():
    ctx = SessionContext()
    return ctx.read_csv(path="testing/data/csv/aggregate_test_100.csv").select("c1")


def test_logical_plan_to_bytes_roundtrip(ctx, df) -> None:
    """Round-trip a LogicalPlan through the session's logical codec."""
    logical_plan_bytes = df.logical_plan().to_bytes()
    logical_plan = LogicalPlan.from_bytes(ctx, logical_plan_bytes)

    df_round_trip = ctx.create_dataframe_from_logical_plan(logical_plan)

    assert df.collect() == df_round_trip.collect()


def test_execution_plan_to_bytes_roundtrip(ctx, df) -> None:
    """Round-trip an ExecutionPlan through the session's physical codec."""
    original_execution_plan = df.execution_plan()
    execution_plan_bytes = original_execution_plan.to_bytes()
    execution_plan = ExecutionPlan.from_bytes(ctx, execution_plan_bytes)

    assert str(original_execution_plan) == str(execution_plan)


def test_logical_plan_to_proto_is_deprecated(ctx, df) -> None:
    """to_proto / from_proto still work but emit DeprecationWarning."""
    plan = df.logical_plan()

    with pytest.warns(DeprecationWarning, match="to_proto"):
        blob = plan.to_proto()
    with pytest.warns(DeprecationWarning, match="from_proto"):
        restored = LogicalPlan.from_proto(ctx, blob)

    df_round_trip = ctx.create_dataframe_from_logical_plan(restored)
    assert df.collect() == df_round_trip.collect()


def test_execution_plan_to_proto_is_deprecated(ctx, df) -> None:
    plan = df.execution_plan()

    with pytest.warns(DeprecationWarning, match="to_proto"):
        blob = plan.to_proto()
    with pytest.warns(DeprecationWarning, match="from_proto"):
        restored = ExecutionPlan.from_proto(ctx, blob)

    assert str(plan) == str(restored)


def test_session_with_logical_extension_codec_roundtrip(ctx, df) -> None:
    """A session with a non-default logical codec still round-trips builtins.

    The codec slot is overridable via with_logical_extension_codec; the
    PythonLogicalCodec wrapper delegates unhandled cases to the inner
    codec, so plans without Python UDFs are unaffected by the swap.
    """
    # Default-routed session should round-trip via to_bytes.
    blob = df.logical_plan().to_bytes()
    restored = LogicalPlan.from_bytes(ctx, blob)
    df_round_trip = ctx.create_dataframe_from_logical_plan(restored)
    assert df.collect() == df_round_trip.collect()


def test_installing_a_physical_codec_preserves_strict_mode() -> None:
    """Installing a physical extension codec must not re-enable inlining.

    `with_physical_extension_codec` builds a replacement `PythonPhysicalCodec`
    around the imported one, and the constructor defaults inlining to on. A
    context that opted out via `with_python_udf_inlining(enabled=False)` has to
    keep its setting, or installing a codec silently starts embedding
    cloudpickled callables in serialized execution plans.

    The logical counterpart lives in `test_pickle_expr.py`; this one needs an
    `ExecutionPlan` because only the physical codec encodes it. `DFPYUDF` is
    the scalar Python-UDF family prefix, shared by both layers; see
    `PY_SCALAR_UDF_FAMILY` in crates/core/src/codec.rs.
    """
    identity = udf(
        lambda arr: arr,
        [pa.string()],
        pa.string(),
        volatility="immutable",
        name="identity_str",
    )

    def plan_bytes(ctx: SessionContext) -> bytes:
        df = ctx.read_csv(path="testing/data/csv/aggregate_test_100.csv").select(
            identity(col("c1"))
        )
        return df.execution_plan().to_bytes(ctx)

    # The inlining default is on, so the strict blob is what has to differ.
    assert b"DFPYUDF" in plan_bytes(SessionContext())

    strict = SessionContext().with_python_udf_inlining(enabled=False)
    assert b"DFPYUDF" not in plan_bytes(strict)

    installed = strict.with_physical_extension_codec(
        strict.__datafusion_physical_extension_codec__()
    )
    assert b"DFPYUDF" not in plan_bytes(installed)


def test_session_codec_capsule_getters(ctx) -> None:
    """SessionContext exposes both logical and physical codec capsules."""
    logical = ctx.ctx.__datafusion_logical_extension_codec__()
    physical = ctx.ctx.__datafusion_physical_extension_codec__()
    assert logical is not None
    assert physical is not None


def test_metrics_tree_walk() -> None:
    ctx = SessionContext()
    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
    df = ctx.sql("SELECT * FROM t WHERE column1 > 1")
    df.collect()
    plan = df.execution_plan()

    results = plan.collect_metrics()
    assert len(results) >= 1
    output_rows_by_op: dict[str, int] = {}
    for name, ms in results:
        assert isinstance(name, str)
        assert isinstance(ms, MetricsSet)
        if ms.output_rows is not None:
            output_rows_by_op[name] = ms.output_rows

    # The filter passes rows where column1 > 1, so exactly
    # 2 rows from (1,'a'),(2,'b'),(3,'c').
    # At least one operator must report exactly 2 output rows (the filter).
    assert 2 in output_rows_by_op.values(), (
        f"Expected an operator with output_rows=2, got {output_rows_by_op}"
    )


def test_metric_properties() -> None:
    ctx = SessionContext()
    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
    df = ctx.sql("SELECT * FROM t WHERE column1 > 1")
    df.collect()
    plan = df.execution_plan()

    found_any_metric = False
    for _, ms in plan.collect_metrics():
        r = repr(ms)
        assert isinstance(r, str)
        for metric in ms.metrics():
            found_any_metric = True
            assert isinstance(metric, Metric)
            assert isinstance(metric.name, str)
            assert len(metric.name) > 0
            assert metric.partition is None or isinstance(metric.partition, int)
            assert metric.value is None or isinstance(
                metric.value, int | datetime.datetime
            )
            assert isinstance(metric.labels(), dict)
            mr = repr(metric)
            assert isinstance(mr, str)
            assert len(mr) > 0
    assert found_any_metric, "Expected at least one metric after execution"


def test_no_meaningful_metrics_before_execution() -> None:
    ctx = SessionContext()
    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
    df = ctx.sql("SELECT * FROM t WHERE column1 > 1")
    plan_before = df.execution_plan()

    # Some plan nodes (e.g. DataSourceExec) eagerly initialize a MetricsSet,
    # so metrics() may return a set even before execution.  However, no rows
    # should have been processed yet — output_rows must be absent or zero.
    for _, ms in plan_before.collect_metrics():
        rows = ms.output_rows
        assert rows is None or rows == 0, (
            f"Expected 0 output_rows before execution, got {rows}"
        )

    # After execution, at least one operator must report rows processed.
    df.collect()
    plan_after = df.execution_plan()
    output_rows_after = [
        ms.output_rows
        for _, ms in plan_after.collect_metrics()
        if ms.output_rows is not None and ms.output_rows > 0
    ]
    assert len(output_rows_after) > 0, "Expected output_rows > 0 after execution"


def test_collect_partitioned_metrics() -> None:
    ctx = SessionContext()
    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
    df = ctx.sql("SELECT * FROM t WHERE column1 > 1")

    df.collect_partitioned()
    plan = df.execution_plan()

    output_rows_values = [
        ms.output_rows for _, ms in plan.collect_metrics() if ms.output_rows is not None
    ]
    assert 2 in output_rows_values, f"Expected 2 in {output_rows_values}"


def test_execute_stream_metrics() -> None:
    ctx = SessionContext()
    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
    df = ctx.sql("SELECT * FROM t WHERE column1 > 1")

    for _ in df.execute_stream():
        pass

    plan = df.execution_plan()
    output_rows_values = [
        ms.output_rows for _, ms in plan.collect_metrics() if ms.output_rows is not None
    ]
    assert 2 in output_rows_values, f"Expected 2 in {output_rows_values}"


def test_execute_stream_partitioned_metrics() -> None:
    ctx = SessionContext()
    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
    df = ctx.sql("SELECT * FROM t WHERE column1 > 1")

    for stream in df.execute_stream_partitioned():
        for _ in stream:
            pass

    plan = df.execution_plan()
    output_rows_values = [
        ms.output_rows for _, ms in plan.collect_metrics() if ms.output_rows is not None
    ]
    assert 2 in output_rows_values, f"Expected 2 in {output_rows_values}"


def test_value_as_datetime() -> None:
    ctx = SessionContext()
    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
    df = ctx.sql("SELECT * FROM t WHERE column1 > 1")
    df.collect()
    plan = df.execution_plan()

    for _, ms in plan.collect_metrics():
        for metric in ms.metrics():
            if metric.name in ("start_timestamp", "end_timestamp"):
                dt = metric.value_as_datetime
                assert dt is None or isinstance(dt, datetime.datetime)
                if dt is not None:
                    assert dt.tzinfo is not None
            else:
                assert metric.value_as_datetime is None


def test_metric_names_and_labels() -> None:
    """Verify that known metric names appear and labels are well-formed."""
    ctx = SessionContext()
    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
    df = ctx.sql("SELECT * FROM t WHERE column1 > 1")
    df.collect()
    plan = df.execution_plan()

    all_metric_names: set[str] = set()
    for _, ms in plan.collect_metrics():
        for metric in ms.metrics():
            all_metric_names.add(metric.name)
            # Labels must be a dict of str->str
            labels = metric.labels()
            for k, v in labels.items():
                assert isinstance(k, str)
                assert isinstance(v, str)

    # After a filter query, we expect at minimum these standard metric names.
    assert "output_rows" in all_metric_names, (
        f"Expected 'output_rows' in {all_metric_names}"
    )
    assert "elapsed_compute" in all_metric_names, (
        f"Expected 'elapsed_compute' in {all_metric_names}"
    )


def test_collect_twice_has_metrics() -> None:
    ctx = SessionContext()
    ctx.sql("CREATE TABLE t AS VALUES (1, 'a'), (2, 'b'), (3, 'c')")
    df = ctx.sql("SELECT * FROM t WHERE column1 > 1")

    df.collect()
    df.collect()

    plan = df.execution_plan()
    output_rows_values = [
        ms.output_rows for _, ms in plan.collect_metrics() if ms.output_rows is not None
    ]
    assert len(output_rows_values) > 0