# ---------------------------------------------------------------------------
# Public HTTPS endpoints that Fusion SaaS calls: POST /tokenize (T1) and
# POST /restore (T2), both on one API Gateway HTTP API. Fusion is shared SaaS, so
# it can only reach PUBLIC HTTPS with auth -- each handler checks a shared secret
# (tokenize_api_key, sent as x-api-key / bearer). These live in the "on-prem" zone
# (tags): they own the reversible token<->PII map, which must never leave it.
# (API Gateway, not a Lambda Function URL: this org's SCP blocks unauthenticated
# Function URLs.)
# ---------------------------------------------------------------------------

data "archive_file" "tokenize" {
  type        = "zip"
  source_dir  = "${path.module}/../gateway_api/tokenize"
  output_path = "${path.module}/tokenize_lambda.zip"
}

resource "aws_iam_role" "tokenize" {
  name = "${local.name}-tokenize-role"
  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect    = "Allow"
      Principal = { Service = "lambda.amazonaws.com" }
      Action    = "sts:AssumeRole"
    }]
  })
}

resource "aws_iam_role_policy" "tokenize" {
  name = "${local.name}-tokenize-policy"
  role = aws_iam_role.tokenize.id
  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        # Mint into the vault; scan customers for best-effort PERSON -> customer_id.
        Effect   = "Allow"
        Action   = ["dynamodb:PutItem"]
        Resource = [aws_dynamodb_table.vault.arn]
      },
      {
        Effect   = "Allow"
        Action   = ["dynamodb:Scan"]
        Resource = [aws_dynamodb_table.customers.arn]
      },
      {
        Effect   = "Allow"
        Action   = ["logs:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"]
        Resource = "arn:aws:logs:*:*:*"
      }
    ]
  })
}

resource "aws_lambda_function" "tokenize" {
  function_name    = "${local.name}-tokenize"
  role             = aws_iam_role.tokenize.arn
  runtime          = "python3.12"
  handler          = "handler.lambda_handler"
  filename         = data.archive_file.tokenize.output_path
  source_code_hash = data.archive_file.tokenize.output_base64sha256
  timeout          = 15
  environment {
    variables = {
      VAULT_TABLE       = aws_dynamodb_table.vault.name
      CUSTOMERS_TABLE   = aws_dynamodb_table.customers.name
      PRESIDIO_URL      = var.presidio_url
      TOKENIZE_API_KEY  = var.tokenize_api_key
      VAULT_TTL_SECONDS = "3600"
    }
  }
  # In-VPC so it can reach the now-private Presidio over 5001. DynamoDB is reached
  # via the gateway endpoint (see ecs.tf); no NAT needed.
  vpc_config {
    subnet_ids         = data.aws_subnets.default.ids
    security_group_ids = [aws_security_group.tokenize_lambda.id]
  }
  tags = local.onprem_tag
}

# VPC-attached Lambdas need ENI management permissions.
resource "aws_iam_role_policy_attachment" "tokenize_vpc" {
  role       = aws_iam_role.tokenize.name
  policy_arn = "arn:aws:iam::aws:policy/service-role/AWSLambdaVPCAccessExecutionRole"
}

# Public HTTPS front door. We use an API Gateway HTTP API rather than a Lambda
# Function URL because this org's SCP blocks unauthenticated (auth_type=NONE)
# Function URLs. The API is public; auth is the shared secret checked in-handler
# (Fusion sends it as x-api-key / bearer). T2 adds a POST /restore route here.
resource "aws_apigatewayv2_api" "gateway" {
  name          = "${local.name}-gateway"
  protocol_type = "HTTP"
  # The static UI (CloudFront origin) calls /demo cross-origin -> allow CORS.
  cors_configuration {
    allow_origins = ["*"]
    allow_methods = ["POST", "OPTIONS"]
    allow_headers = ["content-type", "x-api-key"]
  }
  tags = local.onprem_tag
}

resource "aws_apigatewayv2_integration" "tokenize" {
  api_id                 = aws_apigatewayv2_api.gateway.id
  integration_type       = "AWS_PROXY"
  integration_uri        = aws_lambda_function.tokenize.invoke_arn
  payload_format_version = "2.0"
}

