import csv
import json
import os
import tempfile
from textwrap import dedent
from uuid import uuid4

import pytest

from .datasource import Datasource


def test_create_datasource(datasource, data):
    assert datasource is not None
    assert datasource.name == "test_datasource"
    assert datasource.length == len(data)


def test_create_datasource_unauthenticated(unauthenticated_client):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            Datasource.from_list("test_datasource", [{"value": "test"}])


def test_create_datasource_already_exists_error(data, datasource):
    with pytest.raises(ValueError):
        Datasource.from_list("test_datasource", data, if_exists="error")


def test_create_datasource_already_exists_return(data, datasource):
    returned_dataset = Datasource.from_list("test_datasource", data, if_exists="open")
    assert returned_dataset is not None
    assert returned_dataset.name == "test_datasource"
    assert returned_dataset.length == len(data)


def test_create_datasource_from_hf_dataset(data):
    datasets = pytest.importorskip("datasets")
    hf_dataset = datasets.Dataset.from_list(data)
    datasource = Datasource.from_hf_dataset(f"test_hf_datasource_{uuid4()}", hf_dataset)
    assert datasource is not None
    assert datasource.length == len(hf_dataset)


def test_open_datasource(datasource):
    fetched_datasource = Datasource.open(datasource.name)
    assert fetched_datasource is not None
    assert fetched_datasource.name == datasource.name
    assert fetched_datasource.length == len(datasource)


def test_open_datasource_unauthenticated(unauthenticated_client, datasource):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            Datasource.open("test_datasource")


def test_open_datasource_invalid_input():
    with pytest.raises(ValueError, match=r"Invalid input:.*"):
        Datasource.open("not valid id")


def test_open_datasource_not_found():
    with pytest.raises(LookupError):
        Datasource.open(str(uuid4()))


def test_open_datasource_unauthorized(unauthorized_client, datasource):
    with unauthorized_client.use():
        with pytest.raises(LookupError):
            Datasource.open(datasource.id)


def test_all_datasources(datasource):
    datasources = Datasource.all()
    assert len(datasources) > 0
    assert any(datasource.name == datasource.name for datasource in datasources)


def test_all_datasources_unauthenticated(unauthenticated_client):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            Datasource.all()


def test_drop_datasource():
    Datasource.from_list("datasource_to_delete", [{"value": "test", "label": 0}])
    assert Datasource.exists("datasource_to_delete")
    Datasource.drop("datasource_to_delete")
    assert not Datasource.exists("datasource_to_delete")


def test_drop_datasource_unauthenticated(datasource, unauthenticated_client):
    with unauthenticated_client.use():
        with pytest.raises(ValueError, match="Invalid API key"):
            Datasource.drop(datasource.id)


def test_drop_datasource_not_found():
    with pytest.raises(LookupError):
        Datasource.drop(str(uuid4()))
    # ignores error if specified
    Datasource.drop(str(uuid4()), if_not_exists="ignore")


def test_drop_datasource_unauthorized(datasource, unauthorized_client):
    with unauthorized_client.use():
        with pytest.raises(LookupError):
            Datasource.drop(datasource.id)


def test_drop_datasource_invalid_input():
    with pytest.raises(ValueError, match=r"Invalid input:.*"):
        Datasource.drop("not valid id")


def test_from_list():
    # Test creating datasource from list of dictionaries
    data = [
        {"column1": 1, "column2": "a"},
        {"column1": 2, "column2": "b"},
        {"column1": 3, "column2": "c"},
    ]
    datasource = Datasource.from_list(f"test_list_{uuid4()}", data)
    assert datasource.name.startswith("test_list_")
    assert datasource.length == 3
    assert "column1" in datasource.columns
    assert "column2" in datasource.columns


def test_from_dict():
    # Test creating datasource from dictionary of columns
    data = {
        "column1": [1, 2, 3],
        "column2": ["a", "b", "c"],
    }
    datasource = Datasource.from_dict(f"test_dict_{uuid4()}", data)
    assert datasource.name.startswith("test_dict_")
    assert datasource.length == 3
    assert "column1" in datasource.columns
    assert "column2" in datasource.columns


def test_from_pandas():
    pd = pytest.importorskip("pandas")

    # Test creating datasource from pandas DataFrame
    df = pd.DataFrame(
        {
            "column1": [1, 2, 3],
            "column2": ["a", "b", "c"],
        }
    )
    datasource = Datasource.from_pandas(f"test_pandas_{uuid4()}", df)
    assert datasource.name.startswith("test_pandas_")
    assert datasource.length == 3
    assert "column1" in datasource.columns
    assert "column2" in datasource.columns


