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

Purpose
- Reads failed raw backend events from S3, decodes base64-encoded "rawData" into the inner JSON
  event, and writes newline-delimited JSON back to S3 organized by event timestamp.
- This function only does decode; the rest mirrors the backend parser's behavior (no extra enrichments).
- This function can backfill both batched and realtime event logger, so the destination prefix is determined by the source object filename prefix.

Required environment variables
- SOURCE_DATE: The day to backfill, format "YYYY/MM/DD" (e.g., "2025/02/25"). Must be set.
- BACKEND_EVENT_LOGGER_BUCKET: Source bucket name. Staging: "suno-staging-backend-event-logger",
  Prod: "suno-prod-backend-event-logger".

Input/output locations
- Input prefix: s3://{BACKEND_EVENT_LOGGER_BUCKET}/error/processing-failed/YYYY/MM/DD
  (month/day zero-padded from SOURCE_DATE)
- Output prefix depends on source object filename prefix:
  - If key basename starts with "backend-event-logger-firehose":
    s3://{BACKEND_EVENT_LOGGER_DEST_BUCKET}/batched/0/YYYY/M/D/H/
  - If key basename starts with "backend-realtime-event-logger-firehose":
    s3://{BACKEND_EVENT_LOGGER_DEST_BUCKET}/realtime/0/YYYY/M/D/H/
  Where YYYY/M/D/H come from the inner event's timestamp (month/day not zero-padded in output).
  Files are named as: event-logger-backfill-{AWS_REGION}-1-YYYY-MM-DD-HH-MM-SS-{uuid}
"""

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

import boto3
from botocore.response import StreamingBody


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: Any) -> datetime:
    # Accept unix epoch (int/float) or ISO8601-like strings; returns UTC datetime
    if ts is None:
        raise ValueError("missing timestamp in backend event")
    # Numeric epoch
    if isinstance(ts, (int, float)):
        return datetime.fromtimestamp(float(ts), tz=timezone.utc)
    # String inputs
    if isinstance(ts, str):
        # Try ISO with trailing Z or explicit offset
        try:
            if ts.endswith("Z"):
                return datetime.fromisoformat(ts.replace("Z", "+00:00")).astimezone(timezone.utc)
            return datetime.fromisoformat(ts).astimezone(timezone.utc)
        except Exception:
            # Try parse as float/int epoch encoded as string
            try:
                return datetime.fromtimestamp(float(ts), tz=timezone.utc)
            except Exception:
                pass
            # Common fallback patterns
            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
    raise ValueError(f"unrecognized timestamp format: {ts}")


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 _dest_subprefix_for_key(src_key: str) -> Optional[str]:
    # Determine destination subprefix based on file basename
    base = src_key.rsplit("/", 1)[-1]
    if base.startswith("backend-event-logger-firehose"):
        return "batched/0"
    if base.startswith("backend-realtime-event-logger-firehose"):
        return "realtime/0"
    return None


def _process_object(
    s3,
    src_bucket: str,
    src_key: str,
    dest_bucket: str,
) -> Tuple[int, int]:
    subprefix = _dest_subprefix_for_key(src_key)
    if not subprefix:
        print(f"skip object with unknown prefix: s3://{src_bucket}/{src_key}")
        return 0, 0

    obj = s3.get_object(Bucket=src_bucket, Key=src_key)
    body = obj["Body"]
    assert isinstance(body, StreamingBody)

    processed = 0
    dropped = 0

    # Group by date/hour; also keep HH-MM-SS for filename from first event in group
    by_group: Dict[Tuple[int, int, int, int], List[str]] = {}
    time_by_group: Dict[Tuple[int, int, int, int], str] = {}

    for raw_line in _iter_s3_lines(body):
        line = raw_line.strip()
        if not line:
            continue
        try:
            wrapper = json.loads(line)
        except Exception as exc:
            print(f"bad json line; key={src_key} err={exc}")
            continue

        raw_b64 = wrapper.get("rawData")
        if not isinstance(raw_b64, str):
            print(f"missing rawData; key={src_key}")
            continue
        try:
            inner = json.loads(base64.b64decode(raw_b64).decode("utf-8"))
        except Exception as exc:
            print(f"rawData decode failed; key={src_key} err={exc}")
            continue

        ts = inner.get("timestamp")
        dt = _parse_event_time(ts) 
        if dt is None:
            print(f"missing/invalid timestamp; key={src_key}")
            dropped += 1
            continue

        y, m, d, h = dt.year, dt.month, dt.day, dt.hour
        by_group.setdefault((y, m, d, h), []).append(json.dumps(inner, separators=(",", ":")))
        if (y, m, d, h) not in time_by_group:
            time_by_group[(y, m, d, h)] = dt.strftime("%H-%M-%S")
        processed += 1

    region = _extract_region_from_key(src_key) or os.getenv("AWS_REGION")
    for (y, m, d, h), lines in by_group.items():
        date_str = f"{y}-{m:02d}-{d:02d}"
        time_str = time_by_group.get((y, m, d, h), "00-00-00")
        file_name = f"event-logger-backfill-{region}-1-{date_str}-{time_str}-{uuid.uuid4()}"
        key_out = f"{subprefix}/{y}/{m}/{d}/{h}/{file_name}"
        _s3_put_ndjson(s3, dest_bucket, key_out, lines)

    print(f"object summary bucket={src_bucket} key={src_key} processed={processed} dropped={dropped}")
    return processed, dropped


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

    # Source date
    source_date = _get_env("SOURCE_DATE")  # YYYY/MM/DD
    parts = source_date.strip().split("/")
    if len(parts) != 3:
        raise RuntimeError("SOURCE_DATE must be 'YYYY/MM/DD'")
    y, m, d = int(parts[0]), int(parts[1]), int(parts[2])

    src_bucket = _get_env("BACKEND_EVENT_LOGGER_BUCKET")
    dest_bucket = src_bucket

    prefix = f"error/processing-failed/{y}/{m:02d}/{d:02d}"

    s3 = boto3.client("s3")

    total_processed = 0
    total_dropped = 0

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

    cont: Optional[str] = None
    while True:
        kwargs: Dict[str, Any] = {"Bucket": src_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"]
            p, q = _process_object(s3, src_bucket, key, dest_bucket)
            total_processed += p
            total_dropped += q
        if page.get("IsTruncated"):
            cont = page.get("NextContinuationToken")
        else:
            break

    elapsed = time.monotonic() - started
    print(f"backend backfill summary src_bucket={src_bucket} prefix={prefix} dest_bucket={dest_bucket} processed={total_processed} dropped={total_dropped} elapsed_s={elapsed:.3f}")

    return {
        "status": "ok",
        "src_bucket": src_bucket,
        "src_prefix": prefix,
        "dest_bucket": dest_bucket,
        "processed": total_processed,
        "dropped": total_dropped,
        "elapsed_s": round(elapsed, 3),
    }