import datetime
import functools
import operator
import threading
import time

import pandas as pd
import pyarrow as pa
import pytest
from pyarrow._flight import FlightUnauthenticatedError

import xorq.api as xo
import xorq.expr.datatypes as dt
import xorq.flight.action as A
import xorq.flight.exchanger as E
from xorq.common.utils import classproperty
from xorq.common.utils.rbr_utils import instrument_reader
from xorq.common.utils.tls_utils import TLSKwargs
from xorq.flight import (
    Backend,
    BasicAuth,
    FlightServer,
    FlightUrl,
)
from xorq.flight.action import AddExchangeAction
from xorq.flight.exchanger import EchoExchanger, PandasUDFExchanger
from xorq.flight.tests.conftest import do_agg, field_name, my_udf, return_type
from xorq.loader import load_backend
from xorq.tests.util import assert_frame_equal


def make_flight_url(port, scheme="grpc", auth=None):
    if port is not None:
        assert not FlightUrl.port_in_use(port), f"Port {port} already in use"
    flight_url = (
        FlightUrl(port=port, scheme=scheme)
        if auth is None
        else FlightUrl(
            port=port, scheme=scheme, username=auth.username, password=auth.password
        )
    )
    assert FlightUrl.port_in_use(flight_url.port), (
        f"Port {flight_url.port} should be in use"
    )
    return flight_url


def make_connection_from_name(backend_name):
    fun = load_backend(backend_name).connect if backend_name else xo.connect
    return (
        functools.partial(fun, check_same_thread=False)
        if backend_name == "sqlite"
        else fun
    )


@pytest.mark.parametrize(
    "connection,port",
    [
        pytest.param("duckdb", 5005, id="duckdb"),
        pytest.param("datafusion", 5005, id="datafusion"),
        pytest.param("sqlite", 5005, id="sqlite"),
        pytest.param("", 5005, id="xorq"),
    ],
)
def test_port_in_use(connection, port):
    assert port is not None
    flight_url = make_flight_url(port)
    with pytest.raises(OSError, match="Address already in use"):
        with FlightServer(
            flight_url=flight_url,
            make_connection=make_connection_from_name(connection),
        ) as _:
            # entering the above context releases the port
            # so we won't raise until we enter the second context and try to use it
            flight_url2 = FlightUrl(port=port)  # noqa: F841


class Answer42Action(A.AbstractAction):
    @classproperty
    def name(cls):
        return "answer-42"

    @classmethod
    def description(cls):
        return (
            "the answer to the ultimate question of life, the universe, and everything"
        )

    @classmethod
    def do_action(cls, server, context, action):
        yield A.make_flight_result(42)


def test_list_actions():
    with FlightServer() as flight_server:
        actions = flight_server.client.do_action_one(A.ListActionsAction.name)
        assert actions == tuple(A.actions)


def test_add_action():
    with FlightServer() as flight_server:
        actions = flight_server.client.do_action_one(A.ListActionsAction.name)
        assert Answer42Action.name not in actions
        flight_server.client.do_action(A.AddActionAction.name, Answer42Action)
        actions = flight_server.client.do_action_one(A.ListActionsAction.name)
        assert Answer42Action.name in actions


def test_list_exchanges():
    with FlightServer() as flight_server:
        exchanges = flight_server.client.do_action_one(A.ListExchangesAction.name)
        assert exchanges == tuple(E.exchangers)


@pytest.mark.parametrize(
    "connection,port",
    [
        pytest.param("duckdb", None, id="duckdb"),
        pytest.param("datafusion", None, id="datafusion"),
        pytest.param("sqlite", None, id="sqlite"),
        pytest.param("", None, id="xorq"),
    ],
)
def test_register_and_list_tables(connection, port):
    flight_url = make_flight_url(port)

    with FlightServer(
        flight_url=flight_url,
        verify_client=False,
        make_connection=make_connection_from_name(connection),
    ) as main:
        con = main.con
        assert con.version is not None

        data = pa.table(
            {"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"]}
        ).to_pandas()

        con.register(data, table_name="users")
        t = con.table("users")
        actual = xo.execute(t)

        assert t.schema() is not None
        assert "users" in con.list_tables()
        assert isinstance(actual, pd.DataFrame)


