"""
Backfill Lambda for event-logger (manual trigger only)

Purpose
- Reads failed raw events from S3, decodes the outer base64-encoded "rawData" to extract
  request metadata (request-time, domain-name, x-client-ip) and the inner base64 "request_body".
- The inner decoded JSON is the actual event payload and is processed the same way as the
  event-logger-parser Lambda (field validation, shaping, etc.). This function only adds the
  rawData/request_body decoding step; the rest of the logic mirrors event-logger-parser.

Required environment variables
- SOURCE_DATE: The day to backfill, format "YYYY/MM/DD" (e.g., "2025/09/02"). You must set this
  before invoking; otherwise the function will error.
- EVENT_LOGGER_BUCKET: Bucket name to read from and write to. In staging it should be
  "event-logger-s3-temp"; in prod it should be "event-logger-temp". This is set by CDK per stage.

Notes on input/output locations
- Input (lookup) prefix: s3://{EVENT_LOGGER_BUCKET}/eerror/processing-failed/YYYY/MM/DD
  (month/day are zero-padded and come from SOURCE_DATE)
- Output (write) prefix: s3://{EVENT_LOGGER_BUCKET}/{EventName}/0/YYYY/M/D/H/
  - Year/month/day/hour are derived from the event timestamp (month/day not zero-padded in output)
  - File name format: event-logger-backfill-{AWS_REGION}-1-YYYY-MM-DD-{uuid}
    (AWS_REGION is provided by the Lambda runtime, UUID is randomly generated per file)

Other behaviors
- server_timestamp is derived from request_time (UTC) and formatted like
  "YYYY-MM-DD HH:MM:SS.microseconds+00:00". If request_time is missing/invalid, the function raises.
"""

import base64
import hashlib
import json
import logging
import os
import random
import time
import uuid
from datetime import datetime, timezone
from typing import Any, Dict, Iterable, List, Optional, Tuple

import boto3
from botocore.client import BaseClient
from botocore.response import StreamingBody


# Logging (overridable via LOG_LEVEL)
_LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO").upper()
logging.basicConfig(level=getattr(logging, _LOG_LEVEL, logging.INFO))
logger = logging.getLogger(__name__)


MAX_RECORDS_PER_BATCH = 500
MAX_BATCH_BYTES = 5 * 1024 * 1024  # 5 MiB
PUT_RECORDS_MAX_RETRIES = 5

# Backfill processing configuration (copied from event-logger-parser behavior)
whilte_list_event_name = [
    'Audio-Player-Event',
    'Audio-Creation-Event',
    'Audio-Action-Event',
    'Playlist-Action-Event',
    'Link-Click-Event',
    'Web-Page-Event',
    'Artist-Action-Event',
    'General-Event',
    'Audio-Editing-Event',
    'Account-Action-Event',
    'Search-Event',
    'App-Audio-Player-Event',
    'App-Audio-Creation-Event',
    'App-Audio-Action-Event',
    'App-Playlist-Action-Event',
    'App-Link-Click-Event',
    'App-Artist-Action-Event',
    'App-General-Event',
    'App-Audio-Editing-Event',
    'App-Account-Action-Event',
    'App-Search-Event',
    'Web-User-Event',
    'Hook-Web-Event',
]
required_fields = ['event', 'properties', 'timestamp', 'context']

def is_whitelisted_event(json_value: Dict[str, Any]) -> bool:
    try:
        return json_value.get('event') in whilte_list_event_name
    except Exception:
        return False


def is_required_fields_missing(json_value: Dict[str, Any]) -> bool:
    for field in required_fields:
        if field not in json_value:
            return True
    return False


