#   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 itertools import chain
from json import dumps
from pathlib import Path
from typing import List
from uuid import UUID

import typer
from prettytable import ALL, PrettyTable
from sqlalchemy import update
from sqlmodel import Session, select

import dbgen.cli.styles as styles
from dbgen.cli.options import config_option, model_arg_option
from dbgen.cli.utils import test_connection, validate_model_str
from dbgen.configuration import initialize
from dbgen.core.metadata import ModelEntity

model_app = typer.Typer(name='model', no_args_is_help=True)


@model_app.command('list')
def list_models(config_file: Path = config_option, tags: List[str] = typer.Option(None, '-t')):

    # Notify of config file
    if config_file:
        _, meta_conn = initialize(config_file)
    test_connection(meta_conn)
    meta_engine = meta_conn.get_engine()
    tags = tags or []
    statement = select(  # type: ignore
        ModelEntity.id,
        ModelEntity.name,
        ModelEntity.created_at,
        ModelEntity.last_run,
        ModelEntity.tags,
    )  # type: ignore
    if tags:
        statement = statement.where(ModelEntity.tags.op('&&')(tags))  # type: ignore
    columns = ['id', 'name', 'created_at', 'last_run', 'tags']
    table = PrettyTable(field_names=columns, align='l', hrules=ALL)
    with Session(meta_engine) as session:
        result = session.exec(statement)
        for model_id, model_name, created_at, last_run, tags in result:
            table.add_row((model_id, model_name, created_at, last_run, tags))
    styles.theme_typer_print(str(table))


@model_app.command('tag')
def tag(model_id: UUID, tags: List[str], config_file: Path = config_option):
    # Notify of config file
    if config_file:
        _, meta_conn = initialize(config_file)
    test_connection(meta_conn)
    meta_engine = meta_conn.get_engine()

    with Session(meta_engine) as session:
        existing_tags = session.exec(select(ModelEntity.tags).where(ModelEntity.id == model_id)).one_or_none()
        if existing_tags is None:
            raise typer.BadParameter(f"Invalid model_id, no model with ID {model_id}")
        new_tags = set(chain(existing_tags, tags))
        session.execute(update(ModelEntity).values(tags=new_tags).where(ModelEntity.id == model_id))

        session.commit()


@model_app.command('serialize')
def model_serialize(
    model_str: str = model_arg_option,
    out_file: Path = typer.Option(
        None, '-o', '--out', help='Path to write the serialized model to in json format'
    ),
    config_file: Path = config_option,
):
    model = validate_model_str(model_str)

    # Notify of config file
    if config_file:
        _, meta_conn = initialize(config_file)
    test_connection(meta_conn)
    meta_engine = meta_conn.get_engine()
    with Session(meta_engine) as session:
        model_row = model._get_model_row()
        # Check for existing row and if found grab its created_at
        created_at = session.exec(
            select(ModelEntity.created_at).where(ModelEntity.id == model.uuid)
        ).one_or_none()
        if created_at is None:
            session.merge(model_row)
            session.commit()
            styles.good_typer_print(f"Loaded model {model.name!r} into the database with ID {model.uuid}")
        else:
            styles.good_typer_print(f"Model {model.name!r} already existed with ID {model.uuid}")
    if out_file:
        out_file.write_text(dumps(model_row.graph_json))
        styles.good_typer_print(f"Wrote serialized graph to {out_file}")


@model_app.command('export')
def model_export(
    model_id: UUID,
    out_file: Path = typer.Option(
        'model.json', '-o', '--out', help='Path to write the serialized model to in json format'
    ),
    config_file: Path = config_option,
):

    # Notify of config file
    if config_file:
        _, meta_conn = initialize(config_file)
    test_connection(meta_conn)
    meta_engine = meta_conn.get_engine()

    with Session(meta_engine) as session:
        # Check for existing row and if found grab its created_at
        graph_json = session.exec(
            select(ModelEntity.graph_json).where(ModelEntity.id == model_id)
        ).one_or_none()
        if not graph_json:
            raise ValueError(f"Invalid model_id: No model found with model_id {model_id}")

    out_file.write_text(dumps(graph_json))
    styles.good_typer_print(f"Wrote serialized graph to {out_file}")

