"""
lambda-d2c-sync-shopify-data-connectors

Triggered every 15 minutes by EventBridge. Runs a single sweep over all
in-flight connectors in D2C_ONBOARDING_STATE and advances each one:

  pending_oauth     -> check Fivetran; if CONNECTED: write pending_configure and return
                       if card expired: refresh card
  pending_configure -> apply table config, trigger sync, write pending_sync
  pending_sync      -> check Snowflake for SHOP table; if present: sync_complete
                       + upsert registry
"""

import json
import logging
import time
from datetime import datetime, timedelta, UTC
from pathlib import Path

import httpx
from datadog_lambda.metric import lambda_metric

import config
from src.fivetran_client import (
    FivetranClient,
    ConnectionSetupState,
    ConnectionSyncState,
    SchemaUpdate,
    StandardConfigUpdate,
    TableUpdate,
)
from src.snowflake_client import (
    get_connection,
    get_store_info,
    get_stores_by_statuses,
    resolve_gp,
    resolve_vendor_id,
    set_status_by_connector_id,
    shop_table_exists,
    upsert_state,
    upsert_registry,
    set_synced_to_snowflake,
)

_REGISTRY_SOURCE = "D2C"

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

TABLES_CONFIG_PATH = Path(__file__).parent.parent / "tables_config.json"

# Stop starting new stores when fewer than this many seconds remain in the
# invocation, leaving headroom for the in-flight store to finish and for the
# Snowflake connection to close. Deferred stores are picked up next sweep.
INVOCATION_RESERVE_SECONDS = 120

# Give up on a store that never completes OAuth. After this many days the
# connector is marked oauth_timeout and we stop regenerating Connect Cards, so
# abandoned stores stop churning cards every sweep and surface for manual review.
PENDING_OAUTH_MAX_AGE_DAYS = 14

DEFAULT_STATUSES = ["pending_oauth", "pending_configure", "pending_sync"]


def _sleep(seconds: float, why: str) -> None:
    if seconds > 0:
        logger.info(f"Sleeping {int(seconds)}s ({why})")
        time.sleep(seconds)


def _remaining_seconds(context) -> float:
    """Seconds left in this Lambda invocation, or inf when run locally."""
    if context is None or not hasattr(context, "get_remaining_time_in_millis"):
        return float("inf")
    return context.get_remaining_time_in_millis() / 1000.0


def _age_days(iso_ts: str | None) -> float | None:
    """Age in days of an ISO-8601 timestamp, or None if absent/unparseable."""
    if not iso_ts:
        return None
    try:
        dt = datetime.fromisoformat(iso_ts.replace("Z", "+00:00"))
    except ValueError:
        return None
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=UTC)
    return (datetime.now(UTC) - dt).total_seconds() / 86400


def _fivetran_call(label: str, fn, *args):
    # The client (fivetran_client._request) already retries 429/5xx with
    # backoff and honours Retry-After, so this wrapper only applies the
    # inter-call throttle delay — no second retry layer.
    result = fn(*args)
    _sleep(config.FIVETRAN_CALL_DELAY, f"throttle after {label}")
    return result


def _load_enabled_tables() -> frozenset[str]:
    return frozenset(t["table"].upper() for t in json.loads(TABLES_CONFIG_PATH.read_text()) if t["enabled"])


def _configure(
    client: FivetranClient,
    connector_id: str,
    enabled_tables: frozenset[str],
    sync_state: ConnectionSyncState | None = None,
    context=None,
) -> int:
    # Pause if actively syncing so the schema patch applies cleanly. A connector
    # can also arrive PAUSED — e.g. a prior sweep paused it, then died before
    # resuming (timeout / invocation deadline). Either way we must guarantee it
    # is resumed before triggering the sync below, or it stays paused and never
    # syncs, stranding the store in pending_sync forever.
    was_paused_on_entry = sync_state == ConnectionSyncState.PAUSED
    if sync_state == ConnectionSyncState.SYNCING:
        logger.info(f"Pausing {connector_id} before configuring")
        _fivetran_call("pause", client.pause_connector, connector_id)
        _wait_for_sync_stop(client, connector_id, context=context)
    elif was_paused_on_entry:
        logger.warning(f"{connector_id} arrived PAUSED — will resume after configuring")

    _fivetran_call("reload_schema", client.reload_schema, connector_id)
    existing = _fivetran_call("get_schema", client.get_connection_schema_config, connector_id)

    new_schemas, tables_synced = {}, 0
    for sname, schema_cfg in existing.schemas.items():
        new_tables = {}
        for table_name, table in schema_cfg.tables.items():
            in_subset = table_name.upper() in enabled_tables
            if not table.enabled_patch_settings.allowed:
                if table.enabled and in_subset:
                    tables_synced += 1
                continue
            new_tables[table_name] = TableUpdate(enabled=in_subset)
            if in_subset:
                tables_synced += 1
        new_schemas[sname] = SchemaUpdate(tables=new_tables)

    _fivetran_call(
        "patch_schema", client.patch_connection_schema_config, connector_id, StandardConfigUpdate(schemas=new_schemas)
    )

    if sync_state == ConnectionSyncState.SYNCING or was_paused_on_entry:
        _fivetran_call("resume", client.resume_connector, connector_id)

    _fivetran_call("trigger_sync", client.trigger_sync, connector_id)
    return tables_synced