def test_from_arrow():
    pa = pytest.importorskip("pyarrow")

    # Test creating datasource from pyarrow Table
    table = pa.table(
        {
            "column1": [1, 2, 3],
            "column2": ["a", "b", "c"],
        }
    )
    datasource = Datasource.from_arrow(f"test_arrow_{uuid4()}", table)
    assert datasource.name.startswith("test_arrow_")
    assert datasource.length == 3
    assert "column1" in datasource.columns
    assert "column2" in datasource.columns


def test_from_list_already_exists():
    # Test the if_exists parameter with from_list
    data = [{"column1": 1, "column2": "a"}]
    name = f"test_list_exists_{uuid4()}"

    # Create the first datasource
    datasource1 = Datasource.from_list(name, data)
    assert datasource1.length == 1

    # Try to create again with if_exists="error" (should raise)
    with pytest.raises(ValueError):
        Datasource.from_list(name, data, if_exists="error")

    # Try to create again with if_exists="open" (should return existing)
    datasource2 = Datasource.from_list(name, data, if_exists="open")
    assert datasource2.id == datasource1.id
    assert datasource2.name == datasource1.name


def test_from_dict_already_exists():
    # Test the if_exists parameter with from_dict
    data = {"column1": [1], "column2": ["a"]}
    name = f"test_dict_exists_{uuid4()}"

    # Create the first datasource
    datasource1 = Datasource.from_dict(name, data)
    assert datasource1.length == 1

    # Try to create again with if_exists="error" (should raise)
    with pytest.raises(ValueError):
        Datasource.from_dict(name, data, if_exists="error")

    # Try to create again with if_exists="open" (should return existing)
    datasource2 = Datasource.from_dict(name, data, if_exists="open")
    assert datasource2.id == datasource1.id
    assert datasource2.name == datasource1.name


def test_from_pandas_already_exists():
    pd = pytest.importorskip("pandas")

    # Test the if_exists parameter with from_pandas
    df = pd.DataFrame({"column1": [1], "column2": ["a"]})
    name = f"test_pandas_exists_{uuid4()}"

    # Create the first datasource
    datasource1 = Datasource.from_pandas(name, df)
    assert datasource1.length == 1

    # Try to create again with if_exists="error" (should raise)
    with pytest.raises(ValueError):
        Datasource.from_pandas(name, df, if_exists="error")

    # Try to create again with if_exists="open" (should return existing)
    datasource2 = Datasource.from_pandas(name, df, if_exists="open")
    assert datasource2.id == datasource1.id
    assert datasource2.name == datasource1.name


def test_from_arrow_already_exists():
    pa = pytest.importorskip("pyarrow")

    # Test the if_exists parameter with from_arrow
    table = pa.table({"column1": [1], "column2": ["a"]})
    name = f"test_arrow_exists_{uuid4()}"

    # Create the first datasource
    datasource1 = Datasource.from_arrow(name, table)
    assert datasource1.length == 1

    # Try to create again with if_exists="error" (should raise)
    with pytest.raises(ValueError):
        Datasource.from_arrow(name, table, if_exists="error")

    # Try to create again with if_exists="open" (should return existing)
    datasource2 = Datasource.from_arrow(name, table, if_exists="open")
    assert datasource2.id == datasource1.id
    assert datasource2.name == datasource1.name


def test_from_disk_csv():
    # Test creating datasource from CSV file
    with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
        f.write("column1,column2\n1,a\n2,b\n3,c")
        f.flush()

        try:
            datasource = Datasource.from_disk(f"test_csv_{uuid4()}", f.name)
            assert datasource.length == 3
            assert "column1" in datasource.columns
            assert "column2" in datasource.columns
        finally:
            os.unlink(f.name)


def test_from_disk_json():
    # Test creating datasource from JSON file
    import json

    data = [{"column1": 1, "column2": "a"}, {"column1": 2, "column2": "b"}]

    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
        json.dump(data, f)
        f.flush()

        try:
            datasource = Datasource.from_disk(f"test_json_{uuid4()}", f.name)
            assert datasource.length == 2
            assert "column1" in datasource.columns
            assert "column2" in datasource.columns
        finally:
            os.unlink(f.name)