itertools.chain typer.BadParameter json.dumps dbgen.core.metadata.ModelEntity.tags.op dbgen.cli.utils.validate_model_str sqlmodel.Session typer.Option sqlalchemy.update sqlmodel.select dbgen.configuration.initialize dbgen.cli.utils.test_connection typer.Typer prettytable.PrettyTable dbgen.cli.styles.good_typer_print
[((1102, 1149), 'typer.Typer', 'typer.Typer', ([], {'name': '"""model"""', 'no_args_is_help': '(True)'}), "(name='model', no_args_is_help=True)\n", (1113, 1149), False, 'import typer\n'), ((1248, 1272), 'typer.Option', 'typer.Option', (['None', '"""-t"""'], {}), "(None, '-t')\n", (1260, 1272), False, 'import typer\n'), ((1375, 1401), 'dbgen.cli.utils.test_connection', 'test_connection', (['meta_conn'], {}), '(meta_conn)\n', (1390, 1401), False, 'from dbgen.cli.utils import test_connection, validate_model_str\n'), ((1481, 1589), 'sqlmodel.select', 'select', (['ModelEntity.id', 'ModelEntity.name', 'ModelEntity.created_at', 'ModelEntity.last_run', 'ModelEntity.tags'], {}), '(ModelEntity.id, ModelEntity.name, ModelEntity.created_at,\n    ModelEntity.last_run, ModelEntity.tags)\n', (1487, 1589), False, 'from sqlmodel import Session, select\n'), ((1838, 1893), 'prettytable.PrettyTable', 'PrettyTable', ([], {'field_names': 'columns', 'align': '"""l"""', 'hrules': 'ALL'}), "(field_names=columns, align='l', hrules=ALL)\n", (1849, 1893), False, 'from prettytable import ALL, PrettyTable\n'), ((2372, 2398), 'dbgen.cli.utils.test_connection', 'test_connection', (['meta_conn'], {}), '(meta_conn)\n', (2387, 2398), False, 'from dbgen.cli.utils import test_connection, validate_model_str\n'), ((3008, 3107), 'typer.Option', 'typer.Option', (['None', '"""-o"""', '"""--out"""'], {'help': '"""Path to write the serialized model to in json format"""'}), "(None, '-o', '--out', help=\n    'Path to write the serialized model to in json format')\n", (3020, 3107), False, 'import typer\n'), ((3172, 3201), 'dbgen.cli.utils.validate_model_str', 'validate_model_str', (['model_str'], {}), '(model_str)\n', (3190, 3201), False, 'from dbgen.cli.utils import test_connection, validate_model_str\n'), ((3302, 3328), 'dbgen.cli.utils.test_connection', 'test_connection', (['meta_conn'], {}), '(meta_conn)\n', (3317, 3328), False, 'from dbgen.cli.utils import test_connection, validate_model_str\n'), ((4212, 4319), 'typer.Option', 'typer.Option', (['"""model.json"""', '"""-o"""', '"""--out"""'], {'help': '"""Path to write the serialized model to in json format"""'}), "('model.json', '-o', '--out', help=\n    'Path to write the serialized model to in json format')\n", (4224, 4319), False, 'import typer\n'), ((4472, 4498), 'dbgen.cli.utils.test_connection', 'test_connection', (['meta_conn'], {}), '(meta_conn)\n', (4487, 4498), False, 'from dbgen.cli.utils import test_connection, validate_model_str\n'), ((4951, 5015), 'dbgen.cli.styles.good_typer_print', 'styles.good_typer_print', (['f"""Wrote serialized graph to {out_file}"""'], {}), "(f'Wrote serialized graph to {out_file}')\n", (4974, 5015), True, 'import dbgen.cli.styles as styles\n'), ((1347, 1370), 'dbgen.configuration.initialize', 'initialize', (['config_file'], {}), '(config_file)\n', (1357, 1370), False, 'from dbgen.configuration import initialize\n'), ((1903, 1923), 'sqlmodel.Session', 'Session', (['meta_engine'], {}), '(meta_engine)\n', (1910, 1923), False, 'from sqlmodel import Session, select\n'), ((2344, 2367), 'dbgen.configuration.initialize', 'initialize', (['config_file'], {}), '(config_file)\n', (2354, 2367), False, 'from dbgen.configuration import initialize\n'), ((2450, 2470), 'sqlmodel.Session', 'Session', (['meta_engine'], {}), '(meta_engine)\n', (2457, 2470), False, 'from sqlmodel import Session, select\n'), ((3274, 3297), 'dbgen.configuration.initialize', 'initialize', (['config_file'], {}), '(config_file)\n', (3284, 3297), False, 'from dbgen.configuration import initialize\n'), ((3379, 3399), 'sqlmodel.Session', 'Session', (['meta_engine'], {}), '(meta_engine)\n', (3386, 3399), False, 'from sqlmodel import Session, select\n'), ((4057, 4121), 'dbgen.cli.styles.good_typer_print', 'styles.good_typer_print', (['f"""Wrote serialized graph to {out_file}"""'], {}), "(f'Wrote serialized graph to {out_file}')\n", (4080, 4121), True, 'import dbgen.cli.styles as styles\n'), ((4444, 4467), 'dbgen.configuration.initialize', 'initialize', (['config_file'], {}), '(config_file)\n', (4454, 4467), False, 'from dbgen.configuration import initialize\n'), ((4550, 4570), 'sqlmodel.Session', 'Session', (['meta_engine'], {}), '(meta_engine)\n', (4557, 4570), False, 'from sqlmodel import Session, select\n'), ((4928, 4945), 'json.dumps', 'dumps', (['graph_json'], {}), '(graph_json)\n', (4933, 4945), False, 'from json import dumps\n'), ((2646, 2714), 'typer.BadParameter', 'typer.BadParameter', (['f"""Invalid model_id, no model with ID {model_id}"""'], {}), "(f'Invalid model_id, no model with ID {model_id}')\n", (2664, 2714), False, 'import typer\n'), ((2738, 2764), 'itertools.chain', 'chain', (['existing_tags', 'tags'], {}), '(existing_tags, tags)\n', (2743, 2764), False, 'from itertools import chain\n'), ((3768, 3867), 'dbgen.cli.styles.good_typer_print', 'styles.good_typer_print', (['f"""Loaded model {model.name!r} into the database with ID {model.uuid}"""'], {}), "(\n    f'Loaded model {model.name!r} into the database with ID {model.uuid}')\n", (3791, 3867), True, 'import dbgen.cli.styles as styles\n'), ((3889, 3979), 'dbgen.cli.styles.good_typer_print', 'styles.good_typer_print', (['f"""Model {model.name!r} already existed with ID {model.uuid}"""'], {}), "(\n    f'Model {model.name!r} already existed with ID {model.uuid}')\n", (3912, 3979), True, 'import dbgen.cli.styles as styles\n'), ((4020, 4047), 'json.dumps', 'dumps', (['model_row.graph_json'], {}), '(model_row.graph_json)\n', (4025, 4047), False, 'from json import dumps\n'), ((1714, 1739), 'dbgen.core.metadata.ModelEntity.tags.op', 'ModelEntity.tags.op', (['"""&&"""'], {}), "('&&')\n", (1733, 1739), False, 'from dbgen.core.metadata import ModelEntity\n'), ((2520, 2544), 'sqlmodel.select', 'select', (['ModelEntity.tags'], {}), '(ModelEntity.tags)\n', (2526, 2544), False, 'from sqlmodel import Session, select\n'), ((2790, 2809), 'sqlalchemy.update', 'update', (['ModelEntity'], {}), '(ModelEntity)\n', (2796, 2809), False, 'from sqlalchemy import update\n'), ((3568, 3598), 'sqlmodel.select', 'select', (['ModelEntity.created_at'], {}), '(ModelEntity.created_at)\n', (3574, 3598), False, 'from sqlmodel import Session, select\n'), ((4696, 4726), 'sqlmodel.select', 'select', (['ModelEntity.graph_json'], {}), '(ModelEntity.graph_json)\n', (4702, 4726), False, 'from sqlmodel import Session, select\n')]