"""AWS Lambda: parse WP Engine nginx logs for AI crawler traffic.

Reads   s3://bsd-wpe-logs-private/logs/nginx/*.apachestyle.log.gz
Writes  s3://bsd-wpe-logs-parsed/events/date=YYYY-MM-DD/site=<site>/<stem>.jsonl
        s3://bsd-wpe-logs-parsed/summary/date=YYYY-MM-DD.json
        s3://bsd-wpe-logs-parsed/manifest/processed.json

The date partition comes from each log LINE's timestamp, not the filename:
WP Engine rotates at 00:17, so 20260709-0017-foo.log.gz mostly holds Jul 8
traffic. One source file can therefore land in two date partitions.
"""

import gzip
import io
import json
import os
import re
from collections import defaultdict
from datetime import datetime, timezone

import boto3

SOURCE_BUCKET = os.environ.get("SOURCE_BUCKET", "bsd-wpe-logs-private")
DEST_BUCKET = os.environ.get("DEST_BUCKET", "bsd-wpe-logs-parsed")
SOURCE_PREFIX = os.environ.get("SOURCE_PREFIX", "logs/nginx/")
MANIFEST_KEY = "manifest/processed.json"

# Only the vhost-bearing apachestyle logs. The .access.log.gz siblings are a
# different format and the logs/error/ tree is not access traffic at all.
SOURCE_SUFFIX = ".apachestyle.log.gz"

# Max source files to process in one invocation, so a first run over a large
# backlog cannot blow the Lambda timeout. Leftovers are picked up next run.
MAX_FILES_PER_RUN = int(os.environ.get("MAX_FILES_PER_RUN", "200"))

_SAFE_SITE_RE = re.compile(r"[^a-z0-9._-]+")


def stem_for(key):
    """logs/nginx/20260709-0017-foo.apachestyle.log.gz -> 20260709-0017-foo"""
    name = key.rsplit("/", 1)[-1]
    if name.endswith(SOURCE_SUFFIX):
        name = name[: -len(SOURCE_SUFFIX)]
    return _SAFE_SITE_RE.sub("_", name.lower())


def safe_site(site):
    """Make a vhost safe for use inside an S3 key."""
    cleaned = _SAFE_SITE_RE.sub("_", site.lower()).strip("._-")
    return cleaned or "unknown"


# --------------------------------------------------------------------------
# Sinks: S3 in Lambda, local directory when testing.
# --------------------------------------------------------------------------


class S3Sink(object):
    def __init__(self, bucket, client):
        self.bucket = bucket
        self.s3 = client

    def put(self, key, body):
        self.s3.put_object(
            Bucket=self.bucket, Key=key,
            Body=body.encode("utf-8"),
            ContentType="application/json",
        )

    def get(self, key):
        try:
            resp = self.s3.get_object(Bucket=self.bucket, Key=key)
            return resp["Body"].read().decode("utf-8")
        except self.s3.exceptions.NoSuchKey:
            return None

    def list_keys(self, prefix):
        keys = []
        paginator = self.s3.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=self.bucket, Prefix=prefix):
            for obj in page.get("Contents", []):
                keys.append(obj["Key"])
        return keys


class LocalSink(object):
    """Mirrors S3Sink onto the filesystem so we can test without touching AWS."""

    def __init__(self, root):
        self.root = root

    def _path(self, key):
        return os.path.join(self.root, key)

    def put(self, key, body):
        path = self._path(key)
        os.makedirs(os.path.dirname(path), exist_ok=True)
        with open(path, "w") as fh:
            fh.write(body)

    def get(self, key):
        path = self._path(key)
        if not os.path.exists(path):
            return None
        with open(path) as fh:
            return fh.read()

    def list_keys(self, prefix):
        base = self._path(prefix)
        keys = []
        for dirpath, _dirs, files in os.walk(base):
            for name in files:
                full = os.path.join(dirpath, name)
                keys.append(os.path.relpath(full, self.root))
        return keys


# --------------------------------------------------------------------------
# Core processing
# --------------------------------------------------------------------------


def process_lines(lines, source_key):
    """Group AI-crawler events by (date, site). Returns (grouped, stats)."""
    from logparser import parse_line

    grouped = defaultdict(list)
    stats = {"lines": 0, "hits": 0, "unparsed": 0}

    for line in lines:
        if not line.strip():
            continue
        stats["lines"] += 1
        try:
            event = parse_line(line, source_key)
        except Exception:
            stats["unparsed"] += 1
            continue
        if event is None:
            continue
        stats["hits"] += 1
        date = event["timestamp"][:10]
        grouped[(date, event["site"])].append(event)

    return grouped, stats


def write_events(sink, grouped, source_key):
    """Write one JSONL shard per (date, site). Returns the dates touched."""
    stem = stem_for(source_key)
    dates = set()
    for (date, site), events in grouped.items():
        key = "events/date=%s/site=%s/%s.jsonl" % (date, safe_site(site), stem)
        body = "\n".join(json.dumps(e, sort_keys=True) for e in events) + "\n"
        sink.put(key, body)
        dates.add(date)
    return dates


