#   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 typing import Optional
from uuid import UUID

import sqlalchemy.types as types
from sqlalchemy import Column
from sqlalchemy.orm import registry
from sqlmodel import Field, select

from dbgen.core.decorators import transform
from dbgen.core.entity import Entity
from dbgen.core.generator import Generator
from dbgen.core.model import Model
from dbgen.core.node.extract import Extract
from dbgen.core.node.query import Query

my_registry = registry()


class Parent(Entity, registry=my_registry, table=True):
    __identifying__ = {"label"}
    label: str
    myColumn: Optional[dict] = Field(None, sa_column=Column(types.JSON()))


class Child(Entity, registry=my_registry, table=True):
    __identifying__ = {"label", "parent_id"}
    label: str
    new_col: str = "test"
    parent_id: Optional[UUID] = Field(None, foreign_key="public.parent.id")


class CustomExtractor(Extract):
    n: int = 1000

    def extract(self):
        for i in range(self.n):
            yield {'out': str(i)}

    def length(self, **_):
        return self.n


@transform
def failing_func():
    raise ValueError("Failed")


@transform
def inputs_skipped():
    from dbgen.exceptions import DBgenSkipException

    raise DBgenSkipException(msg="Skip!")


def make_model():
    with Model(name='new_api', registry=my_registry) as model:
        with Generator('add_parent'):
            new_extract = CustomExtractor(n=1000)
            Parent.load(insert=True, label=new_extract["out"], validation='strict', myColumn={'a': 1})

        with Generator('add_parents_v2'):
            Parent.load(insert=True, label="parentier")

        with Generator('add_parents_v3'):
            Parent.load(insert=True, label="parent")

        @transform
        def concise_func(label: str) -> str:
            return f"{label}-test"

        with Generator('add_child'):
            query = Query(select(Parent.id, Parent.label))
            parent_id, parent_label = query.results()
            concise_pyblock = concise_func(query["label"])
            Child.load(insert=True, label=concise_pyblock.results(), parent_id=query["id"])

        with Generator('failing_gen'):
            failing_func()

        with Generator('skip_gen'):
            inputs_skipped()
    return model

sqlalchemy.types.JSON sqlalchemy.orm.registry dbgen.exceptions.DBgenSkipException dbgen.core.model.Model sqlmodel.select dbgen.core.generator.Generator sqlmodel.Field
[((1041, 1051), 'sqlalchemy.orm.registry', 'registry', ([], {}), '()\n', (1049, 1051), False, 'from sqlalchemy.orm import registry\n'), ((1407, 1450), 'sqlmodel.Field', 'Field', (['None'], {'foreign_key': '"""public.parent.id"""'}), "(None, foreign_key='public.parent.id')\n", (1412, 1450), False, 'from sqlmodel import Field, select\n'), ((1805, 1836), 'dbgen.exceptions.DBgenSkipException', 'DBgenSkipException', ([], {'msg': '"""Skip!"""'}), "(msg='Skip!')\n", (1823, 1836), False, 'from dbgen.exceptions import DBgenSkipException\n'), ((1866, 1909), 'dbgen.core.model.Model', 'Model', ([], {'name': '"""new_api"""', 'registry': 'my_registry'}), "(name='new_api', registry=my_registry)\n", (1871, 1909), False, 'from dbgen.core.model import Model\n'), ((1933, 1956), 'dbgen.core.generator.Generator', 'Generator', (['"""add_parent"""'], {}), "('add_parent')\n", (1942, 1956), False, 'from dbgen.core.generator import Generator\n'), ((2125, 2152), 'dbgen.core.generator.Generator', 'Generator', (['"""add_parents_v2"""'], {}), "('add_parents_v2')\n", (2134, 2152), False, 'from dbgen.core.generator import Generator\n'), ((2224, 2251), 'dbgen.core.generator.Generator', 'Generator', (['"""add_parents_v3"""'], {}), "('add_parents_v3')\n", (2233, 2251), False, 'from dbgen.core.generator import Generator\n'), ((2420, 2442), 'dbgen.core.generator.Generator', 'Generator', (['"""add_child"""'], {}), "('add_child')\n", (2429, 2442), False, 'from dbgen.core.generator import Generator\n'), ((2722, 2746), 'dbgen.core.generator.Generator', 'Generator', (['"""failing_gen"""'], {}), "('failing_gen')\n", (2731, 2746), False, 'from dbgen.core.generator import Generator\n'), ((2789, 2810), 'dbgen.core.generator.Generator', 'Generator', (['"""skip_gen"""'], {}), "('skip_gen')\n", (2798, 2810), False, 'from dbgen.core.generator import Generator\n'), ((1217, 1229), 'sqlalchemy.types.JSON', 'types.JSON', ([], {}), '()\n', (1227, 1229), True, 'import sqlalchemy.types as types\n'), ((2470, 2501), 'sqlmodel.select', 'select', (['Parent.id', 'Parent.label'], {}), '(Parent.id, Parent.label)\n', (2476, 2501), False, 'from sqlmodel import Field, select\n')]