# rum.tf

locals {
  rum_domain = var.rum_domain != "" ? var.rum_domain : var.domain_name
}

data "aws_caller_identity" "current" {}

resource "aws_cognito_identity_pool" "rum" {
  identity_pool_name               = "buildonaws-rum"
  allow_unauthenticated_identities = true
}

resource "aws_iam_role" "rum_unauthenticated" {
  name = "buildonaws-rum-unauthenticated"

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect    = "Allow"
      Principal = { Federated = "cognito-identity.amazonaws.com" }
      Action    = "sts:AssumeRoleWithWebIdentity"
      Condition = {
        StringEquals = {
          "cognito-identity.amazonaws.com:aud" = aws_cognito_identity_pool.rum.id
        }
        "ForAnyValue:StringLike" = {
          "cognito-identity.amazonaws.com:amr" = "unauthenticated"
        }
      }
    }]
  })
}

# 최소 권한: PutRumEvents만 허용, 이 app monitor로 스코프 제한
resource "aws_iam_role_policy" "rum_put_events" {
  name = "rum-put-events"
  role = aws_iam_role.rum_unauthenticated.id
  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect   = "Allow"
      Action   = "rum:PutRumEvents"
      Resource = "arn:aws:rum:us-east-1:${data.aws_caller_identity.current.account_id}:appmonitor/${local.rum_domain}"
    }]
  })
}

resource "aws_cognito_identity_pool_roles_attachment" "rum" {
  identity_pool_id = aws_cognito_identity_pool.rum.id
  roles            = { unauthenticated = aws_iam_role.rum_unauthenticated.arn }
}

resource "aws_rum_app_monitor" "site" {
  name   = local.rum_domain
  domain = local.rum_domain

  app_monitor_configuration {
    allow_cookies       = true
    enable_xray         = false
    session_sample_rate = var.rum_sample_rate
    telemetries         = ["performance"]
    guest_role_arn      = aws_iam_role.rum_unauthenticated.arn
    identity_pool_id    = aws_cognito_identity_pool.rum.id
  }
}

<!-- CloudWatch RUM (production only) -->
<script>
  (function(n,i,v,r,s,c,x,z){
    x=window.AwsRumClient={q:[],n:n,i:i,v:v,r:r,c:c};
    window[n]=function(c,p){x.q.push({c:c,p:p})};
    z=document.createElement('script');z.async=true;z.src=s;
    document.head.insertBefore(z,document.head.getElementsByTagName('script')[0]);
  })(
    'cwr',
    'YOUR_APP_MONITOR_ID',
    '1.0.0',
    'us-east-1',
    'https://client.rum.us-east-1.amazonaws.com/1.x/cwr.js',
    {
      sessionSampleRate: 1,
      identityPoolId: 'us-east-1:YOUR_IDENTITY_POOL_ID',
      endpoint: 'https://dataplane.rum.us-east-1.amazonaws.com',
      telemetries: ['performance'],
      allowCookies: true,
      enableXRay: false
    }
  );
</script>

"""CloudWatch RUM page_view 데이터에서 페이지별 speculation rules를 생성한다.

awscrt 필수: CloudFront KeyValueStore data plane API(UpdateKeys)는
SigV4A 서명을 사용하며, botocore는 CRT 확장이 있어야 지원한다.
"""
import json
import logging
import os
import time
from collections import Counter, defaultdict
from urllib.parse import urlsplit

import boto3

logger = logging.getLogger()
logger.setLevel(logging.INFO)

RUM_APP_MONITOR_NAME = os.environ['RUM_APP_MONITOR_NAME']
KVS_ARN = os.environ['KVS_ARN']
S3_BUCKET = os.environ['S3_BUCKET']
S3_PREFIX = os.environ.get('S3_PREFIX', '_speculation')
MIN_TRANSITIONS = int(os.environ.get('MIN_TRANSITIONS', '10'))
CONFIDENCE_THRESHOLD_PRERENDER = float(
    os.environ.get('CONFIDENCE_THRESHOLD_PRERENDER', '0.5'))
MAX_PREFETCH = int(os.environ.get('MAX_PREFETCH', '5'))
EAGERNESS_PRERENDER = os.environ.get('EAGERNESS_PRERENDER', 'moderate')
EAGERNESS_PREFETCH = os.environ.get('EAGERNESS_PREFETCH', 'moderate')
EXCLUDE_PATHS = tuple(
    p for p in os.environ.get('EXCLUDE_PATHS', '').split(',') if p)

