"""
relay-outbound Lambda

Handles SES Configuration Set events forwarded via EventBridge.
Writes to two DynamoDB tables:

  - ses-outbound-messages : EXISTING table for bounce forensics
                            (kept unchanged so existing tooling keeps working)
  - ses-events            : NEW table for billing-style aggregation
                            of sends/bounces/complaints

We list mailadmin domains to filter out events from non-mailadmin
addresses (e.g. SES console test sends).
"""
import json
import os
import time
from datetime import datetime, timedelta, timezone

import boto3

dynamodb = boto3.resource('dynamodb')
ses_client = boto3.client('ses')

# Existing bounce-detail table — preserved as-is
bounce_table = dynamodb.Table('ses-outbound-messages')

# New unified events table for billing
events_table = dynamodb.Table('ses-events')

# ------------------------------------------------------------
# Domain filter
# ------------------------------------------------------------
# We only want to track events for mailboxes managed by mailadmin.
# Easiest source-of-truth: SES verified domain identities. Any address
# whose domain is verified in SES is "ours". Cached for 5 minutes so
# we don't hammer SES on every event.
# ------------------------------------------------------------

_domain_cache = {'expires_at': 0, 'domains': set()}
_DOMAIN_CACHE_TTL_SEC = 300


def get_managed_domains():
    """Return the set of SES verified domain identities (lowercased)."""
    now = time.time()
    if _domain_cache['expires_at'] > now and _domain_cache['domains']:
        return _domain_cache['domains']

    try:
        # ListIdentities is paginated; for ~50 domains one page is enough,
        # but handle pagination defensively.
        domains = set()
        next_token = None
        while True:
            kwargs = {'IdentityType': 'Domain', 'MaxItems': 100}
            if next_token:
                kwargs['NextToken'] = next_token
            resp = ses_client.list_identities(**kwargs)
            for d in resp.get('Identities', []):
                domains.add(d.lower())
            next_token = resp.get('NextToken')
            if not next_token:
                break

        _domain_cache['domains'] = domains
        _domain_cache['expires_at'] = now + _DOMAIN_CACHE_TTL_SEC
        return domains
    except Exception as exc:
        print(f"WARN: could not list SES identities, falling back to cached set: {exc}")
        return _domain_cache['domains']


def domain_of(email):
    if not email or '@' not in email:
        return None
    return email.split('@', 1)[1].strip().lower()


def is_managed_address(email):
    d = domain_of(email)
    if not d:
        return False
    return d in get_managed_domains()


# ------------------------------------------------------------
# DynamoDB helpers
# ------------------------------------------------------------

def _ttl_two_years_from_now():
    return int((datetime.now(timezone.utc) + timedelta(days=730)).timestamp())


def _month_bucket(ts_iso):
    """
    Convert a SES event timestamp (ISO 8601) to 'YYYY-MM'.
    Falls back to current UTC month if parsing fails.
    """
    try:
        if ts_iso.endswith('Z'):
            ts_iso = ts_iso[:-1] + '+00:00'
        dt = datetime.fromisoformat(ts_iso)
    except Exception:
        dt = datetime.now(timezone.utc)
    return dt.strftime('%Y-%m')


def write_event_row(*, event_type, source, recipients, message_id, timestamp_iso,
                    size_bytes=None, subject=None, bounce_type=None, complaint_type=None):
    """
    Write a row to the unified ses-events table.

    Skips silently if the source address isn't from a managed domain
    (per user choice: ignore non-mailadmin events).
    """
    if not is_managed_address(source):
        print(f"Skipping event for non-managed sender: {source}")
        return

    domain = domain_of(source)
    ym = _month_bucket(timestamp_iso)

    pk = f"{domain}#{source}#{ym}"
    sk = f"{timestamp_iso}#{message_id}"

    item = {
        'pk': pk,
        'sk': sk,
        'event_type': event_type,
        'from': source,
        'domain': domain,
        'recipients': recipients or [],
        'message_id': message_id or '',
        'timestamp': timestamp_iso,
        'ttl_unix': _ttl_two_years_from_now(),
    }
    if size_bytes is not None:
        item['size_bytes'] = int(size_bytes)
    if subject:
        item['subject'] = subject[:200]
    if bounce_type:
        item['bounce_type'] = bounce_type
    if complaint_type:
        item['complaint_type'] = complaint_type

    events_table.put_item(Item=item)
    print(f"ses-events <- {event_type} {pk} {sk}")


# ------------------------------------------------------------
# Existing bounce handler (unchanged behavior, plus the new write)
# ------------------------------------------------------------