def _parse_s3_path(path: str) -> Tuple[str, str]:
    if not path:
        raise RuntimeError("S3_PATH must be provided, e.g. s3://bucket/prefix")
    if not path.startswith("s3://"):
        raise RuntimeError("S3_PATH must start with s3://")
    without_scheme = path[len("s3://") :]
    parts = without_scheme.split("/", 1)
    bucket = parts[0]
    prefix = parts[1] if len(parts) > 1 else ""
    if not bucket:
        raise RuntimeError("S3_PATH missing bucket")
    return bucket, prefix


def _get_env(name: str) -> str:
    value = os.getenv(name)
    if not value:
        raise RuntimeError(f"Missing required env var: {name}")
    return value


def _iter_s3_lines(body: StreamingBody, chunk_size: int = 64 * 1024) -> Iterable[str]:
    pending = b""
    for chunk in body.iter_chunks(chunk_size=chunk_size):
        if not chunk:
            continue
        pending += chunk
        while True:
            nl = pending.find(b"\n")
            if nl == -1:
                break
            line = pending[:nl]
            pending = pending[nl + 1 :]
            if line:
                yield line.decode("utf-8")
    if pending:
        yield pending.decode("utf-8")


def _parse_event_time(ts: str) -> datetime:
    # Try a few common formats; force UTC when ambiguous
    if not ts:
        return datetime.now(timezone.utc)
    try:
        # Handle trailing Z
        if ts.endswith('Z'):
            return datetime.fromisoformat(ts.replace('Z', '+00:00')).astimezone(timezone.utc)
        return datetime.fromisoformat(ts).astimezone(timezone.utc)
    except Exception:
        pass
    # Fallbacks
    for fmt in (
        "%Y-%m-%dT%H:%M:%S.%f%z",
        "%Y-%m-%dT%H:%M:%S%z",
        "%Y-%m-%d %H:%M:%S%z",
        "%Y-%m-%d %H:%M:%S",
    ):
        try:
            dt = datetime.strptime(ts, fmt)
            if dt.tzinfo is None:
                dt = dt.replace(tzinfo=timezone.utc)
            return dt.astimezone(timezone.utc)
        except Exception:
            continue
    return datetime.now(timezone.utc)


def _s3_put_ndjson(s3, bucket: str, key: str, lines: List[str]) -> None:
    data = "\n".join(lines) + "\n" if lines else ""
    s3.put_object(Bucket=bucket, Key=key, Body=data.encode('utf-8'))


def _extract_region_from_key(key: str) -> Optional[str]:
    base = key.rsplit('/', 1)[-1]
    parts = base.split('-')
    try:
        idx = parts.index('firehose')
        region_parts = parts[idx + 1: idx + 4]
        if len(region_parts) == 3:
            return '-'.join(region_parts)
    except ValueError:
        pass
    return None


def _make_partition_key(raw_data_b64: str, entry: Dict[str, Any]) -> str:
    meta = entry.get("kinesisRecordMetadata")
    if isinstance(meta, dict):
        pk = meta.get("partitionKey")
        if isinstance(pk, str) and pk:
            return pk
    return hashlib.sha256(raw_data_b64.encode("utf-8")).hexdigest()


def _send_put_records_with_retries(
    kinesis: BaseClient, stream_name: str, records: List[Dict[str, Any]]
) -> Tuple[int, int, int]:
    # Kept for compatibility; not used in new S3 backfill flow.
    total_success = 0
    total_failed = 0
    attempts = 0
    to_send = records

    while to_send:
        attempts += 1
        resp = kinesis.put_records(StreamName=stream_name, Records=to_send)
        failed = int(resp.get("FailedRecordCount", 0))
        recs = resp.get("Records", [])

        if failed == 0:
            total_success += len(to_send)
            to_send = []
            break

        retry_list: List[Dict[str, Any]] = []
        for idx, res in enumerate(recs):
            if res.get("ErrorCode"):
                retry_list.append(to_send[idx])
            else:
                total_success += 1

        total_failed += failed
        if attempts >= PUT_RECORDS_MAX_RETRIES:
            break

        sleep_s = (2 ** attempts) * 0.2 + random.random() * 0.3
        time.sleep(sleep_s)
        to_send = retry_list

    return total_success, total_failed, attempts