def test_from_disk_already_exists():
    # Test the if_exists parameter with from_disk
    with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
        f.write("column1,column2\n1,a")
        f.flush()

        try:
            name = f"test_disk_exists_{uuid4()}"

            # Create the first datasource
            datasource1 = Datasource.from_disk(name, f.name)
            assert datasource1.length == 1

            # Try to create again with if_exists="error" (should raise)
            with pytest.raises(ValueError):
                Datasource.from_disk(name, f.name, if_exists="error")

            # Try to create again with if_exists="open" (should return existing)
            datasource2 = Datasource.from_disk(name, f.name, if_exists="open")
            assert datasource2.id == datasource1.id
            assert datasource2.name == datasource1.name
        finally:
            os.unlink(f.name)


def test_query_datasource_rows():
    """Test querying rows from a datasource with pagination and shuffle."""
    # Create a new dataset with 5 entries for testing
    test_data = [{"id": i, "name": f"item_{i}"} for i in range(5)]
    datasource = Datasource.from_list(name="test_query_datasource", data=test_data)

    # Test basic query
    rows = datasource.query(limit=3)
    assert len(rows) == 3
    assert all(isinstance(row, dict) for row in rows)

    # Test offset
    offset_rows = datasource.query(offset=2, limit=2)
    assert len(offset_rows) == 2
    assert offset_rows[0]["id"] == 2

    # Test shuffle
    shuffled_rows = datasource.query(limit=5, shuffle=True)
    assert len(shuffled_rows) == 5
    assert not all(row["id"] == i for i, row in enumerate(shuffled_rows))

    # Test shuffle with seed
    assert datasource.query(limit=5, shuffle=True, shuffle_seed=42) == datasource.query(
        limit=5, shuffle=True, shuffle_seed=42
    )


def test_query_datasource_with_filters():
    """Test querying datasource rows with various filter operators."""
    # Create a datasource with test data
    test_data = [
        {"name": "Alice", "age": 25, "city": "New York", "score": 85.5},
        {"name": "Bob", "age": 30, "city": "San Francisco", "score": 90.0},
        {"name": "Charlie", "age": 35, "city": "Chicago", "score": 75.5},
        {"name": "Diana", "age": 28, "city": "Boston", "score": 88.0},
        {"name": "Eve", "age": 32, "city": "New York", "score": 92.0},
    ]
    datasource = Datasource.from_list(name=f"test_filter_datasource_{uuid4()}", data=test_data)

    # Test == operator
    rows = datasource.query(filters=[("city", "==", "New York")])
    assert len(rows) == 2
    assert all(row["city"] == "New York" for row in rows)

    # Test > operator
    rows = datasource.query(filters=[("age", ">", 30)])
    assert len(rows) == 2
    assert all(row["age"] > 30 for row in rows)

    # Test >= operator
    rows = datasource.query(filters=[("score", ">=", 88.0)])
    assert len(rows) == 3
    assert all(row["score"] >= 88.0 for row in rows)

    # Test < operator
    rows = datasource.query(filters=[("age", "<", 30)])
    assert len(rows) == 2
    assert all(row["age"] < 30 for row in rows)

    # Test in operator
    rows = datasource.query(filters=[("city", "in", ["New York", "Boston"])])
    assert len(rows) == 3
    assert all(row["city"] in ["New York", "Boston"] for row in rows)

    # Test not in operator
    rows = datasource.query(filters=[("city", "not in", ["New York", "Boston"])])
    assert len(rows) == 2
    assert all(row["city"] not in ["New York", "Boston"] for row in rows)

    # Test like operator
    rows = datasource.query(filters=[("name", "like", "li")])
    assert len(rows) == 2
    assert all("li" in row["name"].lower() for row in rows)

    # Test multiple filters (AND logic)
    rows = datasource.query(filters=[("city", "==", "New York"), ("age", ">", 26)])
    assert len(rows) == 1
    assert rows[0]["name"] == "Eve"

    # Test filter with pagination
    rows = datasource.query(filters=[("age", ">=", 28)], limit=2, offset=1)
    assert len(rows) == 2