def _wait_for_sync_stop(client: FivetranClient, connector_id: str, max_wait: int = 300, context=None) -> None:
    # Never wait past the invocation budget: cap the wait by whatever time is
    # left minus the reserve, so a paused connector can't run the Lambda into
    # its hard timeout mid-configure.
    budget = _remaining_seconds(context) - INVOCATION_RESERVE_SECONDS
    limit = min(max_wait, budget) if budget != float("inf") else max_wait
    if limit <= 0:
        raise TimeoutError(f"{connector_id} still SYNCING — no invocation time left to wait")
    waited = 0
    while waited < limit:
        conn = _fivetran_call("get_status", client.get_connection_status, connector_id)
        if conn.status.sync_state != ConnectionSyncState.SYNCING:
            return
        sleep_time = min(15, limit - waited)
        time.sleep(sleep_time)
        waited += sleep_time + config.FIVETRAN_CALL_DELAY
    raise TimeoutError(f"{connector_id} still SYNCING after {int(limit)}s")


def _handle_pending_oauth(client, sf, store) -> str:
    domain = store["shop_domain"]
    connector_id = store["connector_id"]

    conn = _fivetran_call("get_status", client.get_connection_status, connector_id)

    if conn.status.setup_state == ConnectionSetupState.CONNECTED:
        logger.info(f"[{domain}] CONNECTED — queuing configure for next sweep")
        set_status_by_connector_id(sf, connector_id, "pending_configure", connected_at=datetime.now(UTC).isoformat())
        return "pending_configure"

    if conn.status.setup_state == ConnectionSetupState.BROKEN:
        logger.warning(f"[{domain}] connector is BROKEN — skipping card refresh")
        return "pending_oauth"

    # Give up on stores that never complete OAuth. Anchored on the connector's
    # creation time (stable), not link_created_at, which we overwrite on every
    # card refresh. Abandoned stores become oauth_timeout instead of churning cards.
    age = _age_days(conn.created_at)
    if age is None:
        logger.warning(f"[{domain}] connector created_at missing — skipping OAuth timeout check")
    elif age >= PENDING_OAUTH_MAX_AGE_DAYS:
        logger.warning(f"[{domain}] OAuth not completed after {age:.0f}d — marking oauth_timeout")
        set_status_by_connector_id(
            sf,
            connector_id,
            "oauth_timeout",
            error_message=f"OAuth not completed within {PENDING_OAUTH_MAX_AGE_DAYS} days",
        )
        lambda_metric("d2c.sync_shopify_data_connectors.pending_oauth_abandoned", 1)
        return "oauth_timeout"

    expires_at = store.get("link_expires_at")
    redirect_uri = config.FIVETRAN_REDIRECT_URI or f"https://{domain}"
    if expires_at is not None:
        if isinstance(expires_at, str):
            expires_at = datetime.fromisoformat(expires_at)
        if expires_at.tzinfo is None:
            expires_at = expires_at.replace(tzinfo=UTC)
        if expires_at <= datetime.now(UTC):
            try:
                card = _fivetran_call("get_connect_card", client.get_connect_card, connector_id, redirect_uri)
                new_expiry = datetime.now(UTC) + timedelta(hours=24)
                upsert_state(
                    sf,
                    {
                        "shop_domain": domain,
                        "schema_name": store["schema_name"],
                        "connector_id": connector_id,
                        "status": "pending_oauth",
                        "connect_card_uri": card.uri,
                        "link_created_at": datetime.now(UTC).isoformat(),
                        "link_expires_at": new_expiry.isoformat(),
                    },
                )
                logger.info(f"[{domain}] Connect Card refreshed")
            except Exception as exc:
                logger.error(f"[{domain}] Card refresh failed: {exc}", exc_info=True)
    return "pending_oauth"


