#   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(
            