using System.Net;
using System.Text.Json;
using System.Text.RegularExpressions;
using Amazon.Lambda.Core;
using Amazon.Lambda.S3Events;
using Amazon.S3;
using Amazon.S3.Model;
using Amazon.SimpleSystemsManagement;
using Amazon.SimpleSystemsManagement.Model;
using Amazon.Textract;
using Amazon.Textract.Model;
using SixLabors.ImageSharp;
using SixLabors.ImageSharp.PixelFormats;


[assembly: LambdaSerializer(typeof(Amazon.Lambda.Serialization.SystemTextJson.DefaultLambdaJsonSerializer))]

namespace maskingFunction;

public class Function
{
    private static readonly IAmazonS3 S3 = new AmazonS3Client();
    private static readonly IAmazonTextract Textract = new AmazonTextractClient();
    private static readonly IAmazonSimpleSystemsManagement Ssm = new AmazonSimpleSystemsManagementClient();

    // 以下の項目は、S3バケットのプレフィックス設定および、ParameterStoreの設定項目に合わせてください。
    
    // Parameter Storeで、このLambdaから呼び出すパラメータのパス
    const string ParamPath = "/masking/";
    // 接頭辞一覧を保存するパラメータ名
    const string PrefixParamName = "prefix";
    // 共通の接尾辞を保存するパラメータ名
    const string SuffixParamName = "suffix";
    // 固有の正規表現一覧を保存するパラメータ名
    const string UniqueParamName = "unique";
    // 出力先プレフィックス
    const string DestinationPrefix = "edited/";

    //設定項目に合わせる値はここまでです。

    // パディング幅
    const int Padding = 4;

    // Parameter Storeから取得した設定を保持
    private static MaskingSettings? _settings;

    /// <summary>
    /// Lambdaのエントリーポイント
    /// S3のイベントを受け取り、画像内のPIIをマスキングしてS3に保存する。
    /// </summary>
    /// <param name="evnt">S3から送られるイベント</param>
    /// <param name="context">Lambdaの実行コンテキスト</param>
    public async Task<string> FunctionHandler(S3Event evnt, ILambdaContext context)
    {
        // ParameterStoreから設定値を読み取る、既に読み取り済なら何もしない
        var settings = _settings ??= await LoadSettingsAsync(context);

        // 処理済みのオブジェクトカウント
        var masked = 0;

        foreach (var record in evnt.Records)
        {
            var bucket = record.S3.Bucket.Name;

            // S3イベントからオブジェクトのキー情報を取得
            var key = WebUtility.UrlDecode(record.S3.Object.Key.Replace("+", "%20"));

            // 出力先を設定
            var fileName = key[(key.LastIndexOf('/') + 1)..];
            var outputKey = DestinationPrefix + fileName;

            // マスキング処理
            await MaskObject(bucket, key, bucket, outputKey, settings, context);

            masked++;
        }

        //　処理済みのオブジェクト数を返す
        return $"masked {masked} object(s)";
    }

    /// <summary>
    /// マスク処理を行う関数
    /// </summary>
    /// <param name="bucket"> 元データバケット</param>
    /// <param name="key"> 加工を行う画像データのキー </param>
    /// <param name="destinationBucket"> 加工後画像の保存先バケット </param>
    /// <param name="outputKey"> 保存先のキー </param>
    /// <param name="settings">  </param>
    /// <param name="context"> Lambdaの実行コンテキスト </param>
    private static async Task MaskObject(
        string bucket, string key, string destinationBucket, string outputKey, MaskingSettings settings, ILambdaContext context)
    {
        // TextractにS3オブジェクトを通知し、読み込ませる
        var detected = await Textract.DetectDocumentTextAsync(new DetectDocumentTextRequest
        {
            Document = new Document
            {
                // S3のバケット名とオブジェクトのキーを通知
                S3Object = new Amazon.Textract.Model.S3Object { Bucket = bucket, Name = key }
            }
        });

        var targets = detected.Blocks
            // BoundingBox が無いブロックは除外する。
            .Where(b => b.BlockType == BlockType.WORD && b.Geometry.BoundingBox != null)
            // 登録した正規表現のいずれかにマッチしたらマスキング対象とみなす。
            .Where(b => settings.Patterns.Any(p => p.IsMatch(b.Text)))
            .Select(b => b.Geometry.BoundingBox)
            .ToList();

        // CloudWatch Logsに処理したログを出力する
        context.Logger.LogInformation(
            $"{key}: words={detected.Blocks.Count(b => b.BlockType == BlockType.WORD)}, pii={targets.Count}");

        // 検出が無い場合は処理を終える
        if (targets.Count == 0)
        {
            return;
        }

        // S3から元画像を取得する
        using var response = await S3.GetObjectAsync(bucket, key);
        using var source = new MemoryStream();
        await response.ResponseStream.CopyToAsync(source);
        source.Position = 0; // コピー直後は末尾にいるので先頭に戻す

        // 元画像のフォーマット（PNG / JPEGなど）を判定しておく。
        // 保存時に同じフォーマットで書き戻すために使う。
        var format = await Image.DetectFormatAsync(source);
        source.Position = 0; // 判定でストリームを読み進めたので再度先頭に戻す

        using var image = await Image.LoadAsync<Rgba32>(source);

        // 検出した領域を黒く塗り潰す
        foreach (var box in targets)
        {
            FillBlack(image, box, Padding);
        }

        // マスキング済み画像をS3に保存する。
        var encoder = Configuration.Default.ImageFormatsManager.FindEncoder(format);

        using var output = new MemoryStream();
        await image.SaveAsync(output, encoder);
        output.Position = 0;

        await S3.PutObjectAsync(new PutObjectRequest
        {
            BucketName = destinationBucket,
            Key = outputKey,
            InputStream = output,
            ContentType = format.DefaultMimeType
        });

        context.Logger.LogInformation($"saved: s3://{destinationBucket}/{outputKey}");
    }

