import logging
from types import SimpleNamespace
from typing import Optional

import pyarrow as pa
import pytest

from ray.data._internal.execution.interfaces.ref_bundle import RefBundle
from ray.data._internal.execution.streaming_executor_state import (
    OpState,
    _build_schemas_mismatch_warning,
    dedupe_schemas_with_validation,
)
from ray.data._internal.pandas_block import PandasBlockSchema
from ray.data.block import Schema, _is_empty_schema


@pytest.mark.parametrize(
    "incoming_schema",
    [
        pa.schema([pa.field("uuid", pa.string())]),  # NOTE: diff from old_schema
        pa.schema([]),  # Empty Schema
        PandasBlockSchema(names=["col1"], types=[int]),
        PandasBlockSchema(names=[], types=[]),
        None,  # Null Schema
    ],
)
@pytest.mark.parametrize(
    "old_schema",
    [
        pa.schema([pa.field("id", pa.int64())]),
        pa.schema([]),  # Empty Schema
        PandasBlockSchema(names=["col2"], types=[int]),
        PandasBlockSchema(names=[], types=[]),
        None,  # Null Schema
    ],
)
def test_dedupe_schema_handle_empty(
    old_schema: Optional["Schema"],
    incoming_schema: Optional["Schema"],
):

    incoming_bundle = RefBundle([], owns_blocks=False, schema=incoming_schema)
    out_bundle, diverged = dedupe_schemas_with_validation(
        old_schema, incoming_bundle, enforce_schemas=False
    )

    if _is_empty_schema(old_schema):
        # old_schema is invalid
        assert not diverged, (old_schema, incoming_schema)
        assert out_bundle.schema == incoming_schema, (old_schema, incoming_schema)
    else:
        # old_schema is valid
        assert diverged, (old_schema, incoming_schema)
        assert incoming_schema != old_schema, (old_schema, incoming_schema)
        assert old_schema == out_bundle.schema, (old_schema, incoming_schema)


@pytest.mark.parametrize("enforce_schemas", [False, True])
@pytest.mark.parametrize(
    "incoming_schema", [pa.schema([pa.field("uuid", pa.string())])]
)
@pytest.mark.parametrize("old_schema", [pa.schema([pa.field("id", pa.int64())])])
def test_dedupe_schema_divergence(
    enforce_schemas: bool,
    old_schema: Optional["Schema"],
    incoming_schema: Optional["Schema"],
):

    incoming_bundle = RefBundle([], owns_blocks=False, schema=incoming_schema)
    out_bundle, diverged = dedupe_schemas_with_validation(
        old_schema, incoming_bundle, enforce_schemas=enforce_schemas
    )

    assert diverged

    if enforce_schemas:
        assert out_bundle.schema == pa.schema(list(old_schema) + list(incoming_schema))
    else:
        assert out_bundle.schema == old_schema


@pytest.mark.parametrize(
    "incoming_schema",
    [
        pa.schema([]),  # Empty Schema
        PandasBlockSchema(names=[], types=[]),
        None,  # Null Schema
    ],
)
def test_build_mismatch_warning_empty(incoming_schema: Optional["Schema"]):
    old_schema = pa.schema(
        [
            pa.field("foo", pa.int32()),
            pa.field("bar", pa.string()),
            pa.field("baz", pa.bool_()),
        ]
    )

    msg = _build_schemas_mismatch_warning(old_schema, incoming_schema)

    assert (
        """Operator produced a RefBundle with an empty/unknown schema. (3 total):
    foo: int32
    bar: string
    baz: bool
This may lead to unexpected behavior."""
        == msg
    )