def _handle_pending_configure(client, sf, store, enabled_tables, context=None) -> str:
    domain = store["shop_domain"]
    connector_id = store["connector_id"]

    conn = _fivetran_call("get_status", client.get_connection_status, connector_id)
    setup = conn.status.setup_state

    if setup == ConnectionSetupState.BROKEN:
        set_status_by_connector_id(sf, connector_id, "broken", error_message="connector is BROKEN")
        return "broken"

    if setup != ConnectionSetupState.CONNECTED:
        # A connector can regress between the sweep that set pending_configure
        # and this one — e.g. OAuth revoked / app uninstalled leaves it
        # INCOMPLETE. Don't run reload/patch/trigger against a non-connected
        # connector (it errors on Fivetran's side and churns here). Send it back
        # to pending_oauth so the OAuth handler re-issues a card for re-auth and
        # the 14-day give-up eventually marks it broken.
        logger.warning(f"[{domain}] setup_state={setup} (not connected) — reverting to pending_oauth")
        set_status_by_connector_id(sf, connector_id, "pending_oauth")
        return "pending_oauth"

    connected_at = datetime.now(UTC).isoformat()
    tables_synced = _configure(client, connector_id, enabled_tables, conn.status.sync_state, context=context)

    # Use the connector's real destination schema, not the derived name in state —
    # they can diverge, and the SHOP-table check must look where Fivetran writes.
    schema_name = conn.name

    if shop_table_exists(sf, schema_name):
        _upsert_registry_on_complete(
            sf,
            store,
            connector_id,
            schema_name,
            connected_at=connected_at,
            tables_synced=tables_synced,
            sync_frequency=conn.sync_frequency,
        )
        set_status_by_connector_id(
            sf,
            connector_id,
            "sync_complete",
            connected_at=connected_at,
            tables_synced=tables_synced,
            sync_frequency=conn.sync_frequency,
        )
        set_synced_to_snowflake(sf, domain, True)
        logger.info(f"[{domain}] SHOP present — sync_complete")
        return "sync_complete"

    set_status_by_connector_id(
        sf,
        connector_id,
        "pending_sync",
        connected_at=connected_at,
        tables_synced=tables_synced,
        sync_frequency=conn.sync_frequency,
    )
    # Carry the metadata already on the state record so the registry row stays
    # complete even if Lambda 1's initial registry write failed (the pipeline
    # explicitly allows that). upsert_registry uses COALESCE(staged, target) for
    # most columns, so NULLs are never written over existing data; rep_owner uses
    # COALESCE(target, staged) so Lambda 1's write is always preserved.
    upsert_registry(
        sf,
        {
            "myshopify_domain": domain,
            "schema_name": schema_name,
            "source": _REGISTRY_SOURCE,
            "connector_id": connector_id,
            "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_sync"


def _handle_pending_sync(client, sf, store) -> str:
    domain = store["shop_domain"]
    connector_id = store["connector_id"]

    # Resolve the real schema from the connector itself, keyed on connector_id —
    # never trust the derived schema_name stored in state (it can drift).
    conn = _fivetran_call("get_status", client.get_connection_status, connector_id)
    schema_name = conn.name

    if shop_table_exists(sf, schema_name):
        _upsert_registry_on_complete(sf, store, connector_id, schema_name)
        set_status_by_connector_id(sf, connector_id, "sync_complete")
        set_synced_to_snowflake(sf, domain, True)
        logger.info(f"[{domain}] SHOP table landed — sync_complete")
        return "sync_complete"

    return "pending_sync"


def _upsert_registry_on_complete(
    sf,
    store: dict,
    connector_id: str,
    schema_name: str | None,
    connected_at: str | None = None,
    tables_synced: int | None = None,
    sync_frequency: int | None = None,
) -> None:
    domain = store["shop_domain"]
    merch_company = store.get("merch_company")

    # For pending_sync stores the timing fields are already in the state record;
    # for the configure fast-path they are passed in directly.
    effective_connected_at = connected_at if connected_at is not None else store.get("connected_at")
    effective_tables_synced = tables_synced if tables_synced is not None else store.get("tables_synced")
    effective_sync_frequency = sync_frequency if sync_frequency is not None else store.get("sync_frequency")

    try:
        store_name, shopify_store_id, country_code, currency = get_store_info(sf, schema_name)
    except Exception as exc:
        logger.error(f"[{domain}] store info lookup failed: {exc}", exc_info=True)
        lambda_metric("d2c.sync_shopify_data_connectors.store_info_lookup_failed", 1)
        store_name, shopify_store_id, country_code, currency = None, None, None, None

    # Resolve GP first; vendor is then derived from that same GP so the two
    # stay consistent on the registry row.
    try:
        gp_id, gp_name = resolve_gp(sf, store_name, shop_domain=domain)
    except Exception as exc:
        logger.error(f"[{domain}] GP resolution failed, writing registry without it: {exc}", exc_info=True)
        lambda_metric("d2c.sync_shopify_data_connectors.gp_resolution_failed", 1)
        gp_id, gp_name = None, None

    try:
        vendor_id, vendor_name, vendor_brand, vendor_resolution_method = resolve_vendor_id(
            sf, gp_id, merch_company, shop_domain=domain
        )
    except Exception as exc:
        # A raised failure (statement timeout / error) is distinct from a clean
        # no-match, which returns None without raising. The metric fires only
        # here, so systematic timeouts surface rather than hiding as NULLs.
        logger.error(f"[{domain}] vendor resolution failed, writing registry without it: {exc}", exc_info=True)
        lambda_metric("d2c.sync_shopify_data_connectors.vendor_resolution_failed", 1)
        vendor_id, vendor_name, vendor_brand, vendor_resolution_method = None, None, None, None

    upsert_registry(
        sf,
        {
            "shopify_store_id": shopify_store_id,
            "myshopify_domain": domain,
            "schema_name": schema_name,
            "store_name": store_name,
            "country_code": country_code,
            "currency": currency,
            "source": _REGISTRY_SOURCE,
            "connector_id": connector_id,
            "connected_at": effective_connected_at,
            "merch_company": 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"),
            "sync_frequency": effective_sync_frequency,
            "tables_synced": effective_tables_synced,
            "vendor_id": vendor_id,
            "vendor_name": vendor_name,
            "vendor_brand": vendor_brand,
            "vendor_resolution_method": vendor_resolution_method,
            "global_participant_id": gp_id,
            "global_participant_name": gp_name,
            "is_synced_to_snowflake": True,
        },
    )


def handler(event, context):
    logger.info("Starting connector sync sweep")
    enabled_tables = _load_enabled_tables()

    sf = get_connection()
    counts = {
        "pending_oauth": 0,
        "pending_configure": 0,
        "pending_sync": 0,
        "sync_complete": 0,
        "oauth_timeout": 0,
        "broken": 0,
        "error": 0,
        "deferred": 0,
    }

    try:
        stores = get_stores_by_statuses(sf, DEFAULT_STATUSES)
        logger.info(f"Found {len(stores)} in-flight connector(s)")

        with FivetranClient(config.FIVETRAN_API_KEY, config.FIVETRAN_API_SECRET) as client:
            for i, store in enumerate(stores):
                # Stop starting new stores once the invocation budget is nearly
                # spent; the rest advance on the next scheduled sweep.
                if _remaining_seconds(context) < INVOCATION_RESERVE_SECONDS:
                    deferred = len(stores) - i
                    counts["deferred"] += deferred
                    logger.warning(f"Invocation deadline approaching — deferring {deferred} store(s) to next sweep")
                    break

                domain = store["shop_domain"]
                status = store["status"]
                connector_id = store.get("connector_id")

                if not connector_id:
                    logger.warning(f"[{domain}] No connector ID — skipping")
                    continue

                try:
                    if status == "pending_oauth":
                        result = _handle_pending_oauth(client, sf, store)
                    elif status == "pending_configure":
                        result = _handle_pending_configure(client, sf, store, enabled_tables, context)
                    elif status == "pending_sync":
                        result = _handle_pending_sync(client, sf, store)
                    else:
                        result = status
                    counts[result] = counts.get(result, 0) + 1
                except TimeoutError as exc:
                    # Budget exhausted or sync didn't stop in time — expected on
                    # busy sweeps. Connector may be PAUSED; the next sweep handles
                    # it via the was_paused_on_entry branch in _configure.
                    logger.warning(f"[{domain}] Deferred (timeout): {exc}")
                    counts["deferred"] += 1
                except httpx.HTTPStatusError as exc:
                    logger.error(f"[{domain}] HTTP error: {exc.response.status_code}")
                    counts["error"] += 1
                except Exception as exc:
                    logger.error(f"[{domain}] Error: {exc}", exc_info=True)
                    counts["error"] += 1
    finally:
        sf.close()

    # Include deferred stores: they were not processed this sweep but are still
    # in flight, so excluding them under-reports in_flight on busy sweeps.
    still_in_flight = (
        counts["pending_oauth"] + counts["pending_configure"] + counts["pending_sync"] + counts["deferred"]
    )
    lambda_metric("d2c.sync_shopify_data_connectors.in_flight", still_in_flight)
    lambda_metric("d2c.sync_shopify_data_connectors.sync_complete", counts["sync_complete"])
    lambda_metric("d2c.sync_shopify_data_connectors.errors", counts["error"])
    lambda_metric("d2c.sync_shopify_data_connectors.deferred", counts["deferred"])

    logger.info(f"Sweep complete: {counts}")
    return counts