from __future__ import annotations

import numpy as np
import pandas as pd
import pyarrow as pa
import pytest
from pytest import param

import xorq.api as xo
import xorq.vendor.ibis.expr.operations as ops
from xorq.tests.util import assert_frame_equal


@pytest.fixture
def client():
    return xo.pandas.connect(
        {
            "df": pd.DataFrame({"a": [1, 2, 3], "b": list("abc")}),
            "df_unknown": pd.DataFrame({"array_of_strings": [["a", "b"], [], ["c"]]}),
        }
    )


@pytest.fixture
def table(client):
    return client.table("df")


@pytest.fixture
def test_data():
    return pd.DataFrame({"A": [1, 2, 3, 4, 5], "B": list("abcde")})


def test_connect_no_args():
    con = xo.pandas.connect()
    assert dict(con.tables) == {}


def test_client_table(table):
    assert isinstance(table.op(), ops.DatabaseTable)


@pytest.mark.parametrize(
    "lamduh",
    [(lambda df: df), (lambda df: pa.Table.from_pandas(df))],
    ids=["dataframe", "pyarrow table"],
)
def test_create_table(client, test_data, lamduh):
    test_data = lamduh(test_data)
    client.create_table("testing", obj=test_data)
    assert "testing" in client.list_tables()
    client.create_table("testingschema", schema=client.get_schema("testing"))
    assert "testingschema" in client.list_tables()


def test_literal(client):
    lit = xo.literal(1)
    result = client.execute(lit)
    assert result == 1


def test_list_tables(client):
    assert client.list_tables(like="df_unknown")
    assert not client.list_tables(like="not_in_the_database")
    assert client.list_tables()


def test_drop(table):
    table = table.mutate(c=table.a)
    expr = table.drop("a")
    result = expr.execute()
    expected = table[["b", "c"]].execute()
    assert_frame_equal(result, expected)


@pytest.mark.parametrize(
    "unit",
    [
        "Y",
        "M",
        "D",
        "h",
        "m",
        "s",
        "ms",
        "us",
        "ns",
        param("ps", marks=pytest.mark.xfail),
        param("fs", marks=pytest.mark.xfail),
        param("as", marks=pytest.mark.xfail),
    ],
)
def test_datetime64_infer(client, unit):
    value = np.datetime64("2018-01-02", unit)
    expr = xo.literal(value, type="timestamp")
    result = client.execute(expr)
    assert result == pd.Timestamp(value).to_pydatetime()