@pytest.mark.parametrize("auth", [None, BasicAuth("username", "password")])
@pytest.mark.parametrize("verify_client", [False, True])
def test_tls_encryption(auth, verify_client, tls_kwargs, mtls_kwargs):
    flight_url = make_flight_url(None, scheme="grpc+tls")

    tls_kwargs = TLSKwargs.from_common_name(verify_client=verify_client)

    with FlightServer(
        flight_url=flight_url,
        verify_client=verify_client,
        auth=auth,
        **tls_kwargs.server_kwargs,
    ) as main:
        con = main.con
        assert con.version is not None

        data = pa.table(
            {"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"]}
        ).to_pandas()

        con.register(data, table_name="users")
        t = con.table("users")
        actual = xo.execute(t)

        assert t.schema() is not None
        assert "users" in con.list_tables()
        assert isinstance(actual, pd.DataFrame)


def test_failed_auth(tls_kwargs):
    flight_url = make_flight_url(None, scheme="grpc+tls")

    kwargs = TLSKwargs.from_common_name().server_kwargs

    with FlightServer(
        flight_url=flight_url,
        auth=BasicAuth("username", "password"),
        **kwargs,
    ) as server:
        kwargs = {
            "host": server.flight_url.host,
            "port": server.flight_url.port,
            "username": server.auth.username,
            "password": "not_the_password",
            "tls_root_certs": kwargs["root_certificates"],
        }

        with pytest.raises(FlightUnauthenticatedError):
            instance = Backend()
            instance.do_connect(**kwargs)


@pytest.mark.parametrize(
    "connection,port",
    [
        pytest.param("duckdb", None, id="duckdb"),
        pytest.param("datafusion", None, id="datafusion"),
        pytest.param("sqlite", None, id="sqlite"),
        pytest.param("", None, id="xorq"),
    ],
)
def test_into_backend_flight_server(connection, port, parquet_dir):
    batting = xo.read_parquet(parquet_dir / "batting.parquet")
    flight_url = make_flight_url(port)

    with FlightServer(
        flight_url=flight_url,
        verify_client=False,
        make_connection=make_connection_from_name(connection),
    ) as main:
        con = main.con
        t = batting.filter(batting.yearID == 2015).into_backend(con, "xo_batting")
        expr = (
            t.join(t, "playerID")
            .limit(15)
            .select(player_id="playerID", year_id="yearID_right")
        )

        assert not expr.execute().empty


@pytest.mark.parametrize(
    "connection,port",
    [
        pytest.param("duckdb", None, id="duckdb"),
        pytest.param("datafusion", None, id="datafusion"),
        pytest.param("sqlite", None, id="sqlite"),
        pytest.param("", None, id="xorq"),
    ],
)
def test_read_parquet(connection, port, parquet_dir):
    flight_url = make_flight_url(port)
    with FlightServer(
        flight_url=flight_url,
        verify_client=False,
        make_connection=make_connection_from_name(connection),
    ) as main:
        con = main.con
        batting = con.read_parquet(parquet_dir / "batting.parquet")
        assert xo.execute(batting) is not None


@pytest.mark.parametrize(
    "connection,port",
    [
        pytest.param("duckdb", None, id="duckdb"),
        pytest.param("datafusion", None, id="datafusion"),
        pytest.param("sqlite", None, id="sqlite"),
        pytest.param("", None, id="xorq"),
    ],
)
def test_exchange(connection, port):
    flight_url = make_flight_url(port)

    def my_f(df):
        return df[["a", "b"]].sum(axis=1)

    with FlightServer(
        flight_url=flight_url,
        verify_client=False,
        make_connection=make_connection_from_name(connection),
    ) as main:
        client = main.client
        udf_exchanger = PandasUDFExchanger(
            my_f,
            schema_in=xo.schema({"a": int, "b": int}),
            name="x",
            typ=dt.int64,
            append=True,
        )
        client.do_action(AddExchangeAction.name, udf_exchanger, options=client._options)

        # a small example
        df_in = pd.DataFrame({"a": [1], "b": [2], "c": [100]})
        fut, rbr = client.do_exchange(
            udf_exchanger.command,
            xo.memtable(df_in),
        )
        df_out = rbr.read_pandas()
        writes_reads = fut.result()

        assert writes_reads["n_writes"] == writes_reads["n_reads"]
        assert df_out is not None

        # demonstrate streaming
        df_in = pd.DataFrame(
            {
                "a": range(100_000),
                "b": range(100_000, 200_000),
                "c": range(200_000, 300_000),
            }
        )
        fut, rbr = client.do_exchange_batches(
            udf_exchanger.command,
            instrument_reader(pa.Table.from_pandas(df_in).to_reader(max_chunksize=100)),
        )
        first_batch = next(rbr)
        first_batch_time = datetime.datetime.now()
        assert first_batch is not None, f"must get first batch by {first_batch_time}"

        rest = rbr.read_pandas()
        rest_time = datetime.datetime.now()
        assert rest is not None, f"must get first batch by {rest_time}"
        assert first_batch_time < rest_time

        writes_reads = fut.result()
        assert writes_reads["n_writes"] == 1000  # because
        assert writes_reads["n_reads"] == 1000