def rebuild_summary(sink, date):
    """Recompute summary/date=X.json from every event shard for that date.

    Reading the shards back (rather than summing only this run's events) keeps
    the summary correct on incremental runs and idempotent on re-runs.
    """
    prefix = "events/date=%s/" % date
    by_site = defaultdict(lambda: {"hits": 0, "bytes": 0, "by_platform": defaultdict(int)})
    by_platform = defaultdict(lambda: {"hits": 0, "bytes": 0, "category": None, "sites": set()})
    by_url = defaultdict(lambda: defaultdict(int))
    total_hits = 0
    total_bytes = 0

    for key in sorted(sink.list_keys(prefix)):
        if not key.endswith(".jsonl"):
            continue
        body = sink.get(key)
        if not body:
            continue
        for line in body.splitlines():
            if not line.strip():
                continue
            e = json.loads(line)
            site = e["site"]
            platform = e["ai_platform"]
            nbytes = e.get("bytes", 0)

            total_hits += 1
            total_bytes += nbytes
            by_site[site]["hits"] += 1
            by_site[site]["bytes"] += nbytes
            by_site[site]["by_platform"][platform] += 1
            by_platform[platform]["hits"] += 1
            by_platform[platform]["bytes"] += nbytes
            by_platform[platform]["category"] = e.get("category")
            by_platform[platform]["sites"].add(site)
            by_url[site][e["url"]] += 1

    summary = {
        "date": date,
        "generated_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "totals": {
            "hits": total_hits,
            "bytes": total_bytes,
            "sites": len(by_site),
            "platforms": len(by_platform),
        },
        "by_site": {
            site: {
                "hits": v["hits"],
                "bytes": v["bytes"],
                "by_platform": dict(sorted(v["by_platform"].items(), key=lambda kv: -kv[1])),
            }
            for site, v in sorted(by_site.items(), key=lambda kv: -kv[1]["hits"])
        },
        "by_platform": {
            p: {
                "hits": v["hits"],
                "bytes": v["bytes"],
                "category": v["category"],
                "sites": len(v["sites"]),
            }
            for p, v in sorted(by_platform.items(), key=lambda kv: -kv[1]["hits"])
        },
        "by_url": {
            site: dict(sorted(urls.items(), key=lambda kv: -kv[1]))
            for site, urls in by_url.items()
        },
    }
    sink.put("summary/date=%s.json" % date, json.dumps(summary, indent=2))
    return summary["totals"]


def load_manifest(sink):
    body = sink.get(MANIFEST_KEY)
    if not body:
        return {"processed": {}}
    try:
        data = json.loads(body)
    except ValueError:
        return {"processed": {}}
    data.setdefault("processed", {})
    return data


def save_manifest(sink, manifest):
    manifest["updated_at"] = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    sink.put(MANIFEST_KEY, json.dumps(manifest, indent=2, sort_keys=True))


def list_source_files(s3, bucket, prefix):
    """Return [(key, etag, size)] for apachestyle logs, oldest name first."""
    out = []
    paginator = s3.get_paginator("list_objects_v2")
    for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
        for obj in page.get("Contents", []):
            if not obj["Key"].endswith(SOURCE_SUFFIX):
                continue
            out.append((obj["Key"], obj["ETag"].strip('"'), obj["Size"]))
    out.sort()
    return out


def read_gzip_lines(s3, bucket, key):
    """Stream a gzipped object from S3 and yield decoded lines."""
    resp = s3.get_object(Bucket=bucket, Key=key)
    with gzip.GzipFile(fileobj=io.BytesIO(resp["Body"].read())) as gz:
        for raw in io.TextIOWrapper(gz, encoding="utf-8", errors="replace"):
            yield raw


def run(s3, sink, source_bucket, source_prefix, max_files, log=print):
    manifest = load_manifest(sink)
    processed = manifest["processed"]

    all_files = list_source_files(s3, source_bucket, source_prefix)
    pending = [
        (key, etag, size)
        for key, etag, size in all_files
        if processed.get(key, {}).get("etag") != etag
    ]
    log("source files: %d total, %d pending" % (len(all_files), len(pending)))

    batch = pending[:max_files]
    touched_dates = set()
    totals = {"files": 0, "lines": 0, "hits": 0, "unparsed": 0}

    for key, etag, size in batch:
        grouped, stats = process_lines(read_gzip_lines(s3, source_bucket, key), key)
        dates = write_events(sink, grouped, key)
        touched_dates |= dates

        processed[key] = {
            "etag": etag,
            "size": size,
            "lines": stats["lines"],
            "hits": stats["hits"],
            "unparsed": stats["unparsed"],
            "dates": sorted(dates),
            "processed_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        }
        totals["files"] += 1
        totals["lines"] += stats["lines"]
        totals["hits"] += stats["hits"]
        totals["unparsed"] += stats["unparsed"]
        log("  %s: %d lines, %d hits, %d unparsed -> %s"
            % (key, stats["lines"], stats["hits"], stats["unparsed"], sorted(dates) or "-"))

    for date in sorted(touched_dates):
        t = rebuild_summary(sink, date)
        log("summary date=%s: %d hits across %d sites" % (date, t["hits"], t["sites"]))

    save_manifest(sink, manifest)

    result = dict(totals)
    result["dates"] = sorted(touched_dates)
    result["remaining"] = len(pending) - len(batch)
    return result


def lambda_handler(event, context):
    s3 = boto3.client("s3")
    sink = S3Sink(DEST_BUCKET, s3)
    result = run(s3, sink, SOURCE_BUCKET, SOURCE_PREFIX, MAX_FILES_PER_RUN)
    print(json.dumps(result))
    return result