variable "stg_deployment_tag" {
  description = "Deployment tag used for staging (e.g., 'latest')."
  type        = string
  default     = "latest"
}

module "stg_aws_prefect_backend" {
  count = var.stg_deployment_tag != "" ? 1 : 0

  source = "../../../../../devops/infra/terraform/modules/aws_prefect_backend"

  pipeline                               = local.pipeline
  short_env                              = "stg"
  prefect_task_deployment_repository_url = aws_ecr_repository.pipeline_repo.repository_url
  prefect_task_deployment_tag            = var.stg_deployment_tag

  aws_ecs_cluster      = "${local.prefix}-pipelines-stg"
  aws_ecs_service_name = local.pipeline

  prefect_api_url = local.prefect_api_url
  prefect_api_key = local.prefect_api_key

  prefect_work_pools = [
    {
      name            = "stg--${local.pipeline}:2048-4096"
      kind            = "ecs"
      task_definition = {
        cpu    = 2048
        memory = 4096
        worker_resources = {
          cpu    = 2048
          memory = 4096
        }        
      }
    }
  ]

  prefect_deployments = [
    {
      prefect_flow_name       = "stg--${local.pipeline}:root"
      prefect_deployment_name = "stg--${local.pipeline}"
      pipeline                = tostring(local.pipeline)
      work_pool = {
        name          = "stg--${local.pipeline}:2048-4096"
        job_variables = {}
        queue         = "default"
      }
    }
  ]
}

prefect_deployment

prefect_deployment

resource "prefect_flow" "pipeline_flow" {
  for_each = { for dep in var.prefect_deployments : dep.prefect_deployment_name => dep }
  name     = each.value.prefect_flow_name
  tags     = []
}

resource "prefect_deployment" "pipeline_deployment" {
  for_each = { for dep in var.prefect_deployments : dep.prefect_deployment_name => dep }

  name       = each.value.prefect_deployment_name
  flow_id    = prefect_flow.pipeline_flow[each.key].id
  entrypoint = "src/${replace(each.value.pipeline, "-", "_")}/main.py:root"

  paused     = false
  pull_steps = [
    {
      type      = "set_working_directory",
      directory = "/opt/"
    }
  ]

  job_variables = jsonencode(
    merge(
      lookup(each.value.work_pool, "job_variables", {}),
      {
        "image" : "${var.prefect_task_deployment_repository_url}:${var.prefect_task_deployment_tag}"
      }
    )
  )

  work_pool_name  = each.value.work_pool.name
  work_queue_name = lookup(each.value.work_pool, "queue", "default")
}