@pytest.mark.parametrize(
    "connection",
    (
        "duckdb",
        "datafusion",
        "sqlite",
    ),
)
def test_reentry(connection):
    df_in = pd.DataFrame({"a": [1], "b": [2], "c": [100]})
    with FlightServer(
        verify_client=False,
        make_connection=make_connection_from_name(connection),
    ) as server:
        fut, rbr = server.client.do_exchange_batches(
            EchoExchanger.command,
            pa.RecordBatchReader.from_stream(df_in),
        )
        df_out = rbr.read_pandas()
        assert df_in.equals(df_out)
    with server:
        fut, rbr = server.client.do_exchange_batches(
            EchoExchanger.command,
            pa.RecordBatchReader.from_stream(df_in),
        )
        df_out = rbr.read_pandas()
        assert df_in.equals(df_out)


@pytest.mark.parametrize(
    "connection",
    (
        "duckdb",
        "datafusion",
        "sqlite",
    ),
)
def test_serve_close(connection):
    df_in = pd.DataFrame({"a": [1], "b": [2], "c": [100]})
    server = FlightServer(
        verify_client=False,
        make_connection=make_connection_from_name(connection),
    )

    server.serve()
    fut, rbr = server.client.do_exchange_batches(
        EchoExchanger.command,
        pa.RecordBatchReader.from_stream(df_in),
    )
    df_out = rbr.read_pandas()
    assert df_in.equals(df_out)
    server.close()

    server.serve()
    fut, rbr = server.client.do_exchange_batches(
        EchoExchanger.command,
        pa.RecordBatchReader.from_stream(df_in),
    )
    df_out = rbr.read_pandas()
    assert df_in.equals(df_out)
    server.close()


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

    schema_in = xo.schema({"dummy": "int64"})
    dummy_udxf = E.make_udxf(
        dummy,
        schema_in,
        xo.schema({"row_count": "int64"}),
    )
    flight_server = FlightServer(exchangers=[dummy_udxf])
    with flight_server:
        client = flight_server.client
        available = client.do_action_one(A.ListExchangesAction.name)
        assert dummy_udxf.command in available
        client.do_exchange(
            dummy_udxf.command,
            xo.memtable({"dummy": [0]}, schema=schema_in),
        )


@pytest.mark.parametrize(
    "expr",
    [
        pytest.param(xo.table({"id": int}, name="users").count(), id="scalar"),
        pytest.param(xo.literal(1), id="literal"),
        pytest.param(xo.table({"id": int}, name="users").id, id="column"),
    ],
)
def test_execute_query_non_relation_expr(expr):
    flight_url = make_flight_url(None)
    data = pa.table({"id": [1, 2, 3], "name": ["Alice", "Bob", "Charlie"]}).to_pandas()

    with FlightServer(
        flight_url=flight_url,
        verify_client=False,
        make_connection=xo.duckdb.connect,
    ) as main:
        main.con.register(data, table_name="users")
        actual = main.client.execute(expr)
        assert isinstance(actual, pa.Table)
        assert len(actual) > 0


@pytest.mark.parametrize("block", [True, False])
def test_server_blocks(block):
    flight_url = make_flight_url(None)
    server = FlightServer(
        flight_url=flight_url,
        verify_client=False,
        make_connection=xo.duckdb.connect,
    )

    server_thread = threading.Thread(
        target=functools.partial(server.serve, block=block), daemon=True
    )

    is_blocking = True

    def check_if_still_running():
        nonlocal is_blocking
        time.sleep(1)
        is_blocking = server_thread.is_alive()

    server_thread.start()

    checker_thread = threading.Thread(target=check_if_still_running)
    checker_thread.start()
    checker_thread.join()

    # Try to stop the inner server
    server.server.shutdown()
    server_thread.join(timeout=2.0)

    assert is_blocking == block
    assert not server_thread.is_alive()


def test_exchange_server_from_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})
    command = "diamonds_exchange_command"
    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,
        make_udxf_kwargs={"name": my_udf.__name__, "command": command},
    ).order_by("cut")

    with FlightServer.from_udxf(expr) as server:
        client = server.client
        assert client is not None

        _, rbr = client.do_exchange(
            command,
            input_expr,
        )

        actual = rbr.read_pandas().sort_values("cut", ignore_index=True)
        expected = baseline.sort_values("cut", ignore_index=True)
        assert_frame_equal(
            actual,
            expected,
            check_exact=False,
        )
