import operator

import pandas as pd
import pyarrow as pa
import pytest

import xorq.api as xo
from xorq.flight.exchanger import make_udxf
from xorq.flight.tests.conftest import (
    do_agg,
    field_name,
    my_udf,
    my_udf_on_expr,
    return_type,
)


def test_flight_expr(con, diamonds, baseline):
    unbound_expr = (
        xo.table(diamonds.schema()).pipe(do_agg).mutate(my_udf_on_expr).order_by("cut")
    )
    expr = xo.expr.relations.flight_expr(
        diamonds,
        unbound_expr,
        inner_name="flight-expr",
        name="remote-expr",
        con=con,
    )
    df = expr.execute()
    pd.testing.assert_frame_equal(
        baseline.sort_values("cut", ignore_index=True),
        df.sort_values("cut", ignore_index=True),
        check_exact=False,
    )


def test_flight_udxf(con, diamonds, baseline):
    input_expr = diamonds.pipe(do_agg)
    process_df = operator.methodcaller("assign", **{field_name: my_udf.fn})
    maybe_schema_in = input_expr.schema()
    maybe_schema_out = xo.schema(input_expr.schema() | {field_name: return_type})
    expr = xo.expr.relations.flight_udxf(
        input_expr,
        process_df=process_df,
        maybe_schema_in=maybe_schema_in,
        maybe_schema_out=maybe_schema_out,
        con=con,
        # operator.methodcaller doesn't have name, so must explicitly pass
        make_udxf_kwargs={"name": my_udf.__name__},
    ).order_by("cut")
    df = expr.execute()
    actual = df.sort_values("cut", ignore_index=True)
    expected = baseline.sort_values("cut", ignore_index=True)
    pd.testing.assert_frame_equal(
        actual,
        expected,
        check_exact=False,
    )


def test_make_udxf_fails():
    def dummy(df: pd.DataFrame):
        return pd.DataFrame({"row_count": [42]})

    with pytest.raises(ValueError):
        make_udxf(
            dummy,
            xo.schema({"dummy": "int64"}),
            pa.schema(
                [
                    ("row_count", pa.int64()),
                ]
            ),
        )

    with pytest.raises(ValueError):
        make_udxf(
            dummy,
            pa.schema(
                [
                    ("dummy", pa.int64()),
                ]
            ),
            xo.schema({"row_count": "int64"}),
        )
