"""
lambda-d2c-onboard-shopify-stores

Triggered by SQS (batch size <= 20, max concurrency 2). For each store:
  1. Creates the Fivetran connector
  2. Generates the Connect Card URL
  3. Writes pending_oauth state to Snowflake

The polling loop from the legacy script is intentionally removed.
lambda-d2c-sync-shopify-data-connectors advances state from here.
"""

import json
import logging
from datetime import datetime, timedelta, UTC
from collections import defaultdict
from itertools import batched

from datadog_lambda.metric import lambda_metric

import config
from src.fivetran_client import FivetranClient, ConnectionSetupState
from src.snowflake_client import get_connection, upsert_state, get_contacts
from src.ses_client import SESClient
from src.email_utils import format_onboarding_email, parse_contacts_from_array, get_email_addresses

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

# SES SendEmail caps total recipients (To+Cc+Bcc) at 50; chunk larger
# contact lists so one oversized batch can't fail delivery for everyone.
SES_MAX_RECIPIENTS_PER_EMAIL = 50


def _onboard_store(
    store: dict, client: FivetranClient, sf, existing_by_schema: dict[str, str]
) -> tuple[str, dict | None]:
    domain = store["shop_domain"]
    schema_name = store["schema_name"]
    redirect_uri = config.FIVETRAN_REDIRECT_URI or f"https://{domain}"

    # Idempotency guard: if a connector already exists for this schema in the
    # group, reuse it. Connector creation succeeds before the state write, so a
    # redelivered SQS message (state write failed, or partial batch retry) would
    # otherwise create a duplicate connector for the same store. schema_name is
    # deterministic per store, so it is a safe natural key here.
    connector_id = existing_by_schema.get(schema_name)
    if connector_id:
        logger.info(f"[{domain}] Connector already exists for schema={schema_name}: {connector_id} — reusing")
    else:
        logger.info(f"[{domain}] Creating connector (schema={schema_name})")
        connector = client.create_shopify_connector(
            group_id=config.FIVETRAN_GROUP_ID,
            schema=schema_name,
            shop_domain=domain,
        )
        connector_id = connector.id
        # Guard against a second message in this same batch for the same schema.
        existing_by_schema[schema_name] = connector_id
        logger.info(f"[{domain}] Connector created: {connector_id}")

    current = client.get_connection_status(connector_id)
    if current.status.setup_state == ConnectionSetupState.CONNECTED:
        logger.info(f"[{domain}] Already connected — writing pending_configure")
        upsert_state(
            sf,
            {
                "shop_domain": domain,
                "schema_name": schema_name,
                "connector_id": connector_id,
                "status": "pending_configure",
                "connected_at": datetime.now(UTC).isoformat(),
                "merch_company": store.get("merch_company"),
                "selling_country": store.get("selling_country"),
                "rep_owner": store.get("rep_owner"),
                "alt_myshopify_domain": store.get("alt_myshopify_domain"),
                "custom_domain": store.get("custom_domain"),
            },
        )
        return "pending_configure", None

    card = client.get_connect_card(connector_id, redirect_uri)
    expiry = datetime.now(UTC).replace(microsecond=0) + timedelta(hours=24)

    upsert_state(
        sf,
        {
            "shop_domain": domain,
            "schema_name": schema_name,
            "connector_id": connector_id,
            "status": "pending_oauth",
            "connect_card_uri": card.uri,
            "link_created_at": datetime.now(UTC).isoformat(),
            "link_expires_at": expiry.isoformat(),
            "merch_company": store.get("merch_company"),
            "selling_country": store.get("selling_country"),
            "rep_owner": store.get("rep_owner"),
            "alt_myshopify_domain": store.get("alt_myshopify_domain"),
            "custom_domain": store.get("custom_domain"),
        },
    )
    logger.info(
        f"[{domain}] Connect Card issued (connector_id={connector_id}, "
        f"link_expires_at={expiry.isoformat()}), pending_oauth written"
    )

    # No dedup guard here (unlike existing_by_schema above): a redelivered SQS
    # message that reaches this point re-issues a Connect Card (Fivetran mints a
    # new token/URI per call) and queues a fresh email. This is a known, accepted
    # risk — worst case is a merch company gets a duplicate notification email
    # with a newer link, not data corruption or a duplicate connector.
    email_data = {
        "shop_domain": domain,
        "connect_card_uri": card.uri,
        "link_expires_at": expiry.isoformat(),
        "merch_company": store.get("merch_company"),
        "selling_country": store.get("selling_country"),
    }

    return "pending_oauth", email_data


