"""Truncate test table before running integration tests."""

import base64
import os

import snowflake.connector
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from dotenv import load_dotenv

load_dotenv()


def load_key_bytes():
    """Load Snowflake private key."""
    pem_env = os.getenv('SNOWFLAKE_PRIVATE_KEY')
    passphrase = os.getenv('SNOWFLAKE_KEY_PASSPHRASE')
    if pem_env:
        blob = pem_env.encode()
        if b'BEGIN' not in blob:
            blob = base64.b64decode(pem_env)
        return serialization.load_pem_private_key(
            blob,
            password=passphrase.encode() if passphrase else None,
            backend=default_backend(),
        )

    key_path = os.environ['SNOWFLAKE_PRIVATE_KEY_PATH']
    with open(key_path, 'rb') as f:
        return serialization.load_pem_private_key(
            f.read(),
            password=passphrase.encode() if passphrase else None,
            backend=default_backend(),
        )


private_key = load_key_bytes()
pkcs8 = private_key.private_bytes(
    encoding=serialization.Encoding.DER,
    format=serialization.PrivateFormat.PKCS8,
    encryption_algorithm=serialization.NoEncryption(),
)

ctx = snowflake.connector.connect(
    user=os.environ['SNOWFLAKE_USER'],
    account=os.environ['SNOWFLAKE_ACCOUNT'],
    warehouse=os.environ['SNOWFLAKE_WAREHOUSE'],
    database=os.environ['SNOWFLAKE_DATABASE'],
    schema='TEST',
    role=os.environ['SNOWFLAKE_ROLE'],
    private_key=pkcs8,
    autocommit=False,
    application='integration-tests',
)

try:
    # hardcoded for safety
    sql = """TRUNCATE TABLE ROYALTY_ACCOUNTING.TEST.BOOKED_ACCOUNT_CONTRACT_SNAPSHOT;"""
    with ctx.cursor() as cs:
        cs.execute(sql)
    ctx.commit()
finally:
    ctx.close()