from datetime import datetime
from typing import Any

import pytest

from dvcx.sql.types import (
    JSON,
    Array,
    Boolean,
    DateTime,
    Float,
    Float32,
    Float64,
    Int,
    String,
)
from tests.utils import DEFAULT_TREE, TARRED_TREE, create_tar_dataset

COMPLEX_TREE: dict[str, Any] = {
    **TARRED_TREE,
    **DEFAULT_TREE,
    "nested": {"dir": {"path": {"abc.txt": "abc"}}},
}


@pytest.mark.parametrize("tree", [COMPLEX_TREE], indirect=True)
def test_dir_expansion(cloud_test_catalog, version_aware, cloud_type):
    has_version = version_aware or cloud_type == "gcs"

    ctc = cloud_test_catalog
    catalog = ctc.catalog
    src_uri = ctc.src_uri
    if cloud_type == "file":
        # we don't want to index things in parent directory
        src_uri += "/"

    ds = create_tar_dataset(catalog, ctc.src_uri, "ds2")
    dataset = catalog.get_dataset(ds.name)
    st = catalog.warehouse.clone()
    q = st.dataset_rows(dataset).dir_expansion()
    columns = (
        "id",
        "vtype",
        "is_dir",
        "source",
        "parent",
        "name",
        "version",
        "location",
    )
    result = [dict(zip(columns, r)) for r in st.db.execute(q)]
    to_compare = [
        (r["parent"], r["name"], r["vtype"], r["is_dir"], r["version"] != "")
        for r in result
    ]

    assert all(r["source"] == ctc.storage_uri for r in result)
    if cloud_type == "file":
        prefix = ctc.partial_path + "/"
        prefix_root = ctc.partial_path
    else:
        prefix = ""
        prefix_root = ""

    # Note, we have both a file and a directory entry for expanded tar files
    expected = [
        (f"{prefix_root}", "animals.tar", "", 0, has_version),
        (f"{prefix_root}", "animals.tar", "", 1, False),
        (f"{prefix_root}", "cats", "", 1, False),
        (f"{prefix_root}", "description", "", 0, has_version),
        (f"{prefix_root}", "dogs", "", 1, False),
        (f"{prefix_root}", "nested", "", 1, False),
        (f"{prefix}animals.tar", "cats", "", 1, False),
        (f"{prefix}animals.tar", "description", "tar", 0, False),
        (f"{prefix}animals.tar", "dogs", "", 1, False),
        (f"{prefix}animals.tar/cats", "cat1", "tar", 0, False),
        (f"{prefix}animals.tar/cats", "cat2", "tar", 0, False),
        (f"{prefix}animals.tar/dogs", "dog1", "tar", 0, False),
        (f"{prefix}animals.tar/dogs", "dog2", "tar", 0, False),
        (f"{prefix}animals.tar/dogs", "dog3", "tar", 0, False),
        (f"{prefix}animals.tar/dogs", "others", "", 1, False),
        (f"{prefix}animals.tar/dogs/others", "dog4", "tar", 0, False),
        (f"{prefix}cats", "cat1", "", 0, has_version),
        (f"{prefix}cats", "cat2", "", 0, has_version),
        (f"{prefix}dogs", "dog1", "", 0, has_version),
        (f"{prefix}dogs", "dog2", "", 0, has_version),
        (f"{prefix}dogs", "dog3", "", 0, has_version),
        (f"{prefix}dogs", "others", "", 1, False),
        (f"{prefix}dogs/others", "dog4", "", 0, has_version),
        (f"{prefix}nested", "dir", "", 1, False),
        (f"{prefix}nested/dir", "path", "", 1, False),
        (f"{prefix}nested/dir/path", "abc.txt", "", 0, has_version),
    ]

    if cloud_type == "file":
        # since with file listing, parent is relative path to the root of FS as
        # storage uri is the root of FS, we need to add dirs to the root
        prefix_split = prefix.split("/")
        expected = [
            ("/".join(prefix_split[:i]), prefix_split[i], "", 1, False)
            for i in range(len(prefix_split) - 1)
        ] + expected

    assert to_compare == expected


@pytest.mark.parametrize(
    "cloud_type,version_aware",
    [("s3", True)],
    indirect=True,
)
def test_convert_type(cloud_test_catalog):
    ctc = cloud_test_catalog
    catalog = ctc.catalog
    warehouse = catalog.warehouse
    now = datetime.now()

    # convert int to float
    for f in [Float, Float32, Float64]:
        converted = warehouse.convert_type(1, f())
        assert converted == 1.0
        assert isinstance(converted, float)

    # types match, nothing to convert
    assert warehouse.convert_type(1, Int()) == 1
    assert warehouse.convert_type(1.5, Float()) == 1.5
    assert warehouse.convert_type(True, Boolean()) is True
    assert warehouse.convert_type("s", String()) == "s"
    assert warehouse.convert_type(now, DateTime()) == now
    assert warehouse.convert_type([1, 2], Array(Int)) == [1, 2]
    assert warehouse.convert_type([1.5, 2.5], Array(Float)) == [1.5, 2.5]
    assert warehouse.convert_type(["a", "b"], Array(String)) == ["a", "b"]
    assert warehouse.convert_type([[1, 2], [3, 4]], Array(Array(Int))) == [
        [1, 2],
        [3, 4],
    ]

    # JSON Tests
    assert warehouse.convert_type('{"a": 1}', JSON()) == '{"a": 1}'
    assert warehouse.convert_type({"a": 1}, JSON()) == '{"a": 1}'
    assert warehouse.convert_type([{"a": 1}], JSON()) == '[{"a": 1}]'
    with pytest.raises(ValueError):
        warehouse.convert_type(0.5, JSON())

    # convert array to compatible type
    converted = warehouse.convert_type([1, 2], Array(Float))
    assert converted == [1.0, 2.0]
    assert all(isinstance(c, float) for c in converted)

    # convert nested array to compatible type
    converted = warehouse.convert_type([[1, 2], [3, 4]], Array(Array(Float)))
    assert converted == [[1.0, 2.0], [3.0, 4.0]]
    assert all(isinstance(c, float) for c in converted[0])
    assert all(isinstance(c, float) for c in converted[1])

    # error, float to int
    with pytest.raises(ValueError):
        warehouse.convert_type(1.5, Int())

    # error, float to int in list
    with pytest.raises(ValueError):
        warehouse.convert_type([1.5, 1], Array(Int))