def _send_batched_emails(email_data_list: list[dict], sf) -> None:
    """Batch stores by (merch_company, selling_country) and send grouped emails.

    Args:
        email_data_list: List of store dicts (shop_domain, connect_card_uri, link_expires_at,
            merch_company, selling_country)
        sf: Snowflake connection
    """
    if not email_data_list:
        return

    # Batch by (merch_company, selling_country)
    batches: dict = defaultdict(list)
    for data in email_data_list:
        company = data.get("merch_company")
        country = data.get("selling_country")
        if not company or not country:
            logger.warning(f"Skipping {data.get('shop_domain')}: missing company/country metadata")
            continue
        key = (company, country)
        batches[key].append(data)

    if not batches:
        return

    ses_client = SESClient(region=config.AWS_REGION)

    for (company, country), stores in batches.items():
        logger.info(f"Sending email batch: {company}/{country} ({len(stores)} store(s))")
        try:
            contacts_array = get_contacts(sf, country, company)
            if contacts_array is None:
                logger.warning(f"No contact mapping found for {company}/{country} — skipping email")
                lambda_metric("d2c.onboard_shopify_stores.email_failed", 1, tags=[f"company:{company}"])
                continue

            contacts = parse_contacts_from_array(contacts_array)
            if not contacts:
                logger.warning(f"No usable contacts for {company}/{country} — skipping email")
                lambda_metric("d2c.onboard_shopify_stores.email_failed", 1, tags=[f"company:{company}"])
                continue

            subject = f"Action Required: Complete Shopify OAuth for {len(stores)} store(s) — {company} {country}"
            body = format_onboarding_email(stores, contacts)
            to_addresses = get_email_addresses(contacts)

            if not to_addresses:
                logger.warning(f"No valid email addresses for {company}/{country} — skipping email")
                lambda_metric("d2c.onboard_shopify_stores.email_failed", 1, tags=[f"company:{company}"])
                continue

            message_ids = []
            failed_chunks = 0
            for chunk in batched(to_addresses, SES_MAX_RECIPIENTS_PER_EMAIL):
                try:
                    message_ids.append(
                        ses_client.send_email(
                            to_addresses=list(chunk),
                            subject=subject,
                            body_text=body,
                            from_address=config.SES_FROM_ADDRESS,
                            source_arn=config.SES_SOURCE_ARN,
                        )
                    )
                except Exception as exc:
                    failed_chunks += 1
                    logger.error(
                        f"Failed to send email chunk ({len(chunk)} recipient(s)) for {company}/{country}: {exc}",
                        exc_info=True,
                    )

            if message_ids:
                logger.info(
                    f"Email sent for {company}/{country}: message_ids={message_ids}, "
                    f"recipient_count={len(to_addresses)}, stores={len(stores)}, failed_chunks={failed_chunks}"
                )
                lambda_metric("d2c.onboard_shopify_stores.email_sent", len(stores), tags=[f"company:{company}"])
            if failed_chunks:
                lambda_metric(
                    "d2c.onboard_shopify_stores.email_chunk_failed", failed_chunks, tags=[f"company:{company}"]
                )
            if not message_ids and failed_chunks:
                # Every chunk failed — no recipient got notified for this batch.
                lambda_metric("d2c.onboard_shopify_stores.email_failed", 1, tags=[f"company:{company}"])
        except Exception as exc:
            logger.error(f"Failed to send email for {company}/{country}: {exc}", exc_info=True)
            lambda_metric("d2c.onboard_shopify_stores.email_failed", 1, tags=[f"company:{company}"])


def handler(event: dict, context) -> dict:
    records = event.get("Records", [])
    logger.info(f"Processing {len(records)} SQS record(s)")

    if not records:
        return {"batchItemFailures": []}

    sf = get_connection()
    batch_failures: list[dict[str, str]] = []
    email_data_list: list[dict] = []

    try:
        with FivetranClient(config.FIVETRAN_API_KEY, config.FIVETRAN_API_SECRET) as client:
            # Fetch existing connectors once per invocation so each store can be
            # deduplicated against them without an API call per message.
            existing_by_schema = {
                c.schema_name: c.id
                for c in client.get_group_connectors(config.FIVETRAN_GROUP_ID)
                if c.service == "shopify"
            }
            logger.info(f"{len(existing_by_schema)} existing connector(s) in group")

            for record in records:
                store = json.loads(record["body"])
                domain = store.get("shop_domain", "unknown")
                schema_name = store.get("schema_name", "unknown")
                message_id = record.get("messageId")
                try:
                    logger.info(f"[{domain}] Processing store (schema={schema_name}, message_id={message_id})")
                    status, email_data = _onboard_store(store, client, sf, existing_by_schema)
                    if email_data:
                        email_data_list.append(email_data)
                    logger.info(f"[{domain}] Done — status={status}")
                    lambda_metric(
                        "d2c.onboard_shopify_stores.onboarded",
                        1,
                        tags=[f"status:{status}"],
                    )
                except Exception as exc:
                    logger.error(f"[{domain}] Failed: {exc}", exc_info=True)
                    lambda_metric("d2c.onboard_shopify_stores.failed", 1)
                    if message_id:
                        batch_failures.append({"itemIdentifier": message_id})
                    else:
                        # Fall back to failing the whole batch if the message id is missing.
                        raise

        # Send batched emails after all stores are processed — best-effort;
        # failures must not cause SQS to retry already-onboarded records.
        if email_data_list:
            try:
                _send_batched_emails(email_data_list, sf)
            except Exception as exc:
                logger.error(f"Email batch sending failed: {exc}", exc_info=True)
                lambda_metric("d2c.onboard_shopify_stores.email_batch_failed", 1)
    finally:
        sf.close()

    return {"batchItemFailures": batch_failures}