    /// <summary>
    /// BoundingBoxが指定する座標を塗りつぶす処理を行う
    /// </summary>
    /// <param name="image">加工対象の画像</param>
    /// <param name="box">Textractが返した矩形。値は0.0〜1.0の相対座標</param>
    /// <param name="padding">文字の端が残らないように矩形を上下左右に広げるピクセル数</param>
    private static void FillBlack(Image<Rgba32> image, BoundingBox box, int padding)
    {
        double boxLeft = box.Left.GetValueOrDefault();
        double boxTop = box.Top.GetValueOrDefault();
        double boxWidth = box.Width.GetValueOrDefault();
        double boxHeight = box.Height.GetValueOrDefault();

        // TextractのBoundingBoxは画像サイズに依存しない割合で返ってくるため、画像の幅・高さを掛けて実ピクセル座標に変換する
        // paddingで矩形を少し広げ、文字の縁が黒塗りの外に残るのを防ぐ
        var left = Math.Clamp((int)(boxLeft * image.Width) - padding, 0, image.Width - 1);
        var top = Math.Clamp((int)(boxTop * image.Height) - padding, 0, image.Height - 1);
        var right = Math.Clamp((int)((boxLeft + boxWidth) * image.Width) + padding, left + 1, image.Width);
        var bottom = Math.Clamp((int)((boxTop + boxHeight) * image.Height) + padding, top + 1, image.Height);

        image.ProcessPixelRows(accessor =>
        {
            for (var y = top; y < bottom; y++)
            {
                accessor.GetRowSpan(y)[left..right].Fill(Color.Black.ToPixel<Rgba32>());
            }
        });
    }

    /// <summary>
    /// ParameterStoreから設定を取得する
    /// </summary>
    private static async Task<MaskingSettings> LoadSettingsAsync(ILambdaContext context)
    {
        var values = new Dictionary<string, string>();
        string? nextToken = null;

        // GetParametersByPathメソッドで、パス配下のパラメータを全て取得する
        do {
            var response = await Ssm.GetParametersByPathAsync(new GetParametersByPathRequest
            {
                Path = ParamPath,
                Recursive = true, 
                NextToken = nextToken
            });

            foreach (var parameter in response.Parameters)
            {
                values[parameter.Name[(parameter.Name.LastIndexOf('/') + 1)..]] = parameter.Value;
            }
            nextToken = response.NextToken;
        } while(!string.IsNullOrEmpty(nextToken));

        // PIIの判定用パラメータを変数に代入
        var resourceIdPrefixes = values[PrefixParamName].Split(','); 
        // AWSリソースの共通接尾辞を取得
        var resourceIdSuffix = values[SuffixParamName];
        // 上記2つを結合する
        var resourceIdPattern = "(?:" + string.Join("|", resourceIdPrefixes.Select(Regex.Escape)) + ")" + resourceIdSuffix;

        // 固有の正規表現を取得
        var patterns = JsonSerializer.Deserialize<string[]>(values[UniqueParamName])!
            .Append(resourceIdPattern)
            .Select(p => new Regex(p, RegexOptions.Compiled | RegexOptions.IgnoreCase))
            .ToArray();

        // 取得した設定値を返す
        return new MaskingSettings(Patterns: patterns);
    }

    /// <summary>
    /// Parameter Storeから取得した値を保持する構造体
    /// </summary>
    private sealed record MaskingSettings
    (
        IReadOnlyCollection<Regex> Patterns // PIIを判定する正規表現のCollection
    );
}

  {
    "Version": "2012-10-17",
    "Statement": [
        {
            "Sid": "ParameterStore",
            "Effect": "Allow",
            "Action": "ssm:GetParametersByPath",
            "Resource": "{YOUR-SSM-ARN}/*"
        },
        {
            "Sid": "S3ReadRaw",
            "Effect": "Allow",
            "Action": "s3:GetObject",
            "Resource": "{YOUR-S3BUCKET-ARN}/raw/*"
        },
        {
            "Sid": "S3WriteEdited",
            "Effect": "Allow",
            "Action": "s3:PutObject",
            "Resource": "{YOUR-S3BUCKET-ARN}/edited/*"
        },
        {
            "Sid": "Textract",
            "Effect": "Allow",
            "Action": "textract:DetectDocumentText",
            "Resource": "*"
        }
    ]
}