PAGE_VIEW_EVENT = 'com.amazon.rum.page_view_event'
KVS_BATCH_SIZE = 100

rum = boto3.client('rum')
s3 = boto3.client('s3')
kvs = boto3.client('cloudfront-keyvaluestore')


def normalize_path(value):
    """URL에서 path만 추출, trailing slash 제거, 빈 값은 '/'."""
    if not value:
        return None
    path = urlsplit(value).path or '/'
    if len(path) > 1:
        path = path.rstrip('/') or '/'
    return path


def collect_transitions():
    """지난 24시간 RUM page_view_event에서 parent→child 전환 추출."""
    after = int((time.time() - 86400) * 1000)
    transitions = defaultdict(Counter)
    next_token = None

    while True:
        kwargs = {
            'Name': RUM_APP_MONITOR_NAME,
            'TimeRange': {'After': after},
            'MaxResults': 100,
        }
        if next_token:
            kwargs['NextToken'] = next_token
        resp = rum.get_app_monitor_data(**kwargs)

        for raw in resp.get('Events', []):
            event = json.loads(raw)
            if event.get('event_type') != PAGE_VIEW_EVENT:
                continue
            details = event.get('event_details', {})
            if isinstance(details, str):
                details = json.loads(details)
            source = normalize_path(details.get('parentPageId'))
            target = normalize_path(details.get('pageId'))
            if source and target and source != target:
                transitions[source][target] += 1

        next_token = resp.get('NextToken')
        if not next_token:
            break
    return transitions


def build_rules(destinations):
    """1위 >= threshold → prerender, 나머지 → prefetch."""
    total = sum(destinations.values())
    if total < MIN_TRANSITIONS:
        return None
    ranked = destinations.most_common()
    prerender, prefetch = [], []
    top_path, top_count = ranked[0]
    if (top_count / total >= CONFIDENCE_THRESHOLD_PRERENDER
            and not any(top_path.startswith(p) for p in EXCLUDE_PATHS)):
        prerender = [top_path]
        ranked = ranked[1:]
    prefetch = [path for path, _ in ranked[:MAX_PREFETCH]]
    rules = {}
    if prefetch:
        rules['prefetch'] = [{'urls': prefetch, 'eagerness': EAGERNESS_PREFETCH}]
    if prerender:
        rules['prerender'] = [{'urls': prerender, 'eagerness': EAGERNESS_PRERENDER}]
    return rules or None


def rules_file_path(page_path):
    if page_path == '/':
        return f'/{S3_PREFIX}/index.json'
    return f'/{S3_PREFIX}{page_path}.json'


def upload_rules(file_path, rules):
    s3.put_object(
        Bucket=S3_BUCKET, Key=file_path.lstrip('/'),
        Body=json.dumps(rules).encode(),
        ContentType='application/speculationrules+json',
        CacheControl='public, max-age=86400',
    )


def update_kvs(entries):
    """ETag 기반 optimistic locking으로 배치 업데이트."""
    etag = kvs.describe_key_value_store(KvsARN=KVS_ARN)['ETag']
    for i in range(0, len(entries), KVS_BATCH_SIZE):
        batch = entries[i:i + KVS_BATCH_SIZE]
        resp = kvs.update_keys(
            KvsARN=KVS_ARN, IfMatch=etag,
            Puts=[{'Key': k, 'Value': v} for k, v in batch],
        )
        etag = resp['ETag']


def handler(event, context):
    transitions = collect_transitions()
    kvs_entries = [('__default__', f'/{S3_PREFIX}/default.json')]
    pages_updated = 0

    for page_path, destinations in transitions.items():
        rules = build_rules(destinations)
        if rules is None:
            continue
        file_path = rules_file_path(page_path)
        upload_rules(file_path, rules)
        kvs_entries.append((page_path, file_path))
        pages_updated += 1

    update_kvs(kvs_entries)
    logger.info(json.dumps({'pages_updated': pages_updated}))
    return {'pages_updated': pages_updated}

# speculation-rules.tf

resource "aws_cloudfront_key_value_store" "speculation" {
  name = var.speculation_kvs_name
}

