#   Copyright 2021 Modelyst LLC
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

from random import shuffle
from typing import Optional, cast

import pytest
from sqlalchemy.future import Engine
from sqlmodel import Session, func, select

import tests.example.entities as entities
from dbgen.core.args import Const
from dbgen.core.entity import Entity
from dbgen.core.func import Import
from dbgen.core.generator import Generator
from dbgen.core.metadata import RunEntity
from dbgen.core.node.load import Load
from dbgen.core.node.query import BaseQuery
from dbgen.core.node.transforms import PyBlock


def transform_func(x):
    return f"{x}-child"


@pytest.fixture(scope='function')
def basic_generator() -> Generator:
    Parent = entities.Parent
    Child = entities.Child
    select_stmt = select(Parent.label)
    query = BaseQuery.from_select_statement(select_stmt)
    assert isinstance(query.hash, str)
    pyblock = PyBlock(function=transform_func, inputs=[query["label"]], outputs=["newnames"])

    load = Child.load(insert=True, label=pyblock["newnames"], type=Const("child_type"))
    assert isinstance(load.hash, str)
    gen = Generator(name="test", extract=query, transforms=[pyblock], loads=[load])
    return gen


def test_basic_graph_sort(basic_generator: Generator):
    """Ensure a simple Query->PyBlock->Load is sorted correctly."""
    graph = basic_generator._computational_graph()
    assert len(graph) == 3
    sorted_nodes = basic_generator._sort_graph()
    query, transform, load = sorted_nodes
    assert isinstance(query, BaseQuery)
    assert isinstance(transform, PyBlock)
    assert isinstance(load, Load)


def test_basic_graph_in_place(basic_generator: Generator):
    """Ensure that changes to the output of ._sort_graph() are in place and affect the generator as well."""
    query, transform, load = basic_generator._sort_graph()
    assert isinstance(load, Load)
    load.run({transform.hash: {"newnames": ("1", "2")}})
    assert load._output == basic_generator._sorted_loads()[0]._output
    assert isinstance(query, BaseQuery)
    query.outputs.append("test")
    assert basic_generator.extract == query
    assert isinstance(transform, PyBlock)
    import_to_add = Import(lib="numpy", lib_alias="np")
    transform.env.imports.append(import_to_add)
    assert basic_generator.transforms[0] == transform
    assert basic_generator.transforms[0].env.imports == [import_to_add]


def test_sorted_loads():
    """Shuffle around the loads and make sure sorted_loads still works."""
    val = Const("test")
    gp_load = entities.GrandParent.load(label=val, type=val)
    u_load = entities.Parent.load(label=val, type=Const("uncle"), grand_parent_id=gp_load)
    p_load = entities.Parent.load(label=val, type=val, grand_parent_id=gp_load)
    c_load = entities.Child.load(label=val, type=val, parent_id=p_load, uncle_id=u_load)
    loads = [gp_load, c_load, p_load, u_load]
    for _ in range(10):
        shuffle(loads)
        gen = Generator(name="test", loads=loads)
        assert gen._sorted_loads() == [
            gp_load,
            *sorted((u_load, p_load), key=lambda x: x.hash),
            c_load,
        ]


@pytest.mark.skip
def test_no_extractor(sql_engine: Engine, raw_connection):
    """Shuffle around the loads and make sure sorted_loads still works."""
    entities.Parent.metadata.create_all(sql_engine)
    pyblock = PyBlock(function=transform_func, inputs=[Const("test")], outputs=["newnames"])
    p_load = entities.GrandParent.load(insert=True, label=pyblock["newnames"], type=Const("gp_type"))
    gen = Generator(name="test", transforms=[pyblock], loads=[p_load])
    gen.run(sql_engine)

    with Session(sql_engine) as session:
        session = cast(Session, session)
        statement = select(entities.GrandParent).where(entities.GrandParent.label == "test-child")
        result = session.exec(statement)
        assert result.one()