@pytest.mark.parametrize("truncation_length", [20, 4, 3, 2])
def test_build_mismatch_warning_truncation(truncation_length: int):
    old_schema = pa.schema(
        [
            pa.field("foo", pa.int32()),
            pa.field("bar", pa.string()),
            pa.field("baz", pa.bool_()),
            pa.field("quux", pa.uint16()),
            pa.field("corge", pa.int64()),
            pa.field("grault", pa.int32()),
        ]
    )
    incoming_schema = pa.schema(
        [
            pa.field("foo", pa.int64()),
            pa.field("baz", pa.bool_()),
            pa.field("qux", pa.float64()),
            pa.field("quux", pa.uint32()),
            pa.field("corge", pa.int32()),
            pa.field("grault", pa.int64()),
        ]
    )

    msg = _build_schemas_mismatch_warning(
        old_schema, incoming_schema, truncation_length=truncation_length
    )

    if truncation_length >= 4:
        assert (
            """Operator produced a RefBundle with a different schema than the previous one.
Fields exclusive to the incoming schema (1 total):
    qux: double
Fields exclusive to the old schema (1 total):
    bar: string
Fields that have different types across the old and the incoming schemas (4 total):
    foo: int32 => int64
    quux: uint16 => uint32
    corge: int64 => int32
    grault: int32 => int64
This may lead to unexpected behavior."""
            == msg
        )
    elif truncation_length == 3:
        assert (
            """Operator produced a RefBundle with a different schema than the previous one.
Fields exclusive to the incoming schema (1 total):
    qux: double
Fields exclusive to the old schema (1 total):
    bar: string
Fields that have different types across the old and the incoming schemas (4 total):
    foo: int32 => int64
    quux: uint16 => uint32
    corge: int64 => int32
    ... and 1 more
This may lead to unexpected behavior."""
            == msg
        )
    else:
        assert (
            """Operator produced a RefBundle with a different schema than the previous one.
Fields exclusive to the incoming schema (1 total):
    qux: double
Fields exclusive to the old schema (1 total):
    bar: string
Fields that have different types across the old and the incoming schemas (4 total):
    foo: int32 => int64
    quux: uint16 => uint32
    ... and 2 more
This may lead to unexpected behavior."""
            == msg
        )


def test_build_mismatch_warning_disordered():
    old_schema = pa.schema(
        [
            pa.field("foo", pa.int32()),
            pa.field("bar", pa.string()),
            pa.field("baz", pa.bool_()),
            pa.field("qux", pa.float64()),
        ]
    )
    incoming_schema = pa.schema(
        [
            pa.field("foo", pa.int32()),
            pa.field("baz", pa.bool_()),
            pa.field("bar", pa.string()),
            pa.field("qux", pa.float64()),
        ]
    )

    msg = _build_schemas_mismatch_warning(old_schema, incoming_schema)

    assert (
        """Operator produced a RefBundle with a different schema than the previous one.
Some fields are ordered differently across the old and the incoming schemas.
This may lead to unexpected behavior."""
        == msg
    )


class _DummyOp:
    def __init__(self, enforce_schemas=True):
        self.name = "dummy"
        self.data_context = SimpleNamespace(enforce_schemas=enforce_schemas)
        self.input_dependencies = []
        self.output_dependencies = []
        self.metrics = SimpleNamespace(
            num_alive_actors=0,
            num_restarting_actors=0,
            num_pending_actors=0,
            num_external_inqueue_blocks=0,
            num_external_inqueue_bytes=0,
            num_external_outqueue_blocks=0,
            num_external_outqueue_bytes=0,
        )

    def num_output_splits(self):
        return 1

    def get_actor_info(self):
        return SimpleNamespace(
            running=0,
            restarting=0,
            pending=0,
            active=0,
            idle=0,
            pool_utilization=0,
            tasks_in_flight=0,
        )


@pytest.mark.parametrize("enforce_schemas", [False, True])
def test_add_output_emits_warning(enforce_schemas, caplog, propagate_logs):
    op = _DummyOp(enforce_schemas=enforce_schemas)
    state = OpState(op, [])

    old_schema = pa.schema(
        [
            pa.field("foo", pa.int32()),
            pa.field("bar", pa.bool_()),
        ]
    )
    new_schema = pa.schema(
        [
            pa.field("foo", pa.int32()),
            pa.field("baz", pa.uint32()),
        ]
    )

    # First output seeds schema (no warning)
    state.add_output(RefBundle([], owns_blocks=False, schema=old_schema))

    assert not caplog.messages

    # Second output diverges, should warn once
    with caplog.at_level(logging.WARNING):
        state.add_output(RefBundle([], owns_blocks=False, schema=new_schema))

    assert not (enforce_schemas ^ len(caplog.messages))

    if enforce_schemas:
        msg = "\n".join(caplog.messages)
        assert (
            msg
            == """Operator produced a RefBundle with a different schema than the previous one.
Fields exclusive to the incoming schema (1 total):
    baz: uint32
Fields exclusive to the old schema (1 total):
    bar: bool
This may lead to unexpected behavior."""
        )


if __name__ == "__main__":
    import sys

    sys.exit(pytest.main(["-v", __file__]))