"""
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.
"""

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"
REGISTRY_TABLE = f"{config.SNOWFLAKE_DATABASE}.{config.SNOWFLAKE_SCHEMA}.ECOMMERCE_STORE_REGISTRY"

# QA uses the new location; Prod keeps the current location
if config.ENVIRONMENT == "qa":
    ARTIST_STORES_TABLE = f"{config.SNOWFLAKE_DATABASE}.{config.SNOWFLAKE_SCHEMA}.ARTIST_STORES"
else:
    ARTIST_STORES_TABLE = "ORCHARD_APP_REPORTING_V2.D2C.ARTIST_STORES"

# Default guard against runaway statements.
DEFAULT_STATEMENT_TIMEOUT_SECONDS = 30


def _sql_norm(col: str) -> str:
    return (
        f"lower(regexp_replace(regexp_replace({col}, '^\\\\s+|\\\\s+$', ''), "
        f"'^https?://(www\\\\.)?|^www\\\\.|/+$', ''))"
    )


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 get_artist_stores_metadata(conn: snowflake.connector.SnowflakeConnection) -> list[dict]:
    with cursor(conn) as cur:
        cur.execute(f"""
            SELECT
                {_sql_norm("SHOPIFY_SHOP")}        AS shopify_shop,
                {_sql_norm("SHOPIFY_STORE_URL")}   AS shopify_store_url,
                {_sql_norm("PUBLIC_DOMAIN")}        AS public_domain,
                {_sql_norm("ALTERNATE_DOMAIN")}     AS alternate_domain,
                SELLING_COUNTRY, REP_OWNER, MERCH_COMPANY, ACTIVE_
            FROM {ARTIST_STORES_TABLE}
            WHERE SHOPIFY_SHOP IS NOT NULL OR SHOPIFY_STORE_URL IS NOT NULL
        """)
        cols = [d[0].lower() for d in cur.description]
        return [dict(zip(cols, row)) for row in cur.fetchall()]


def get_onboarding_domains(conn: snowflake.connector.SnowflakeConnection) -> frozenset[str]:
    with cursor(conn) as cur:
        cur.execute(f"""
            SELECT SPLIT_PART(lower(SHOP_DOMAIN), '/', 1) FROM {ONBOARDING_STATE_TABLE}
            WHERE SHOP_DOMAIN IS NOT NULL
            UNION
            SELECT SPLIT_PART(lower(ALT_MYSHOPIFY_DOMAIN), '/', 1) FROM {ONBOARDING_STATE_TABLE}
            WHERE ALT_MYSHOPIFY_DOMAIN IS NOT NULL
        """)
        return frozenset(row[0] for row in cur.fetchall() if row[0])


def get_registry_domains(conn: snowflake.connector.SnowflakeConnection) -> frozenset[str]:
    with cursor(conn) as cur:
        cur.execute(f"""
            SELECT SPLIT_PART(lower(MYSHOPIFY_DOMAIN), '/', 1) FROM {REGISTRY_TABLE}
            WHERE MYSHOPIFY_DOMAIN IS NOT NULL
            UNION
            SELECT SPLIT_PART(lower(ALT_MYSHOPIFY_DOMAIN), '/', 1) FROM {REGISTRY_TABLE}
            WHERE ALT_MYSHOPIFY_DOMAIN IS NOT NULL
        """)
        return frozenset(row[0] for row in cur.fetchall())