def _should_split_batch(current_bytes: int, next_data_len: int, num_records: int) -> bool:
    if num_records >= MAX_RECORDS_PER_BATCH:
        return True
    overhead = 8 * num_records  # conservative overhead estimate
    return (current_bytes + next_data_len + overhead) > MAX_BATCH_BYTES


def _process_object_to_s3(
    s3,
    bucket: str,
    key: str,
    dest_bucket: str,
    dest_prefix: str,
) -> Tuple[int, int, int]:
    obj = s3.get_object(Bucket=bucket, Key=key)
    body = obj["Body"]
    assert isinstance(body, StreamingBody)

    line_no = 0
    processed = 0
    dropped = 0

    # Collect results and then write grouped by event name/date
    records: List[Dict[str, Any]] = []

    for raw_line in _iter_s3_lines(body):
        line_no += 1
        line = raw_line.strip()
        if not line:
            continue

        dedup_key = f"{bucket}:{key}:{line_no}"
        try:
            entry = json.loads(line)
        except Exception as exc:
            print(f"Bad JSON at line {line_no}; dedup={dedup_key} err={exc}")
            continue

        raw_b64 = entry.get("rawData")
        if not isinstance(raw_b64, str):
            print(f"Missing rawData at line {line_no}; dedup={dedup_key}")
            continue

        try:
            outer_bytes = base64.b64decode(raw_b64, validate=True)
            outer = json.loads(outer_bytes.decode('utf-8'))
        except Exception as exc:
            print(f"rawData decode/parse error at line {line_no}; dedup={dedup_key} err={exc}")
            continue

        client_ip = outer.get('x-client-ip', '')
        request_time = outer.get('request-time', '')
        domain_name = outer.get('domain-name', '')
        request_body_b64 = outer.get('requst_body')   # notice the typo requst_body
        if not isinstance(request_body_b64, str):
            print(f"Missing request_body after rawData decode at line {line_no}; dedup={dedup_key}")
            continue

        try:
            inner = json.loads(base64.b64decode(request_body_b64).decode('utf-8'))
        except Exception as exc:
            print(f"request_body decode/parse error at line {line_no}; dedup={dedup_key} err={exc}")
            continue

        if is_required_fields_missing(inner) or not is_whitelisted_event(inner):
            dropped += 1
            continue

        result: Dict[str, Any] = {}
        # server_timestamp must reflect request_time; raise if missing or invalid
        if not request_time:
            raise ValueError("request_time missing for event")
        dt_req = datetime.strptime(request_time, "%d/%b/%Y:%H:%M:%S %z").astimezone(timezone.utc)
        tz = dt_req.strftime("%z")
        tz_formatted = tz[:3] + ":" + tz[3:] if len(tz) == 5 else tz
        server_timestamp_str = f"{dt_req.strftime('%Y-%m-%d %H:%M:%S')}.{dt_req.microsecond:06d}{tz_formatted}"
        result['server_timestamp'] = server_timestamp_str
        result['client_ip'] = client_ip
        result['request_time'] = request_time
        result['domain_name'] = domain_name
        for field in required_fields:
            result[field] = inner[field]

        anonymous_id = inner.get('anonymousId')
        if anonymous_id is not None:
            result['anonymous_id'] = anonymous_id

        records.append({'result': result})
        processed += 1

    # Group by (event name, event date & hour) and write to S3 as NDJSON
    by_event_hour: Dict[Tuple[str, int, int, int, int], List[str]] = {}
    time_by_group: Dict[Tuple[str, int, int, int, int], str] = {}
    for rec in records:
        evt = rec['result']
        ts = evt.get('timestamp', '')
        dt = _parse_event_time(ts) if isinstance(ts, str) else datetime.now(timezone.utc)
        y, m, d, h = dt.year, dt.month, dt.day, dt.hour
        event_name = evt.get('event', 'Unknown-Event')
        line = json.dumps(evt, separators=(",", ":"))
        group_key = (event_name, y, m, d, h)
        by_event_hour.setdefault(group_key, []).append(line)
        if group_key not in time_by_group:
            time_by_group[group_key] = dt.strftime("%H-%M-%S")

    region = _extract_region_from_key(key) or os.getenv("AWS_REGION")
    for (event_name, y, m, d, h), lines in by_event_hour.items():
        # Key format: {EventName}/0/year/month/date/hour/event-logger-backfill-<region>-1-YYYY-MM-DD-<uuid>
        date_str = f"{y}-{m:02d}-{d:02d}"
        time_str = time_by_group.get((event_name, y, m, d, h), "00-00-00")
        file_name = f"event-logger-backfill-{region}-1-{date_str}-{time_str}-{uuid.uuid4()}"
        key_out = f"{event_name}/0/{y}/{m}/{d}/{h}/{file_name}"
        _s3_put_ndjson(s3, dest_bucket, key_out, lines)

    print(f"object summary bucket={bucket} key={key} written={sum(len(v) for v in by_event_hour.values())} dropped={dropped}")
    # Detailed per-event counts
    event_counts: Dict[str, int] = {}
    for (ename, _yy, _mm, _dd, _hh), lines in by_event_hour.items():
        event_counts[ename] = event_counts.get(ename, 0) + len(lines)
    if event_counts:
        details = ", ".join(f"{k}:{v}" for k, v in sorted(event_counts.items()))
        print(f"object event_counts: {details}")

    return processed, dropped, 0