@pytest.mark.database
def test_dumb_extractor(connection, sql_engine, recreate_meta):
    class User(Entity, table=True):
        __identifying__ = {"label"}
        label: Optional[str]
        new_label: Optional[str] = None

    User.metadata.create_all(connection)
    num_users = 100
    sess = Session(connection)
    users = [User(label=f"user_{i}") for i in range(num_users)]
    user_le = User._get_load_entity()
    for user in users:
        user.id = user_le._get_hash(user.dict())
        sess.add(user)
    count = sess.exec(select(func.count(User.id))).one()
    assert count == num_users
    connection.commit()
    statement = select(User.id, User.label)
    query = BaseQuery.from_select_statement(statement)
    assert query.length(connection=connection) == num_users
    pyblock = PyBlock(function=transform_func, inputs=[query["label"]])
    u_load = User.load(user=query["id"], new_label=pyblock["out"])
    run = RunEntity()
    sess.add(run)
    sess.commit()
    sess.refresh(run)
    gen = Generator(
        name="test",
        extract=query,
        transforms=[pyblock],
        loads=[u_load],
        batch_size=10000,
    )
    connection.commit()
    gen.run(sql_engine, sql_engine, run_id=run.id, ordering=0)

tests.example.entities.Child.load dbgen.core.args.Const random.shuffle sqlmodel.func.count tests.example.entities.Parent.load tests.example.entities.GrandParent.load dbgen.core.metadata.RunEntity typing.cast dbgen.core.node.query.BaseQuery.from_select_statement tests.example.entities.Parent.metadata.create_all sqlmodel.select pytest.fixture sqlmodel.Session dbgen.core.func.Import dbgen.core.node.transforms.PyBlock dbgen.core.generator.Generator
[((1168, 1200), 'pytest.fixture', 'pytest.fixture', ([], {'scope': '"""function"""'}), "(scope='function')\n", (1182, 1200), False, 'import pytest\n'), ((1311, 1331), 'sqlmodel.select', 'select', (['Parent.label'], {}), '(Parent.label)\n', (1317, 1331), False, 'from sqlmodel import Session, func, select\n'), ((1344, 1388), 'dbgen.core.node.query.BaseQuery.from_select_statement', 'BaseQuery.from_select_statement', (['select_stmt'], {}), '(select_stmt)\n', (1375, 1388), False, 'from dbgen.core.node.query import BaseQuery\n'), ((1442, 1521), 'dbgen.core.node.transforms.PyBlock', 'PyBlock', ([], {'function': 'transform_func', 'inputs': "[query['label']]", 'outputs': "['newnames']"}), "(function=transform_func, inputs=[query['label']], outputs=['newnames'])\n", (1449, 1521), False, 'from dbgen.core.node.transforms import PyBlock\n'), ((1659, 1732), 'dbgen.core.generator.Generator', 'Generator', ([], {'name': '"""test"""', 'extract': 'query', 'transforms': '[pyblock]', 'loads': '[load]'}), "(name='test', extract=query, transforms=[pyblock], loads=[load])\n", (1668, 1732), False, 'from dbgen.core.generator import Generator\n'), ((2727, 2762), 'dbgen.core.func.Import', 'Import', ([], {'lib': '"""numpy"""', 'lib_alias': '"""np"""'}), "(lib='numpy', lib_alias='np')\n", (2733, 2762), False, 'from dbgen.core.func import Import\n'), ((3049, 3062), 'dbgen.core.args.Const', 'Const', (['"""test"""'], {}), "('test')\n", (3054, 3062), False, 'from dbgen.core.args import Const\n'), ((3077, 3123), 'tests.example.entities.GrandParent.load', 'entities.GrandParent.load', ([], {'label': 'val', 'type': 'val'}), '(label=val, type=val)\n', (3102, 3123), True, 'import tests.example.entities as entities\n'), ((3228, 3294), 'tests.example.entities.Parent.load', 'entities.Parent.load', ([], {'label': 'val', 'type': 'val', 'grand_parent_id': 'gp_load'}), '(label=val, type=val, grand_parent_id=gp_load)\n', (3248, 3294), True, 'import tests.example.entities as entities\n'), ((3308, 3383), 'tests.example.entities.Child.load', 'entities.Child.load', ([], {'label': 'val', 'type': 'val', 'parent_id': 'p_load', 'uncle_id': 'u_load'}), '(label=val, type=val, parent_id=p_load, uncle_id=u_load)\n', (3327, 3383), True, 'import tests.example.entities as entities\n'), ((3837, 3884), 'tests.example.entities.Parent.metadata.create_all', 'entities.Parent.metadata.create_all', (['sql_engine'], {}), '(sql_engine)\n', (3872, 3884), True, 'import tests.example.entities as entities\n'), ((4090, 4150), 'dbgen.core.generator.Generator', 'Generator', ([], {'name': '"""test"""', 'transforms': '[pyblock]', 'loads': '[p_load]'}), "(name='test', transforms=[pyblock], loads=[p_load])\n", (4099, 4150), False, 'from dbgen.core.generator import Generator\n'), ((4728, 4747), 'sqlmodel.Session', 'Session', (['connection'], {}), '(connection)\n', (4735, 4747), False, 'from sqlmodel import Session, func, select\n'), ((5072, 5099), 'sqlmodel.select', 'select', (['User.id', 'User.label'], {}), '(User.id, User.label)\n', (5078, 5099), False, 'from sqlmodel import Session, func, select\n'), ((5112, 5154), 'dbgen.core.node.query.BaseQuery.from_select_statement', 'BaseQuery.from_select_statement', (['statement'], {}), '(statement)\n', (5143, 5154), False, 'from dbgen.core.node.query import BaseQuery\n'), ((5229, 5286), 'dbgen.core.node.transforms.PyBlock', 'PyBlock', ([], {'function': 'transform_func', 'inputs': "[query['label']]"}), "(function=transform_func, inputs=[query['label']])\n", (5236, 5286), False, 'from dbgen.core.node.transforms import PyBlock\n'), ((5364, 5375), 'dbgen.core.metadata.RunEntity', 'RunEntity', ([], {}), '()\n', (5373, 5375), False, 'from dbgen.core.metadata import RunEntity\n'), ((5444, 5541), 'dbgen.core.generator.Generator', 'Generator', ([], {'name': '"""test"""', 'extract': 'query', 'transforms': '[pyblock]', 'loads': '[u_load]', 'batch_size': '(10000)'}), "(name='test', extract=query, transforms=[pyblock], loads=[u_load],\n    batch_size=10000)\n", (5453, 5541), False, 'from dbgen.core.generator import Generator\n'), ((3462, 3476), 'random.shuffle', 'shuffle', (['loads'], {}), '(loads)\n', (3469, 3476), False, 'from random import shuffle\n'), ((3491, 3526), 'dbgen.core.generator.Generator', 'Generator', ([], {'name': '"""test"""', 'loads': 'loads'}), "(name='test', loads=loads)\n", (3500, 3526), False, 'from dbgen.core.generator import Generator\n'), ((4185, 4204), 'sqlmodel.Session', 'Session', (['sql_engine'], {}), '(sql_engine)\n', (4192, 4204), False, 'from sqlmodel import Session, func, select\n'), ((4235, 4257), 'typing.cast', 'cast', (['Session', 'session'], {}), '(Session, session)\n', (4239, 4257), False, 'from typing import Optional, cast\n'), ((1590, 1609), 'dbgen.core.args.Const', 'Const', (['"""child_type"""'], {}), "('child_type')\n", (1595, 1609), False, 'from dbgen.core.args import Const\n'), ((3174, 3188), 'dbgen.core.args.Const', 'Const', (['"""uncle"""'], {}), "('uncle')\n", (3179, 3188), False, 'from dbgen.core.args import Const\n'), ((4062, 4078), 'dbgen.core.args.Const', 'Const', (['"""gp_type"""'], {}), "('gp_type')\n", (4067, 4078), False, 'from dbgen.core.args import Const\n'), ((3940, 3953), 'dbgen.core.args.Const', 'Const', (['"""test"""'], {}), "('test')\n", (3945, 3953), False, 'from dbgen.core.args import Const\n'), ((4278, 4306), 'sqlmodel.select', 'select', (['entities.GrandParent'], {}), '(entities.GrandParent)\n', (4284, 4306), False, 'from sqlmodel import Session, func, select\n'), ((4974, 4993), 'sqlmodel.func.count', 'func.count', (['User.id'], {}), '(User.id)\n', (4984, 4993), False, 'from sqlmodel import Session, func, select\n')]