#   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 List, Optional
from uuid import UUID, uuid4

import requests
from pydantic import HttpUrl
from pydantic.tools import parse_obj_as
from sqlalchemy.sql.expression import text
from sqlmodel import Session, select

from dbgen import Const, Entity, Extract, Generator, Model, Query
from dbgen.configuration import config, get_engines
from dbgen.core.node.transforms import PyBlock


class CustomJsonExtract(Extract):
    url: HttpUrl = parse_obj_as(HttpUrl, 'https://jsonplaceholder.typicode.com/posts')
    outputs: List[str] = ['out', 'uuid']

    def setup(self, **_):
        self._response = requests.get(self.url).json()
        self._response += [{}]

    def extract(self):
        for row in self._response:
            row['uuid'] = uuid4()
            yield {'out': row, 'uuid': row['uuid']}

    def length(self, **_):
        return len(self._response)


class JSONEntityBase(Entity):
    __tablename__ = 'json_entity'
    tags: Optional[List[dict]]
    my_uuid: Optional[UUID]


class JSONEntity(JSONEntityBase, table=True):
    __tablename__ = 'json_entity'
    __identifying__ = {
        'json_val',
    }
    json_val: Optional[dict]


model = Model(name='test_json')
load_json = Generator(name='load_json', loads=[JSONEntity.load(insert=True, json_val=Const({}))])
model.add_gen(load_json)

extract = CustomJsonExtract()
load = JSONEntity.load(insert=True, json_val=extract['out'], my_uuid=extract['uuid'])
load_http_json = Generator(name='load_http_json', extract=extract, loads=[load])
model.add_gen(load_http_json)

query = Query(select(JSONEntity.id, JSONEntity.json_val.op('->')(text("'title'")).label('title')))


def get_title_words(text: str):
    if text:
        return [{'word': word} for word in text.split(' ')]


pb = PyBlock(function=get_title_words, inputs=[query['title']])
load = JSONEntity.load(json_entity=query['id'], tags=pb['out'])
add_tags = Generator(name='add_tags', extract=query, transforms=[pb], loads=[load])
model.add_gen(add_tags)

if __name__ == '__main__':
    main_engine, _ = get_engines(config)
    with Session(main_engine) as session:
        json_entity = session.exec(select(JSONEntity)).first()
        if json_entity:
            print(json_entity.dict().keys())
            print(json_entity.json_val)

dbgen.Generator requests.get dbgen.core.node.transforms.PyBlock sqlalchemy.sql.expression.text pydantic.tools.parse_obj_as sqlmodel.Session uuid.uuid4 dbgen.configuration.get_engines sqlmodel.select dbgen.Model dbgen.Const sqlalchemy.sql.expression.text.split
[((1773, 1796), 'dbgen.Model', 'Model', ([], {'name': '"""test_json"""'}), "(name='test_json')\n", (1778, 1796), False, 'from dbgen import Const, Entity, Extract, Generator, Model, Query\n'), ((2054, 2117), 'dbgen.Generator', 'Generator', ([], {'name': '"""load_http_json"""', 'extract': 'extract', 'loads': '[load]'}), "(name='load_http_json', extract=extract, loads=[load])\n", (2063, 2117), False, 'from dbgen import Const, Entity, Extract, Generator, Model, Query\n'), ((2362, 2420), 'dbgen.core.node.transforms.PyBlock', 'PyBlock', ([], {'function': 'get_title_words', 'inputs': "[query['title']]"}), "(function=get_title_words, inputs=[query['title']])\n", (2369, 2420), False, 'from dbgen.core.node.transforms import PyBlock\n'), ((2496, 2568), 'dbgen.Generator', 'Generator', ([], {'name': '"""add_tags"""', 'extract': 'query', 'transforms': '[pb]', 'loads': '[load]'}), "(name='add_tags', extract=query, transforms=[pb], loads=[load])\n", (2505, 2568), False, 'from dbgen import Const, Entity, Extract, Generator, Model, Query\n'), ((1047, 1114), 'pydantic.tools.parse_obj_as', 'parse_obj_as', (['HttpUrl', '"""https://jsonplaceholder.typicode.com/posts"""'], {}), "(HttpUrl, 'https://jsonplaceholder.typicode.com/posts')\n", (1059, 1114), False, 'from pydantic.tools import parse_obj_as\n'), ((2642, 2661), 'dbgen.configuration.get_engines', 'get_engines', (['config'], {}), '(config)\n', (2653, 2661), False, 'from dbgen.configuration import config, get_engines\n'), ((2671, 2691), 'sqlmodel.Session', 'Session', (['main_engine'], {}), '(main_engine)\n', (2678, 2691), False, 'from sqlmodel import Session, select\n'), ((1354, 1361), 'uuid.uuid4', 'uuid4', ([], {}), '()\n', (1359, 1361), False, 'from uuid import UUID, uuid4\n'), ((1208, 1230), 'requests.get', 'requests.get', (['self.url'], {}), '(self.url)\n', (1220, 1230), False, 'import requests\n'), ((2338, 2353), 'sqlalchemy.sql.expression.text.split', 'text.split', (['""" """'], {}), "(' ')\n", (2348, 2353), False, 'from sqlalchemy.sql.expression import text\n'), ((1882, 1891), 'dbgen.Const', 'Const', (['{}'], {}), '({})\n', (1887, 1891), False, 'from dbgen import Const, Entity, Extract, Generator, Model, Query\n'), ((2214, 2229), 'sqlalchemy.sql.expression.text', 'text', (['"""\'title\'"""'], {}), '("\'title\'")\n', (2218, 2229), False, 'from sqlalchemy.sql.expression import text\n'), ((2739, 2757), 'sqlmodel.select', 'select', (['JSONEntity'], {}), '(JSONEntity)\n', (2745, 2757), False, 'from sqlmodel import Session, select\n')]