def lambda_handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
    started = time.monotonic()

    # Source (fixed bucket and path pattern)
    source_date = _get_env("SOURCE_DATE")  
    try:
        parts = source_date.strip().split("/")
        if len(parts) != 3:
            raise ValueError("SOURCE_DATE must be 'YYYY/MM/DD'")
        y, m, d = int(parts[0]), int(parts[1]), int(parts[2])
    except Exception:
        raise RuntimeError("Invalid SOURCE_DATE; expected 'YYYY/MM/DD', e.g. '2025/09/02'")

    bucket_name = _get_env("EVENT_LOGGER_BUCKET")
    bucket = bucket_name
    # Source requires zero-padded month/day
    prefix = f"eerror/processing-failed/{y}/{m:02d}/{d:02d}"

    # Destination (same bucket; per-event prefixes are determined per record)
    dest_bucket = bucket_name
    dest_prefix = ""  # unused; kept for signature compatibility

    s3 = boto3.client("s3")

    total_written = 0
    total_dropped = 0

    print(f"backfill start src_bucket={bucket} src_prefix={prefix} dest_bucket={dest_bucket}")

    cont: Optional[str] = None
    while True:
        kwargs: Dict[str, Any] = {"Bucket": bucket, "Prefix": prefix}
        if cont:
            kwargs["ContinuationToken"] = cont
        page = s3.list_objects_v2(**kwargs)
        contents = page.get("Contents", [])
        for obj in contents:
            key = obj["Key"]
            written, dropped, _ = _process_object_to_s3(s3, bucket, key, dest_bucket, dest_prefix)
            total_written += written
            total_dropped += dropped
        if page.get("IsTruncated"):
            cont = page.get("NextContinuationToken")
        else:
            break

    elapsed = time.monotonic() - started
    print(f"backfill summary src_bucket={bucket} src_prefix={prefix} dest_bucket={dest_bucket} written={total_written} dropped={total_dropped} elapsed_s={elapsed:.3f}")

    return {
        "status": "ok",
        "src_bucket": bucket,
        "src_prefix": prefix,
        "dest_bucket": dest_bucket,
        "dest_prefix": None,
        "written": total_written,
        "dropped": total_dropped,
        "elapsed_s": round(elapsed, 3),
    }