def get_crm_domains(conn: snowflake.connector.SnowflakeConnection) -> frozenset[str]:
    with cursor(conn) as cur:
        cur.execute("""
            SELECT TABLE_SCHEMA
            FROM CRM_ECOMMERCE_DATA.INFORMATION_SCHEMA.TABLES
            WHERE TABLE_NAME = 'SHOP'
            ORDER BY TABLE_SCHEMA
        """)
        schemas = [row[0] for row in cur.fetchall()]

    if not schemas:
        return frozenset()

    unions = "\nUNION ALL\n".join(
        f"SELECT SPLIT_PART(lower(MYSHOPIFY_DOMAIN), '/', 1) AS domain "
        f'FROM CRM_ECOMMERCE_DATA."{esc}"."SHOP" WHERE MYSHOPIFY_DOMAIN IS NOT NULL'
        for s in schemas
        for esc in [s.replace('"', '""')]
    )
    with cursor(conn) as cur:
        cur.execute(f"SELECT DISTINCT domain FROM ({unions})")
        return frozenset(row[0] for row in cur.fetchall() if row[0])


def get_fansifter_domains(conn: snowflake.connector.SnowflakeConnection) -> frozenset[str]:
    with cursor(conn) as cur:
        cur.execute(f"""
            SELECT SPLIT_PART({_sql_norm("MYSHOPIFY_DOMAIN")}, '/', 1)
            FROM FANSIFTER_APP_REPORTING.PROD_SHOPIFY.SHOPIFY_SHOP
            WHERE MYSHOPIFY_DOMAIN IS NOT NULL
        """)
        return frozenset(row[0] for row in cur.fetchall() if row[0])


def upsert_registry(conn: snowflake.connector.SnowflakeConnection, data: dict) -> None:
    # Discovery-time write: owns IS_ACTIVE (read from ARTIST_STORES). This is
    # intentionally NOT the same column set as sync-shopify-data-connectors'
    # upsert_registry, which owns the vendor_id / global_participant enrichment
    # and never touches IS_ACTIVE. Keep the two asymmetric.
    with cursor(conn) as cur:
        cur.execute(
            f"""
            MERGE INTO {REGISTRY_TABLE} AS target
            USING (
                SELECT
                    %(myshopify_domain)s       AS myshopify_domain,
                    %(schema_name)s            AS schema_name,
                    %(source)s                 AS source,
                    %(connector_id)s           AS connector_id,
                    %(selling_country)s        AS selling_country,
                    %(merch_company)s          AS merch_company,
                    %(rep_owner)s              AS rep_owner,
                    %(alt_myshopify_domain)s   AS alt_myshopify_domain,
                    %(custom_domain)s          AS custom_domain,
                    %(is_active)s              AS is_active,
                    %(is_synced_to_snowflake)s AS is_synced_to_snowflake
            ) AS staged ON {_sql_norm("target.myshopify_domain")} = {_sql_norm("staged.myshopify_domain")}
            WHEN MATCHED THEN UPDATE SET
                schema_name            = COALESCE(staged.schema_name, target.schema_name),
                source                 = COALESCE(staged.source, target.source),
                connector_id           = COALESCE(staged.connector_id, target.connector_id),
                selling_country        = COALESCE(staged.selling_country, target.selling_country),
                merch_company          = COALESCE(staged.merch_company, target.merch_company),
                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),
                is_active              = COALESCE(staged.is_active, target.is_active),
                is_synced_to_snowflake = COALESCE(staged.is_synced_to_snowflake, target.is_synced_to_snowflake)
            WHEN NOT MATCHED THEN INSERT (
                myshopify_domain, schema_name, source,
                connector_id, selling_country, merch_company, rep_owner,
                alt_myshopify_domain, custom_domain, is_active, is_synced_to_snowflake
            ) VALUES (
                staged.myshopify_domain, staged.schema_name, staged.source,
                staged.connector_id, staged.selling_country, staged.merch_company, staged.rep_owner,
                staged.alt_myshopify_domain, staged.custom_domain,
                staged.is_active, staged.is_synced_to_snowflake
            )
        """,
            {
                k: data.get(k)
                for k in [
                    "myshopify_domain",
                    "schema_name",
                    "source",
                    "connector_id",
                    "selling_country",
                    "merch_company",
                    "rep_owner",
                    "alt_myshopify_domain",
                    "custom_domain",
                    "is_active",
                    "is_synced_to_snowflake",
                ]
            },
        )