"""
Snowflake connection. In Lambda, reads the PEM private key from config
(loaded at cold start from Secrets Manager). Falls back to
SNOWFLAKE_PRIVATE_KEY_PATH for local development.

This Lambda only writes onboarding state; discovery/registry/enrichment
queries live in the connect and sync Lambdas.
"""

import os
from contextlib import contextmanager
from pathlib import Path

import snowflake.connector
from cryptography.hazmat.primitives import serialization

import config

ONBOARDING_STATE_TABLE = f"{config.SNOWFLAKE_DATABASE}.{config.SNOWFLAKE_SCHEMA}.D2C_ONBOARDING_STATE"
CONTACT_MAPPINGS_TABLE = f"{config.SNOWFLAKE_DATABASE}.{config.SNOWFLAKE_SCHEMA}.D2C_CONTACT_MAPPINGS"

# Guard against a runaway statement. This Lambda only runs a single fast MERGE,
# but the tight default matches the sync Lambda for consistency.
DEFAULT_STATEMENT_TIMEOUT_SECONDS = 30


def _load_private_key() -> bytes:
    if config.SNOWFLAKE_PRIVATE_KEY_PEM:
        pem_data = config.SNOWFLAKE_PRIVATE_KEY_PEM.encode()
    else:
        key_path = os.environ.get("SNOWFLAKE_PRIVATE_KEY_PATH", str(Path.home() / ".ssh/snowflake/rsa_key.p8"))
        with open(key_path, "rb") as f:
            pem_data = f.read()

    passphrase_raw = os.environ.get("SNOWFLAKE_KEY_PASSPHRASE", "").strip("'\"") or None
    password = passphrase_raw.encode() if passphrase_raw else None

    try:
        p_key = serialization.load_pem_private_key(pem_data, password=password)
    except TypeError:
        p_key = serialization.load_pem_private_key(pem_data, password=None)

    return p_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )


def get_connection() -> snowflake.connector.SnowflakeConnection:
    return snowflake.connector.connect(
        account=config.SNOWFLAKE_ACCOUNT,
        user=config.SNOWFLAKE_USER,
        role=config.SNOWFLAKE_ROLE,
        warehouse=config.SNOWFLAKE_WAREHOUSE,
        database=config.SNOWFLAKE_DATABASE,
        schema=config.SNOWFLAKE_SCHEMA,
        private_key=_load_private_key(),
        session_parameters={
            "QUERY_TAG": "shopify-data-connector-onboarding",
            "STATEMENT_TIMEOUT_IN_SECONDS": DEFAULT_STATEMENT_TIMEOUT_SECONDS,
        },
    )


@contextmanager
def cursor(conn: snowflake.connector.SnowflakeConnection):
    cur = conn.cursor()
    try:
        yield cur
    finally:
        cur.close()


def upsert_state(conn: snowflake.connector.SnowflakeConnection, data: dict) -> None:
    with cursor(conn) as cur:
        cur.execute(
            f"""
            MERGE INTO {ONBOARDING_STATE_TABLE} AS target
            USING (
                SELECT
                    %(shop_domain)s          AS shop_domain,
                    %(schema_name)s          AS schema_name,
                    %(connector_id)s         AS connector_id,
                    %(status)s               AS status,
                    %(error_message)s        AS error_message,
                    %(connect_card_uri)s     AS connect_card_uri,
                    %(link_created_at)s      AS link_created_at,
                    %(link_expires_at)s      AS link_expires_at,
                    %(connected_at)s         AS connected_at,
                    %(merch_company)s        AS merch_company,
                    %(selling_country)s      AS selling_country,
                    %(rep_owner)s            AS rep_owner,
                    %(alt_myshopify_domain)s AS alt_myshopify_domain,
                    %(custom_domain)s        AS custom_domain,
                    CURRENT_TIMESTAMP()      AS updated_at
            ) AS staged ON LOWER(target.shop_domain) = LOWER(staged.shop_domain)
            WHEN MATCHED THEN UPDATE SET
                schema_name          = COALESCE(staged.schema_name, target.schema_name),
                connector_id         = COALESCE(staged.connector_id, target.connector_id),
                status               = COALESCE(staged.status, target.status),
                error_message        = staged.error_message,
                connect_card_uri     = staged.connect_card_uri,
                link_created_at      = staged.link_created_at,
                link_expires_at      = staged.link_expires_at,
                connected_at         = staged.connected_at,
                merch_company        = COALESCE(staged.merch_company, target.merch_company),
                selling_country      = COALESCE(staged.selling_country, target.selling_country),
                rep_owner            = COALESCE(staged.rep_owner, target.rep_owner),
                alt_myshopify_domain = COALESCE(staged.alt_myshopify_domain, target.alt_myshopify_domain),
                custom_domain        = COALESCE(staged.custom_domain, target.custom_domain),
                updated_at           = staged.updated_at
            WHEN NOT MATCHED THEN INSERT (
                shop_domain, schema_name, connector_id, status,
                error_message, connect_card_uri, link_created_at, link_expires_at,
                connected_at, merch_company, selling_country, rep_owner,
                alt_myshopify_domain, custom_domain, updated_at
            ) VALUES (
                staged.shop_domain, staged.schema_name, staged.connector_id, staged.status,
                staged.error_message, staged.connect_card_uri, staged.link_created_at, staged.link_expires_at,
                staged.connected_at, staged.merch_company, staged.selling_country, staged.rep_owner,
                staged.alt_myshopify_domain, staged.custom_domain, staged.updated_at
            )
        """,
            {
                k: data.get(k)
                for k in [
                    "shop_domain",
                    "schema_name",
                    "connector_id",
                    "status",
                    "error_message",
                    "connect_card_uri",
                    "link_created_at",
                    "link_expires_at",
                    "connected_at",
                    "merch_company",
                    "selling_country",
                    "rep_owner",
                    "alt_myshopify_domain",
                    "custom_domain",
                ]
            },
        )


def get_contacts(
    conn: snowflake.connector.SnowflakeConnection,
    selling_country: str,
    merch_company: str,
) -> list[dict] | str | None:
    """Fetch contact list for a country/company combination.

    Args:
        conn: Snowflake connection
        selling_country: Country code (e.g., 'EU')
        merch_company: The full merch company name as stored in
            ARTIST_STORES.MERCH_COMPANY (e.g., 'Kings Road', 'IndieMerch'),
            not an abbreviated code

    Returns:
        Raw CONTACTS value from Snowflake — a list of dicts with uppercase
        keys (NAME, ROLE, COMPANY, EMAIL) or a JSON string of the same,
        depending on driver deserialization. Callers must pass this through
        parse_contacts_from_array before use. None if no active mapping exists.
    """
    with cursor(conn) as cur:
        cur.execute(
            f"""
            SELECT CONTACTS
            FROM {CONTACT_MAPPINGS_TABLE}
            WHERE SELLING_COUNTRY = %(selling_country)s
              AND MERCH_COMPANY = %(merch_company)s
              AND ACTIVE = TRUE
            LIMIT 1
            """,
            {
                "selling_country": selling_country,
                "merch_company": merch_company,
            },
        )
        row = cur.fetchone()
        if row:
            return row[0]  # CONTACTS array
    return None