def test_query_datasource_with_none_filters():
    """Test filtering for None values."""
    test_data = [
        {"name": "Alice", "age": 25, "label": "A"},
        {"name": "Bob", "age": 30, "label": None},
        {"name": "Charlie", "age": 35, "label": "C"},
        {"name": "Diana", "age": None, "label": "D"},
        {"name": "Eve", "age": 32, "label": None},
    ]
    datasource = Datasource.from_list(name=f"test_none_filter_{uuid4()}", data=test_data)

    # Test == None
    rows = datasource.query(filters=[("label", "==", None)])
    assert len(rows) == 2
    assert all(row["label"] is None for row in rows)

    # Test != None
    rows = datasource.query(filters=[("label", "!=", None)])
    assert len(rows) == 3
    assert all(row["label"] is not None for row in rows)

    # Test that None values are excluded from comparison operators
    rows = datasource.query(filters=[("age", ">", 25)])
    assert len(rows) == 3
    assert all(row["age"] is not None and row["age"] > 25 for row in rows)


def test_query_datasource_filter_invalid_column():
    """Test that querying with an invalid column raises an error."""
    test_data = [{"name": "Alice", "age": 25}]
    datasource = Datasource.from_list(name=f"test_invalid_filter_{uuid4()}", data=test_data)

    with pytest.raises(ValueError):
        datasource.query(filters=[("invalid_column", "==", "test")])


def test_to_list(data, datasource):
    assert datasource.to_list() == data


def test_download_datasource_hf_dataset(data, datasource: Datasource):
    Dataset = pytest.importorskip("datasets").Dataset
    with tempfile.TemporaryDirectory() as temp_dir:
        datasource.download(temp_dir, file_type="hf_dataset")
        downloaded_hf_dataset_dir = f"{temp_dir}/{datasource.name}"
        assert os.path.exists(downloaded_hf_dataset_dir)
        assert os.path.isdir(downloaded_hf_dataset_dir)
        assert not os.path.exists(f"{downloaded_hf_dataset_dir}.zip")
        dataset_from_downloaded_hf_dataset = Dataset.load_from_disk(downloaded_hf_dataset_dir)
        assert set(dataset_from_downloaded_hf_dataset.column_names) == set(data[0].keys())
        assert dataset_from_downloaded_hf_dataset.to_list() == data


def test_download_datasource_json(data, datasource: Datasource):
    with tempfile.TemporaryDirectory() as temp_dir:
        datasource.download(temp_dir, file_type="json")
        downloaded_json_file = f"{temp_dir}/{datasource.name}.json"
        assert os.path.exists(downloaded_json_file)
        with open(downloaded_json_file, "r") as f:
            content = json.load(f)
            assert content == data


def test_download_datasource_csv(data, datasource: Datasource):
    with tempfile.TemporaryDirectory() as temp_dir:
        datasource.download(temp_dir, file_type="csv")
        downloaded_csv_file = f"{temp_dir}/{datasource.name}.csv"
        assert os.path.exists(downloaded_csv_file)
        with open(downloaded_csv_file, "r") as f:
            reader = csv.DictReader(f)
            rows = list(reader)
        assert set(rows[0].keys()) == set(data[0].keys())
        assert len(rows) == len(data)
        # Compare all columns, handling CSV's string representation
        for i, (csv_row, expected_row) in enumerate(zip(rows, data)):
            for col in expected_row.keys():
                csv_val = csv_row[col]
                expected_val = expected_row[col]
                # CSV represents None as empty string
                if expected_val is None:
                    assert csv_val == "", f"Row {i}, column {col}: expected empty string, got {csv_val!r}"
                elif isinstance(expected_val, float):
                    assert float(csv_val) == pytest.approx(expected_val), f"Row {i}, column {col}"
                elif isinstance(expected_val, int):
                    assert int(csv_val) == expected_val, f"Row {i}, column {col}"
                else:
                    assert csv_val == str(expected_val), f"Row {i}, column {col}"


def test_datasource_repr():
    ds = Datasource(
        {  # type: ignore[typeddict-item]
            "id": "ds-id",
            "org_id": "org-id",
            "name": "my_datasource",
            "description": None,
            "storage_path": "",
            "length": 100,
            "columns": [
                {"name": "value", "type": "STRING"},
                {"name": "label", "type": "ENUM", "enum_options": ["soup", "cats"]},
                {"name": "score", "type": "FLOAT"},
            ],
            "created_at": "2025-01-01T00:00:00",
            "updated_at": "2025-01-01T00:00:00",
        }
    )
    assert repr(ds) == dedent(
        """\
        Datasource({
            name: 'my_datasource',
            length: 100,
            columns: ({
                value: str
                label: enum('soup', 'cats')
                score: float
            })
        })"""
    )
