from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
from jinja2 import Template

from lib.config import OUTFILE_AWS_KEY_ID
from lib.config import OUTFILE_AWS_SECRET_KEY
from lib.config import OUTFILE_BUCKET
from lib.config import SF_SCHEMA


QUERY_TEMPLATE = Template("""
    CREATE OR REPLACE TRANSIENT TABLE ROYALTY_ACCOUNTING.{{schema}}.STMT_DB_SALES_DISTRO_INGEST (
        BATCH_ID VARCHAR(255),
        STATEMENT_ID NUMBER(32,0),
        VENDOR_ID NUMBER(32,0),
        CUSTOMER_MASTER_MASTER_ID NUMBER(32,0),
        SUBDISTRIBUTOR VARCHAR(255),
        ORIGINAL_CURRENCY_ISO VARCHAR(30),
        ACTIVITY_RATE NUMBER(18,6),
        COUNTRY_ID NUMBER(12,0),
        START_DATE DATE,
        DATE DATE,
        UPC NUMBER(38,0),
        CD NUMBER(32,0),
        TRACK_ID NUMBER(32,0),
        ISRC VARCHAR(255),
        TRACK_NAME VARCHAR(255),
        QTY NUMBER(18,0),
        UNIT_PRICE NUMBER(18,6),
        TOTAL NUMBER(18,6),
        WHT NUMBER(18,6),
        CONFIGURATION VARCHAR(255),
        TRANS_TYPE VARCHAR(2),
        TRANS_SUBTYPE NUMBER(18,0),
        RETAIL_PRICE NUMBER(18,6),
        ORIGINAL_PRICE NUMBER(18,6),
        DISCOUNT NUMBER(18,6),
        PHYS_PPD NUMBER(18,6),
        VIDEO_ID VARCHAR(255),
        ACTUAL_STATEMENT_NO VARCHAR(128),
        SAP_EXCLUDE BOOLEAN,
        UNIQUE_DETAIL_ID NUMBER(38,0)
    );

    COPY INTO ROYALTY_ACCOUNTING.{{schema}}.STMT_DB_SALES_DISTRO_INGEST
    FROM {{outfile_location}}
    CREDENTIALS=(
        AWS_KEY_ID='{{aws_key_id}}'
        AWS_SECRET_KEY='{{aws_secret_key}}'
    )
    FILE_FORMAT=(
        TYPE='CSV'
        FIELD_DELIMITER='\t'
        RECORD_DELIMITER='\n'
        COMPRESSION='GZIP'
        NULL_IF=('NULL')
    )
    ON_ERROR=ABORT_STATEMENT
""")


def load_sales(**kwargs):
    batch_id = kwargs.get('params').get('batch_id')
    outfile_location = f's3://{OUTFILE_BUCKET}/stmtdb-to-sf/extract-sales/distro/{batch_id}/parts/' 

    query = QUERY_TEMPLATE.render(
        schema=SF_SCHEMA,
        outfile_location=outfile_location,
        aws_key_id=OUTFILE_AWS_KEY_ID,
        aws_secret_key=OUTFILE_AWS_SECRET_KEY
    )

    sf_hook = SnowflakeHook()
    sf_hook.run(query, autocommit=True)