# variables.tf
variable "prefix" {
  description = "Prefix for all resource names"
  type        = string
}

variable "databricks_account_id" {
  description = "Databricks account ID for Unity Catalog"
  type        = string
}

variable "aws_account_id" {
  description = "AWS account ID where resources will be created"
  type        = string
}

variable "storage_credential_name" {
  description = "Name for the Databricks storage credential"
  type        = string
  default     = "storage-credential"
}

variable "s3_bucket_name" {
  description = "S3 bucket name for external location"
  type        = string
}

variable "tags" {
  description = "Tags to apply to all resources"
  type        = map(string)
  default     = {}
}

variable "iam_propagation_delay" {
  description = "Time to wait for IAM changes to propagate"
  type        = string
  default     = "30s"
}

variable "credential_propagation_delay" {
  description = "Time to wait for credential creation to complete"
  type        = string
  default     = "20s"
}

# Local values for commonly used references
locals {
  databricks_master_role_arn = "arn:aws:iam::414351767826:role/unity-catalog-prod-UCMasterRole-14S5ZJVKOTYTL"
  unity_catalog_role_name    = "${var.prefix}-unity-catalog-access-role"
  unity_catalog_policy_name  = "${var.prefix}-unity-catalog-access-policy"
  external_location_name     = "${var.s3_bucket_name}-external-location"
}

# main.tf
# IAM role for Unity Catalog data access
resource "aws_iam_role" "unity_catalog_access_role" {
  name = local.unity_catalog_role_name

  assume_role_policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      # Allow Databricks master role to assume this role from the Databricks Customer AWS Account
      {
        Sid    = "AllowDatabricksMasterRole"
        Effect = "Allow"
        Principal = {
          AWS = local.databricks_master_role_arn
        }
        Action = "sts:AssumeRole"
        Condition = {
          StringEquals = {
            "sts:ExternalId" = var.databricks_account_id
          }
        }
      },
      # Allow self-assumption for role chaining
      {
        Sid    = "AllowSelfAssumption"
        Effect = "Allow"
        Principal = {
          AWS = "arn:aws:iam::${var.aws_account_id}:root"
        }
        Action = "sts:AssumeRole"
        Condition = {
          ArnEquals = {
            "aws:PrincipalArn" = "arn:aws:iam::${var.aws_account_id}:role/${local.unity_catalog_role_name}"
          }
        }
      }
    ]
  })

  tags = merge(var.tags, {
    Name        = "${var.prefix}-unity-catalog-iam-role"
    Environment = var.prefix
    Purpose     = "Unity Catalog Data Access"
  })
}

# IAM policy for S3 bucket access
resource "aws_iam_policy" "unity_catalog_s3_access_policy" {
  name        = local.unity_catalog_policy_name
  description = "Policy for Unity Catalog to access S3 bucket ${var.s3_bucket_name}"

  policy = jsonencode({
    Version = "2012-10-17"
    Statement = [
      {
        Sid    = "S3BucketAccess"
        Effect = "Allow"
        Action = [
          "s3:GetObject",
          "s3:PutObject",
          "s3:DeleteObject",
          "s3:ListBucket",
          "s3:GetBucketLocation",
          "s3:GetBucketPolicy",
          "s3:ListBucketMultipartUploads",
          "s3:GetObjectVersion"
        ]
        Resource = [
          "arn:aws:s3:::${var.s3_bucket_name}",
          "arn:aws:s3:::${var.s3_bucket_name}/*"
        ]
      },
      {
        Sid      = "AllowRoleAssumption"
        Effect   = "Allow"
        Action   = "sts:AssumeRole"
        Resource = aws_iam_role.unity_catalog_access_role.arn
      }
    ]
  })

  tags = merge(var.tags, {
    Name        = "${var.prefix}-unity-catalog-iam-policy"
    Environment = var.prefix
    Purpose     = "Unity Catalog S3 Access"
  })
}

# Attach the policy to the role
resource "aws_iam_role_policy_attachment" "unity_catalog_policy_attachment" {
  role       = aws_iam_role.unity_catalog_access_role.name
  policy_arn = aws_iam_policy.unity_catalog_s3_access_policy.arn
}

# Wait for IAM resources to propagate
resource "time_sleep" "wait_for_iam_propagation" {
  create_duration = var.iam_propagation_delay

  depends_on = [
    aws_iam_role_policy_attachment.unity_catalog_policy_attachment
  ]

  triggers = {
    role_arn   = aws_iam_role.unity_catalog_access_role.arn
    policy_arn = aws_iam_policy.unity_catalog_s3_access_policy.arn
  }
}

# Databricks storage credential
resource "databricks_storage_credential" "unity_catalog_credential" {
  name = "${var.prefix}-${var.storage_credential_name}"

  aws_iam_role {
    role_arn = aws_iam_role.unity_catalog_access_role.arn
  }

  depends_on = [
    time_sleep.wait_for_iam_propagation
  ]
}

# Wait for credential creation to complete
resource "time_sleep" "wait_for_credential_propagation" {
  create_duration = var.credential_propagation_delay

  depends_on = [
    databricks_storage_credential.unity_catalog_credential
  ]

  triggers = {
    credential_id = databricks_storage_credential.unity_catalog_credential.id
  }
}

# Databricks external location
resource "databricks_external_location" "s3_external_location" {
  name            = local.external_location_name
  url             = "s3://${var.s3_bucket_name}"
  credential_name = databricks_storage_credential.unity_catalog_credential.name

  comment = "External location for bucket ${var.s3_bucket_name}"

  depends_on = [
    time_sleep.wait_for_credential_propagation
  ]
}

# outputs.tf
output "unity_catalog_role_arn" {
  description = "ARN of the Unity Catalog IAM role"
  value       = aws_iam_role.unity_catalog_access_role.arn
}

output "unity_catalog_policy_arn" {
  description = "ARN of the Unity Catalog IAM policy"
  value       = aws_iam_policy.unity_catalog_s3_access_policy.arn
}

output "storage_credential_id" {
  description = "ID of the Databricks storage credential"
  value       = databricks_storage_credential.unity_catalog_credential.id
}

output "external_location_id" {
  description = "ID of the Databricks external location"
  value       = databricks_external_location.s3_external_location.id
}

output "external_location_url" {
  description = "URL of the external location"
  value       = databricks_external_location.s3_external_location.url
}