# Granular detector is now INTERNAL at /tokenize-raw (x-api-key protected); the
# orchestrator calls it. Public POST /tokenize does the full round trip (below).
resource "aws_apigatewayv2_route" "tokenize" {
  api_id    = aws_apigatewayv2_api.gateway.id
  route_key = "POST /tokenize-raw"
  target    = "integrations/${aws_apigatewayv2_integration.tokenize.id}"
}

resource "aws_apigatewayv2_stage" "default" {
  api_id      = aws_apigatewayv2_api.gateway.id
  name        = "$default"
  auto_deploy = true
  # Access logs capture the RAW method + path AWS received -- for debugging what
  # Fusion actually sends (404 route-no-match never reaches a Lambda log).
  access_log_settings {
    destination_arn = aws_cloudwatch_log_group.apigw_access.arn
    format = jsonencode({
      requestId        = "$context.requestId"
      ip               = "$context.identity.sourceIp"
      method           = "$context.httpMethod"
      path             = "$context.path"
      routeKey         = "$context.routeKey"
      status           = "$context.status"
      protocol         = "$context.protocol"
      userAgent        = "$context.identity.userAgent"
      integrationError = "$context.integrationErrorMessage"
    })
  }
  tags = local.onprem_tag
}

resource "aws_cloudwatch_log_group" "apigw_access" {
  name              = "/apigw/${local.name}-gateway-access"
  retention_in_days = 7
  tags              = local.onprem_tag
}

resource "aws_lambda_permission" "tokenize_apigw" {
  statement_id  = "AllowApiGatewayInvoke"
  action        = "lambda:InvokeFunction"
  function_name = aws_lambda_function.tokenize.function_name
  principal     = "apigateway.amazonaws.com"
  source_arn    = "${aws_apigatewayv2_api.gateway.execution_arn}/*/*"
}

########################################
# T2: /restore -- egress re-identification (vault lookup by session_id)
########################################
data "archive_file" "restore" {
  type        = "zip"
  source_dir  = "${path.module}/../gateway_api/restore"
  output_path = "${path.module}/restore_lambda.zip"
}

resource "aws_iam_role" "restore" {
  name = "${local.name}-restore-role"
  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect    = "Allow"
      Principal = { Service = "lambda.amazonaws.com" }
      Action    = "sts:AssumeRole"
    }]
  })
}

resource "aws_iam_role_policy" "restore" {
  name = "${local.name}-restore-policy"
  role = aws_iam_role.restore.id
  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        # Read-only: look up this session's tokens to re-attach identity.
        Effect   = "Allow"
        Action   = ["dynamodb:Scan"]
        Resource = [aws_dynamodb_table.vault.arn]
      },
      {
        Effect   = "Allow"
        Action   = ["logs:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"]
        Resource = "arn:aws:logs:*:*:*"
      }
    ]
  })
}

resource "aws_lambda_function" "restore" {
  function_name    = "${local.name}-restore"
  role             = aws_iam_role.restore.arn
  runtime          = "python3.12"
  handler          = "handler.lambda_handler"
  filename         = data.archive_file.restore.output_path
  source_code_hash = data.archive_file.restore.output_base64sha256
  timeout          = 15
  environment {
    variables = {
      VAULT_TABLE      = aws_dynamodb_table.vault.name
      TOKENIZE_API_KEY = var.tokenize_api_key
    }
  }
  tags = local.onprem_tag
}

resource "aws_apigatewayv2_integration" "restore" {
  api_id                 = aws_apigatewayv2_api.gateway.id
  integration_type       = "AWS_PROXY"
  integration_uri        = aws_lambda_function.restore.invoke_arn
  payload_format_version = "2.0"
}

resource "aws_apigatewayv2_route" "restore" {
  api_id    = aws_apigatewayv2_api.gateway.id
  route_key = "POST /restore"
  target    = "integrations/${aws_apigatewayv2_integration.restore.id}"
}

resource "aws_lambda_permission" "restore_apigw" {
  statement_id  = "AllowApiGatewayInvoke"
  action        = "lambda:InvokeFunction"
  function_name = aws_lambda_function.restore.function_name
  principal     = "apigateway.amazonaws.com"
  source_arn    = "${aws_apigatewayv2_api.gateway.execution_arn}/*/*"
}