def handle_bounce(detail):
    bounce = detail.get('bounce', {}) or {}
    mail = detail.get('mail', {}) or {}

    feedback_id = bounce.get('feedbackId')
    bounce_type = bounce.get('bounceType', 'Unknown')
    bounce_subtype = bounce.get('bounceSubType', 'Unknown')
    bounced_recipients = [r.get('emailAddress') for r in bounce.get('bouncedRecipients', []) if r.get('emailAddress')]
    original_source = mail.get('source')
    original_message_id = mail.get('messageId')
    timestamp = bounce.get('timestamp') or mail.get('timestamp') or datetime.now(timezone.utc).isoformat()

    print(f"Bounce: id={feedback_id} from={original_source} type={bounce_type}/{bounce_subtype}")

    # 1) EXISTING: store full bounce forensics in ses-outbound-messages
    bounce_table.put_item(Item={
        'MessageId': feedback_id,
        'original_message_id': original_message_id,
        'original_source': original_source,
        'bounceType': bounce_type,
        'bounceSubType': bounce_subtype,
        'bouncedRecipients': bounced_recipients,
        'timestamp': timestamp,
        'event_type': 'bounce',
    })

    # 2) NEW: also record into ses-events for billing aggregation
    write_event_row(
        event_type='bounce',
        source=original_source,
        recipients=bounced_recipients,
        message_id=original_message_id or feedback_id,
        timestamp_iso=timestamp,
        bounce_type=f"{bounce_type}/{bounce_subtype}",
    )


# ------------------------------------------------------------
# New: send handler
# ------------------------------------------------------------

def handle_send(detail):
    mail = detail.get('mail', {}) or {}
    source = mail.get('source')
    message_id = mail.get('messageId')
    timestamp = mail.get('timestamp') or datetime.now(timezone.utc).isoformat()

    # SES events do not include the message size directly. The closest
    # info available is mail.headersTruncated and mail.commonHeaders, but
    # neither contains the byte size. We approximate from the sum of all
    # known header values + 0 (we'll log 0 if unknown). If size_bytes is
    # critical, the alternative is to request the 'sendingPoolMetadata'
    # via SES API, but for billing-grade tracking the SES Send Quota
    # itself is per-message anyway.
    common = mail.get('commonHeaders', {}) or {}
    subject = common.get('subject')

    # Recipients: prefer 'to' from commonHeaders, fall back to destination
    to_field = common.get('to') or mail.get('destination') or []
    if isinstance(to_field, str):
        to_field = [to_field]

    # Try to get size from SES if exposed (newer event format includes it
    # under 'mail.headersTruncated' — not always present)
    size_bytes = mail.get('messageSize') or mail.get('size') or 0

    print(f"Send: id={message_id} from={source} to={to_field} size={size_bytes}")

    write_event_row(
        event_type='send',
        source=source,
        recipients=to_field,
        message_id=message_id,
        timestamp_iso=timestamp,
        size_bytes=size_bytes,
        subject=subject,
    )


# ------------------------------------------------------------
# New: complaint handler
# ------------------------------------------------------------

def handle_complaint(detail):
    complaint = detail.get('complaint', {}) or {}
    mail = detail.get('mail', {}) or {}

    source = mail.get('source')
    message_id = mail.get('messageId')
    timestamp = complaint.get('timestamp') or mail.get('timestamp') or datetime.now(timezone.utc).isoformat()
    complaint_type = complaint.get('complaintFeedbackType') or 'unknown'
    complained_recipients = [r.get('emailAddress') for r in complaint.get('complainedRecipients', []) if r.get('emailAddress')]

    print(f"Complaint: id={message_id} from={source} type={complaint_type}")

    write_event_row(
        event_type='complaint',
        source=source,
        recipients=complained_recipients,
        message_id=message_id,
        timestamp_iso=timestamp,
        complaint_type=complaint_type,
    )


# ------------------------------------------------------------
# Entry point
# ------------------------------------------------------------

def lambda_handler(event, context):
    print(f"Received event: {json.dumps(event)[:500]}")

    event_type = event.get('detail-type')
    detail = event.get('detail', {}) or {}

    try:
        if event_type == 'Email Bounced':
            handle_bounce(detail)
        elif event_type == 'Email Sent':
            handle_send(detail)
        elif event_type == 'Email Complaint Received':
            handle_complaint(detail)
        else:
            print(f"Ignoring event type: {event_type}")
    except Exception as exc:
        # Don't fail the Lambda — losing one event is preferable to
        # endless EventBridge retries that pile up DLQs.
        print(f"ERROR processing {event_type}: {exc}")

    return {'statusCode': 200}