resource "aws_lambda_function" "speculation_generator" {
  function_name    = "speculation-rules-generator"
  role             = aws_iam_role.speculation_lambda.arn
  runtime          = "python3.12"
  handler          = "index.handler"
  filename         = data.archive_file.speculation_lambda.output_path
  source_code_hash = data.archive_file.speculation_lambda.output_base64sha256
  timeout          = 300
  memory_size      = 256

  environment {
    variables = {
      RUM_APP_MONITOR_NAME           = aws_rum_app_monitor.site.name
      KVS_ARN                        = aws_cloudfront_key_value_store.speculation.arn
      S3_BUCKET                      = aws_s3_bucket.site.bucket
      S3_PREFIX                      = var.speculation_s3_prefix
      MIN_TRANSITIONS                = var.speculation_min_transitions
      CONFIDENCE_THRESHOLD_PRERENDER = var.speculation_confidence_threshold
      MAX_PREFETCH                   = var.speculation_max_prefetch
      EXCLUDE_PATHS                  = ""
    }
  }
}

# EventBridge: 매일 UTC 03:00
resource "aws_cloudwatch_event_rule" "daily_speculation" {
  name                = "daily-speculation-rules-update"
  schedule_expression = var.speculation_schedule
}

resource "aws_cloudwatch_event_target" "speculation_lambda" {
  rule = aws_cloudwatch_event_rule.daily_speculation.name
  arn  = aws_lambda_function.speculation_generator.arn
}

resource "aws_lambda_permission" "eventbridge" {
  statement_id  = "AllowEventBridgeInvoke"
  action        = "lambda:InvokeFunction"
  function_name = aws_lambda_function.speculation_generator.function_name
  principal     = "events.amazonaws.com"
  source_arn    = aws_cloudwatch_event_rule.daily_speculation.arn
}

# 기본 규칙 JSON (KVS에 엔트리가 없는 페이지에 제공)
resource "aws_s3_object" "default_speculation_rules" {
  bucket       = aws_s3_bucket.site.id
  key          = "${var.speculation_s3_prefix}/default.json"
  content      = jsonencode({
    prefetch = [{
      where     = { href_matches = "/*" }
      eagerness = "conservative"
    }]
  })
  content_type  = "application/speculationrules+json"
  cache_control = "public, max-age=86400"
}

import cf from 'cloudfront';

const kvs = cf.kvs();
const DEFAULT_RULES_PATH = '/_speculation/default.json';

function normalizeUri(uri) {
  // viewer-request의 '/index.html' rewrite를 되돌려 KVS key와 매칭
  if (uri.endsWith('/index.html')) {
    uri = uri.slice(0, uri.length - 11);
  }
  if (uri.length > 1 && uri.endsWith('/')) {
    uri = uri.slice(0, uri.length - 1);
  }
  return uri === '' ? '/' : uri;
}

async function handler(event) {
  const request = event.request;
  const response = event.response;

  // HTML 응답에만 삽입
  const contentType = response.headers['content-type']
    ? response.headers['content-type'].value : '';
  if (contentType.indexOf('text/html') === -1) {
    return response;
  }

  // Markdown 응답(verified bot)에는 삽입하지 않음
  const wantsMarkdown = request.headers['x-wants-markdown'];
  if (wantsMarkdown && wantsMarkdown.value === 'true') {
    return response;
  }

  let rulesPath = DEFAULT_RULES_PATH;
  try {
    rulesPath = await kvs.get(normalizeUri(request.uri));
  } catch (e) {
    // Key 미스 또는 KVS 오류 → 기본 규칙
  }

  // Speculation-Rules 헤더는 quoted URL 문자열이어야 함
  response.headers['speculation-rules'] = { value: '"' + rulesPath + '"' };
  return response;
}

  default_cache_behavior {
    # ... 기존 viewer-request association ...

    function_association {
      event_type   = "viewer-response"
      function_arn = aws_cloudfront_function.speculation_rules.arn
    }
  }

{
  "prefetch": [{
    "where": { "href_matches": "/*" },
    "eagerness": "conservative"
  }]
}

if (document.prerendering) {
  document.addEventListener('prerenderingchange', () => {
    trackPageView(); // 사용자가 실제로 이 페이지에 도착할 때만 실행
  });
} else {
  trackPageView();
}

cd infra/lambda/speculation-rules-generator
pip install -r requirements.txt -t build/
cp index.py build/
cd build && zip -r ../speculation-rules-generator.zip .

# HTML 응답에서 Speculation-Rules 헤더 확인
curl -sI https://your-site.com/posts/example/ | grep -i speculation

# 규칙 JSON 파일 확인 (Lambda 실행 후)
curl -s https://your-site.com/_speculation/posts/example.json | jq .

# 정적 리소스에 헤더 없는지 확인
curl -sI https://your-site.com/css/style.css | grep -i speculation