########################################
# T6: /demo -- Fusion-less orchestrator the static UI calls (tokenize->agent->restore)
########################################
data "archive_file" "orchestrator" {
  type        = "zip"
  source_dir  = "${path.module}/../gateway_api/orchestrator"
  output_path = "${path.module}/orchestrator_lambda.zip"
}

resource "aws_iam_role" "orchestrator" {
  name = "${local.name}-orchestrator-role"
  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [{
      Effect    = "Allow"
      Principal = { Service = "lambda.amazonaws.com" }
      Action    = "sts:AssumeRole"
    }]
  })
}

resource "aws_iam_role_policy" "orchestrator" {
  name = "${local.name}-orchestrator-policy"
  role = aws_iam_role.orchestrator.id
  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        # Invoke the AgentCore Runtime. Scoped to the configured runtime ARN when set.
        Effect   = "Allow"
        Action   = ["bedrock-agentcore:InvokeAgentRuntime"]
        Resource = var.agent_runtime_arn != "" ? [var.agent_runtime_arn, "${var.agent_runtime_arn}/*"] : ["*"]
      },
      {
        Effect   = "Allow"
        Action   = ["logs:CreateLogGroup", "logs:CreateLogStream", "logs:PutLogEvents"]
        Resource = "arn:aws:logs:*:*:*"
      }
    ]
  })
}

resource "aws_lambda_function" "orchestrator" {
  function_name    = "${local.name}-orchestrator"
  role             = aws_iam_role.orchestrator.arn
  runtime          = "python3.12"
  handler          = "handler.lambda_handler"
  filename         = data.archive_file.orchestrator.output_path
  source_code_hash = data.archive_file.orchestrator.output_base64sha256
  timeout          = 120
  environment {
    variables = {
      REGION            = var.region
      TOKENIZE_URL      = "${trimsuffix(aws_apigatewayv2_stage.default.invoke_url, "/")}/tokenize-raw"
      RESTORE_URL       = "${trimsuffix(aws_apigatewayv2_stage.default.invoke_url, "/")}/restore"
      TOKENIZE_API_KEY  = var.tokenize_api_key
      AGENT_RUNTIME_ARN = var.agent_runtime_arn
    }
  }
  tags = local.onprem_tag
}

resource "aws_apigatewayv2_integration" "orchestrator" {
  api_id                 = aws_apigatewayv2_api.gateway.id
  integration_type       = "AWS_PROXY"
  integration_uri        = aws_lambda_function.orchestrator.invoke_arn
  payload_format_version = "2.0"
}

# Full round trip (tokenize -> agent -> restore). Public entrypoint is /tokenize
# (what Fusion mirrors); /demo kept as an alias so the existing UI keeps working.
resource "aws_apigatewayv2_route" "orchestrator" {
  api_id    = aws_apigatewayv2_api.gateway.id
  route_key = "POST /demo"
  target    = "integrations/${aws_apigatewayv2_integration.orchestrator.id}"
}

resource "aws_apigatewayv2_route" "orchestrator_tokenize" {
  api_id     = aws_apigatewayv2_api.gateway.id
  route_key  = "POST /tokenize"
  target     = "integrations/${aws_apigatewayv2_integration.orchestrator.id}"
  depends_on = [aws_apigatewayv2_route.tokenize] # free up "POST /tokenize" first
}

resource "aws_lambda_permission" "orchestrator_apigw" {
  statement_id  = "AllowApiGatewayInvoke"
  action        = "lambda:InvokeFunction"
  function_name = aws_lambda_function.orchestrator.function_name
  principal     = "apigateway.amazonaws.com"
  source_arn    = "${aws_apigatewayv2_api.gateway.execution_arn}/*/*"
}

output "tokenize_url" { value = "${trimsuffix(aws_apigatewayv2_stage.default.invoke_url, "/")}/tokenize" }
output "restore_url" { value = "${trimsuffix(aws_apigatewayv2_stage.default.invoke_url, "/")}/restore" }
output "demo_url" { value = "${trimsuffix(aws_apigatewayv2_stage.default.invoke_url, "/")}/demo" }