#!/usr/bin/env python3
"""Create a 500-fan test subset from the DDB export for smoke testing.

Reads the first .json.gz data file from the existing DDB export, filters
for spotify-presave records, takes the first N (default 500), and uploads
to a smoke-test prefix in the same S3 bucket with a matching manifest.

Usage:
    awsume aws_dev
    python3 scripts/create_test_subset.py [--count 500]

Prerequisites:
    pip install boto3
"""

import argparse
import gzip
import io
import json
import sys

import boto3

BUCKET = "dev-mymac80"
SOURCE_PREFIX = "resonance-engine/ddb-export/01771804222395-8dd0c6a7/"
DEST_PREFIX = "resonance-engine/ddb-export/smoke-test/"
DEST_DATA_KEY = f"{DEST_PREFIX}data/test-500.json.gz"
DEST_MANIFEST_KEY = f"{DEST_PREFIX}manifest-files.json"


def list_data_files(s3):
    """List .json.gz data files in the source export."""
    prefix = f"{SOURCE_PREFIX}data/"
    response = s3.list_objects_v2(Bucket=BUCKET, Prefix=prefix)
    files = []
    for obj in response.get("Contents", []):
        if obj["Key"].endswith(".json.gz"):
            files.append(obj["Key"])
    return sorted(files)


def is_spotify_presave(line_bytes):
    """Quick check if a DDB export line is a spotify-presave record.

    Avoids full JSON parse + deserialization for non-matching records.
    """
    return b"task:spotify-presave" in line_bytes


def extract_presave_records(s3, source_key, count):
    """Stream a .json.gz file and extract the first N presave records."""
    print(f"Reading s3://{BUCKET}/{source_key} ...")
    response = s3.get_object(Bucket=BUCKET, Key=source_key)

    records = []
    total_lines = 0
    with gzip.GzipFile(fileobj=response["Body"]) as gz:
        for line in gz:
            total_lines += 1
            if not is_spotify_presave(line):
                continue

            line_str = line.decode("utf-8").strip()
            if not line_str:
                continue

            try:
                raw = json.loads(line_str)
            except json.JSONDecodeError:
                continue

            records.append(raw)
            if len(records) >= count:
                break

            if len(records) % 100 == 0:
                print(
                    f"  Found {len(records)}/{count} presave records "
                    f"({total_lines} lines scanned) ..."
                )

    print(
        f"  Extracted {len(records)} presave records from "
        f"{total_lines} lines"
    )
    return records


def upload_test_data(s3, records):
    """Compress records to .json.gz and upload to S3."""
    buf = io.BytesIO()
    with gzip.GzipFile(fileobj=buf, mode="wb") as gz:
        for record in records:
            gz.write(json.dumps(record).encode("utf-8"))
            gz.write(b"\n")

    buf.seek(0)
    size_kb = buf.getbuffer().nbytes / 1024
    print(f"Uploading {len(records)} records ({size_kb:.1f} KB compressed)")
    print(f"  -> s3://{BUCKET}/{DEST_DATA_KEY}")
    s3.put_object(Bucket=BUCKET, Key=DEST_DATA_KEY, Body=buf.getvalue())


def upload_manifest(s3, record_count):
    """Create and upload manifest-files.json for the test subset."""
    manifest_entry = {
        "itemCount": record_count,
        "dataFileS3Key": DEST_DATA_KEY,
    }
    manifest_body = json.dumps(manifest_entry) + "\n"

    print(f"Uploading manifest -> s3://{BUCKET}/{DEST_MANIFEST_KEY}")
    s3.put_object(
        Bucket=BUCKET,
        Key=DEST_MANIFEST_KEY,
        Body=manifest_body.encode("utf-8"),
    )


def main():
    parser = argparse.ArgumentParser(
        description="Create a test subset from DDB export for smoke testing"
    )
    parser.add_argument(
        "--count",
        type=int,
        default=500,
        help="Number of spotify-presave records to extract (default: 500)",
    )
    args = parser.parse_args()

    s3 = boto3.client("s3")

    # Find the first data file
    data_files = list_data_files(s3)
    if not data_files:
        print(f"ERROR: No .json.gz files found at s3://{BUCKET}/{SOURCE_PREFIX}data/")
        sys.exit(1)

    print(f"Found {len(data_files)} data file(s)")
    source_key = data_files[0]

    # Extract presave records
    records = extract_presave_records(s3, source_key, args.count)
    if not records:
        print("ERROR: No spotify-presave records found")
        sys.exit(1)

    # Upload test data + manifest
    upload_test_data(s3, records)
    upload_manifest(s3, len(records))

    print()
    print("Smoke test data ready!")
    print(f"  Data:     s3://{BUCKET}/{DEST_DATA_KEY}")
    print(f"  Manifest: s3://{BUCKET}/{DEST_MANIFEST_KEY}")
    print()
    print("Invoke the Manifest Parser with:")
    print(f'  aws lambda invoke \\')
    print(f'    --function-name dev-lambda-resonance-manifest-parser \\')
    print(f'    --payload \'{{"bucket": "{BUCKET}", '
          f'"export_prefix": "{DEST_PREFIX}"}}\' \\')
    print(f'    --cli-binary-format raw-in-base64-out \\')
    print(f'    /tmp/manifest-parser-response.json')


if __name__ == "__main__":
    main()