terraform {
  required_version = ">= 1.0"
  required_providers {
    aws = {
      source  = "hashicorp/aws"
      version = ">= 6.0"
    }
    archive = {
      source  = "hashicorp/archive"
      version = ">= 2.0"
    }
  }
}

provider "aws" {
  region = var.aws_region
}

data "archive_file" "lambda_zip" {
  type        = "zip"
  output_path = "${path.module}/lambda_function.zip"
  source {
    content  = file("${path.module}/lambda_function.py")
    filename = "lambda_function.py"
  }
}

resource "aws_iam_role" "lambda_execution_role" {
  name = "mirror-ball-lambda-execution-role"

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

resource "aws_iam_role_policy_attachment" "lambda_basic_execution" {
  policy_arn = "arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
  role       = aws_iam_role.lambda_execution_role.name
}

resource "aws_lambda_function" "mirror_ball_controller" {
  filename         = data.archive_file.lambda_zip.output_path
  function_name    = "mirror-ball-controller"
  role             = aws_iam_role.lambda_execution_role.arn
  handler          = "lambda_function.lambda_handler"
  runtime          = "python3.13"
  timeout          = 30
  source_code_hash = data.archive_file.lambda_zip.output_base64sha256

  environment {
    variables = {
      SWITCHBOT_TOKEN  = var.switchbot_token
      SWITCHBOT_SECRET = var.switchbot_secret
      DEVICE_ID        = var.device_id
    }
  }

  depends_on = [
    aws_iam_role_policy_attachment.lambda_basic_execution,
  ]
}

resource "aws_iam_user" "okta_workflow" {
  name = "okta-workflow-user"
}

resource "aws_iam_policy" "okta_workflow_lambda" {
  name        = "okta-workflow-lambda-policy"
  description = "Allows Okta Workflows Lambda connector to list and invoke the mirror ball controller"

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Sid      = "ListLambdaFunctions"
        Effect   = "Allow"
        Action   = "lambda:ListFunctions"
        Resource = "*"
      },
      {
        Sid      = "InvokeMirrorBallController"
        Effect   = "Allow"
        Action   = "lambda:InvokeFunction"
        Resource = aws_lambda_function.mirror_ball_controller.arn
      }
    ]
  })
}

resource "aws_iam_user_policy_attachment" "okta_workflow_lambda" {
  user       = aws_iam_user.okta_workflow.name
  policy_arn = aws_iam_policy.okta_workflow_lambda.arn
}

resource "aws_iam_access_key" "okta_workflow" {
  user = aws_iam_user.okta_workflow.name
}

variable "aws_region" {
  description = "AWS region"
  type        = string
  default     = "ap-northeast-1"
}

variable "switchbot_token" {
  description = "SwitchBot API token"
  type        = string
  sensitive   = true
}

variable "switchbot_secret" {
  description = "SwitchBot API secret"
  type        = string
  sensitive   = true
}

variable "device_id" {
  description = "SwitchBot device ID for the mirror ball"
  type        = string
}

output "lambda_function_name" {
  description = "Name of the mirror ball controller Lambda function"
  value       = aws_lambda_function.mirror_ball_controller.function_name
}

output "lambda_function_arn" {
  description = "ARN of the mirror ball controller Lambda function"
  value       = aws_lambda_function.mirror_ball_controller.arn
}

output "okta_workflow_access_key_id" {
  description = "Access key ID for the Okta Workflows IAM user"
  value       = aws_iam_access_key.okta_workflow.id
}

output "okta_workflow_secret_access_key" {
  description = "Secret access key for the Okta Workflows IAM user (sensitive)"
  value       = aws_iam_access_key.okta_workflow.secret
  sensitive   = true
}

import time
import hashlib
import hmac
import base64
import json
import os
import urllib.request
import urllib.error

def lambda_handler(event, context):
    token = os.environ.get("SWITCHBOT_TOKEN")
    secret = os.environ.get("SWITCHBOT_SECRET")
    device_id = os.environ.get("DEVICE_ID")

    # Support both direct invocation and API Gateway proxy integration
    if isinstance(event.get("body"), str):
        body = json.loads(event["body"])
    elif isinstance(event, dict):
        body = event
    else:
        body = {}

    action = body.get("action", "").lower()

    if action == "on":
        command = "turnOn"
    elif action == "off":
        command = "turnOff"
    else:
        return {
            "statusCode": 400,
            "body": json.dumps(
                {
                    "message": f'Invalid action: "{action}". Expected "on" or "off".',
                    "success": False,
                },
                ensure_ascii=False,
            ),
        }

    try:
        t = int(round(time.time() * 1000))
        nonce = ""
        string_to_sign = bytes(f"{token}{t}{nonce}", "utf-8")
        secret_bytes = bytes(secret, "utf-8")
        sign = base64.b64encode(
            hmac.new(secret_bytes, msg=string_to_sign, digestmod=hashlib.sha256).digest()
        ).decode("utf-8")

        headers = {
            "Authorization": token,
            "t": str(t),
            "sign": sign,
            "nonce": nonce,
            "Content-Type": "application/json",
        }

        payload = json.dumps({"command": command, "commandType": "command"}).encode("utf-8")
        url = f"https://api.switch-bot.com/v1.1/devices/{device_id}/commands"

        req = urllib.request.Request(url, data=payload, headers=headers, method="POST")
        with urllib.request.urlopen(req, timeout=10) as response:
            status_code = response.getcode()
            response_body = json.loads(response.read().decode("utf-8"))

        return {
            "statusCode": status_code,
            "body": json.dumps(
                {
                    "message": f'Command "{command}" sent successfully.',
                    "device_id": device_id,
                    "switchbot_response": response_body,
                    "success": True,
                },
                ensure_ascii=False,
            ),
        }

    except urllib.error.HTTPError as e:
        error_body = e.read().decode("utf-8")
        return {
            "statusCode": e.code,
            "body": json.dumps(
                {
                    "message": f"HTTP error occurred: {e.code}",
                    "error": error_body,
                    "success": False,
                },
                ensure_ascii=False,
            ),
        }

    except Exception as e:
        return {
            "statusCode": 500,
            "body": json.dumps(
                {
                    "message": "Unexpected error occurred.",
                    "error": str(e),
                    "success": False,
                },
                ensure_ascii=False,
            ),
        }