from uuid import uuid4
import snowflake.connector
from aiokafka import AIOKafkaConsumer
from aiokafka.helpers import create_ssl_context

from dataexport.conf import config

DATABASE_CONFIG = {
    "database_name": config.SNOWFLAKE_DATABASE,
    "schema_name": config.SNOWFLAKE_SCHEMA,
}

if not config.SNOWFLAKE_PASSWORD:
    raise Exception("Need to setup .env creds")

connection = snowflake.connector.connect(
    user=config.SNOWFLAKE_USER,
    password=config.SNOWFLAKE_PASSWORD,
    account=config.SNOWFLAKE_ACCOUNT,
    warehouse=config.SNOWFLAKE_WAREHOUSE,
    database=config.SNOWFLAKE_DATABASE,
    schema=config.SNOWFLAKE_SCHEMA,
    role=config.SNOWFLAKE_ROLE,
    private_key=config.SNOWFLAKE_CONNECT_ARGS.get("private_key")
)

cursor = connection.cursor()


def get_kafka_consumer() -> AIOKafkaConsumer:
    """Factory which returns configured kafka consumer."""
    return AIOKafkaConsumer(
        config.TOPICS,
        value_deserializer=lambda v: str(v.decode()),
        bootstrap_servers=config.BOOTSTRAP_SERVERS,
        security_protocol=config.SECURITY_PROTOCOL,
        ssl_context=create_ssl_context(),
        group_id=config.GROUP_ID,
        client_id=str(uuid4())
    )