import argparse
import datetime
import json
import logging
import os
from pathlib import Path

import snowflake.connector
import yaml
from snowflake.connector.converter import SnowflakeConverter

logger = logging.getLogger(__name__)


def snowflake_connection():
    logger.info('Getting snowflake connection')
    return snowflake.connector.connect(
        account=os.environ['SNOWFLAKE_ACCOUNT'],
        user=os.environ['SNOWFLAKE_USER'],
        role=os.environ.get('SNOWFLAKE_ROLE'),
        password=os.environ['SNOWFLAKE_PASSWORD'],
        database=os.environ['SNOWFLAKE_DATABASE'],
        schema=os.environ['SNOWFLAKE_SCHEMA']
    )


DBT_EXECUTIONS_TABLE = os.environ.get('DBT_EXECUTIONS_TABLE', 'dbt_executions')

COLUMNS = {
    "label": "varchar",
    "command_line": "varchar",
    "started_at": "datetime",
    "exit_status": "number",
    "project_name": "varchar",
    "project_id": "varchar",
    "invocation_id": "varchar",
    "generated_at": "datetime",
    "manifest_json": "variant",
    "sources_json": "variant",
    "run_results_json": "variant",
}


def get_ddl(overwrite=False):
    columns_sql = ',\n'.join([f'{k} {v}' for k, v in COLUMNS.items()])
    query_sql = f"CREATE {'OR REPLACE TABLE' if overwrite else 'TABLE IF NOT EXISTS'} {DBT_EXECUTIONS_TABLE} ({columns_sql})"
    return query_sql

def insert_into_executions(connection: snowflake.connector.SnowflakeConnection, row: dict):
    columns = []
    values = []
    logger.info(f'Saving data {row.keys()} ...')
    for column, type in COLUMNS.items():
        columns.append(column)
        value = row.get(column)
        if isinstance(value, list) or isinstance(value, dict):
            value = f"PARSE_JSON({SnowflakeConverter.quote(SnowflakeConverter.escape(json.dumps(value)))})"
            values.append(value)
        else:
            value = SnowflakeConverter.quote(value)
            values.append(value)
    logger.info(f'Insert into {DBT_EXECUTIONS_TABLE} ...')
    sql_query = f"insert into {DBT_EXECUTIONS_TABLE} ({', '.join(columns)}) select {', '.join(values)}"
    with connection.cursor() as cursor:
        cursor.execute(sql_query)


def _load_json_if_exists(dir_path: Path, filename: str):
    path = dir_path / filename
    if not path.exists():
        logger.info(f"File {filename} not found.")
        return None
    with path.open() as manifest_file:
        logger.info(f"Loading from {filename} ...")
        return json.load(manifest_file)


def target2snowflake(dbt_dir, dbt_cli=None, dbt_exit_status=None, label=None, started_at=None):
    dbt_dir_path = Path(dbt_dir)
    assert dbt_dir_path.exists() and dbt_dir_path.is_dir()
    dbt_project_yml_path = dbt_dir_path / 'dbt_project.yml'
    with dbt_project_yml_path.open() as dbt_project_yml_file:
        dbt_project = yaml.safe_load(dbt_project_yml_file)

    target_dir_path = dbt_dir_path / 'target'
    assert target_dir_path.exists() and target_dir_path.is_dir()

    manifest = _load_json_if_exists(target_dir_path, 'manifest.json')
    sources = _load_json_if_exists(target_dir_path, 'sources.json')
    run_results = _load_json_if_exists(target_dir_path, 'run_results.json')

    # filter out jsons from old invocations
    if manifest and sources and manifest['metadata']['invocation_id'] != sources['metadata']['invocation_id']:
        sources = None
    if manifest and run_results and manifest['metadata']['invocation_id'] != run_results['metadata']['invocation_id']:
        run_results = None

    started_at_sql = started_at if started_at else datetime.datetime.now().isoformat()

    result = {
        'label': label,
        'project_name': dbt_project['name'],
        'command_line': dbt_cli,
        'project_id': manifest['metadata']['project_id'],
        'invocation_id': manifest['metadata']['invocation_id'],
        'generated_at': manifest['metadata']['generated_at'],
        'manifest_json': manifest,
        'sources_json': sources,
        'run_results_json': run_results,
        'exit_status': dbt_exit_status,
        'started_at': started_at_sql,
    }
    return result


def parse_args(args):
    parser = argparse.ArgumentParser(description='Upload DBT execution results to Snowflake.')
    parser.add_argument('-c', '--command', help='Command line for dbt. Should start with word dbt. Example "dbt run -m model"')
    parser.add_argument('-p', '--project', help='dbt project root directory. By default is current directory.')
    parser.add_argument('-s', '--status', help='dbt command exit code. 0 - succes, != 0 - error')
    parser.add_argument('-l', '--label', help='Optional label for execution')
    parser.add_argument('--started-at', help='Optional execution datetime')
    parser.add_argument('--ddl', action='store_true', help='Print DDL statement for the table.')
    return vars(parser.parse_args(args))


def main():
    parsed = parse_args(args=None)
    if parsed.get('ddl'):
        ddl = get_ddl(overwrite=True)
        print(ddl)
        exit(0)
    project_dir = Path(parsed.get('project') or '.')
    cli = parsed.get('command')
    connection = snowflake_connection()
    row = target2snowflake(
        dbt_dir=project_dir,
        dbt_cli=cli,
        dbt_exit_status=parsed.get('status'),
        label=parsed.get('label'),
        started_at=parsed.get('started_at'),
    )
    insert_into_executions(connection, row)


if __name__ == '__main__':
    logging.basicConfig(level=logging.INFO)
    main()