terraform {
  required_providers {
    coder = {
      source = "coder/coder"
    }
    docker = {
      source = "kreuzwerker/docker"
    }
  }
}

module "jetbrains_gateway" {
  source         = "registry.coder.com/modules/jetbrains-gateway/coder"
  version        = "1.0.13"
  agent_id       = coder_agent.main.id
  agent_name     = "main"
  folder         = "/home/coder/data"
  jetbrains_ides = ["PY"]
  default        = "PY"
}

module "filebrowser" {
  source   = "registry.coder.com/modules/filebrowser/coder"
  version  = "1.0.8"
  agent_id = coder_agent.main.id
  folder   = "/home/coder/data"
}

locals {
  jupyter-path     = data.coder_parameter.framework.value == "conda" ? "/home/coder/.conda/envs/DL/bin/jupyter" : "/home/coder/.local/bin/jupyter"
  jupyter-count    = (data.coder_parameter.framework.value == "conda" || data.coder_parameter.jupyter.value == "false") ? 0 : 1
  code-server-count = data.coder_parameter.code_server.value == "false" ? 0 : 1
}

data "coder_parameter" "ram" {
  name         = "ram"
  display_name = "RAM (GB)"
  description  = "Choose amount of RAM (min: 16 GB, max: 128 GB)"
  type         = "number"
  icon         = "https://raw.githubusercontent.com/matifali/logos/main/memory.svg"
  mutable      = true
  default      = "32"
  order        = 2
  validation {
    min = 16
    max = 128
  }
}

data "coder_parameter" "folder_loc" {
  name         = "folder_loc"
  display_name = "Folder name"
  description  = "Folder name for folder mount on host machine"
  type         = "string"
  icon         = "https://www.svgrepo.com/show/788/folder.svg"
  mutable      = false
  order        = 3
}


data "coder_parameter" "gpu_index" {
  name         = "gpu_index"
  display_name = "GPU Index"
  description  = "Select which GPU to use (0, 1, 2, etc.)"
  type         = "string"
  mutable      = false
  default      = "0"
  icon         = "https://www.svgrepo.com/show/474382/gpu.svg"
  order        = 1
  
  # Optional: Add validation to ensure the index is valid
  option {
    name        = "0️⃣"
    description = "Gpu Index 0"
    value       = "0"
  }
  option {
    name        = "1️⃣"
    description = "Gpu Index 1"
    value       = "1"
  }
  option {
    name        = "2️⃣"
    description = "Gpu Index 2"
    value       = "2"
  }
  option {
    name        = "3️⃣"
    description = "Gpu Index 3"
    value       = "3"
  }
  option {
    name        = "4️⃣"
    description = "Gpu Index 4"
    value       = "4"
  }
  option {
    name        = "5️⃣"
    description = "Gpu Index 5"
    value       = "5"
  }
  option {
    name        = "6️⃣"
    description = "Gpu Index 6"
    value       = "6"
  }
}

data "coder_parameter" "framework" {
  name         = "framework"
  display_name = "Deep Learning Framework"
  icon         = "https://raw.githubusercontent.com/matifali/logos/main/memory.svg"
  description  = "Choose your preffered framework"
  type         = "string"
  mutable      = false
  default      = "torch"
  order        = 4
  option {
    name        = "PyTorch"
    description = "PyTorch"
    value       = "torch"
    icon        = "https://raw.githubusercontent.com/matifali/logos/main/pytorch.svg"
  }
  option {
    name        = "Tensorflow"
    description = "Tensorflow"
    value       = "tf"
    icon        = "https://raw.githubusercontent.com/matifali/logos/main/tensorflow.svg"
  }
  option {
    name        = "Tensorflow + PyTorch"
    description = "Tensorflow + PyTorch"
    value       = "tf-torch"
    icon        = "https://raw.githubusercontent.com/matifali/logos/main/tf-torch.svg"
  }
  option {
    name        = "Tensorflow + PyTorch + conda"
    description = "Tensorflow + PyTorch + conda"
    value       = "tf-torch-conda"
    icon        = "https://raw.githubusercontent.com/matifali/logos/main/tf-torch-conda.svg"
  }
  option {
    name        = "Conda"
    description = "Only conda (install whatever you need)"
    value       = "conda"
    icon        = "https://raw.githubusercontent.com/matifali/logos/main/conda.svg"
  }
}

resource "coder_metadata" "workspace_info" {
  count       = data.coder_workspace.me.start_count
  resource_id = docker_image.deeplearning.id
  icon        = data.coder_parameter.framework.option[index(data.coder_parameter.framework.option.*.value, data.coder_parameter.framework.value)].icon
  item {
    key   = "Framework"
    value = data.coder_parameter.framework.option[index(data.coder_parameter.framework.option.*.value, data.coder_parameter.framework.value)].name
  }
  item {
    key   = "RAM (GB)"
    value = data.coder_parameter.ram.value
  }
}

data "coder_parameter" "jupyter" {
  name        = "Jupyter"
  icon        = "https://raw.githubusercontent.com/matifali/logos/main/jupyter.svg"
  description = "Do you want Jupyter Lab?"
  type        = "bool"
  mutable     = true
  default     = "false"
  order       = 3
}

data "coder_parameter" "code_server" {
    name        = "Code Server"
    description = "Do you want Code Server?"
    type        = "bool"
    default     = "true"
    mutable     = true
    icon        = "/icon/code.svg"
    order       = 4
}

provider "docker" {}

provider "coder" {}

data "coder_workspace" "me" {}

data "coder_workspace_owner" "me" {}

module "coder-login" {
  count    = data.coder_workspace.me.start_count
  source   = "registry.coder.com/modules/coder-login/coder"
  version  = "1.0.15"
  agent_id = coder_agent.main.id
}

module "code-server" {
  count         = local.code-server-count
  source        = "registry.coder.com/modules/code-server/coder"
  version       = "1.0.30"
  extensions    = ["ms-python.python", "ms-toolsai.jupyter", "dracula-theme.theme-dracula"]
  agent_id      = coder_agent.main.id
  share         = "public"
}

module "jupyterlab" {
  count    = local.jupyter-count
  source   = "registry.coder.com/modules/jupyterlab/coder"
  version  = "1.0.8"
  agent_id = coder_agent.main.id
}

resource "coder_agent" "main" {
  arch           = "amd64"
  os             = "linux"
  startup_script = <<EOT
    #!/bin/bash
    set -euo pipefail
    # Create user data directory
    mkdir -p ~/data
    # make user share directory
    mkdir -p ~/share
    sudo chown -R coder:coder /home/coder/data
    EOT

  metadata {
    display_name = "CPU Usage"
    interval     = 10
    order        = 1
    key          = "cpu_usage"
    script       = "coder stat cpu --host"
  }

  metadata {
    display_name = "RAM Usage"
    interval     = 10
    order        = 2
    key          = "ram_usage"
    script       = "coder stat mem --host"
  }

  metadata {
    display_name = "GPU Usage"
    interval     = 10
    order        = 3
    key          = "gpu_usage"
    script       = <<EOT
      nvidia-smi --query-gpu=utilization.gpu --format=csv,noheader,nounits | awk '{printf "%s%%", $1}'
    EOT
  }

  metadata {
    display_name = "GPU Memory Usage"
    interval     = 10
    order        = 4
    key          = "gpu_memory_usage"
    script       = <<EOT
      nvidia-smi --query-gpu=utilization.memory --format=csv,noheader,nounits | awk '{printf "%s%%", $1}'
    EOT
  }

  metadata {
    display_name = "Disk Usage"
    interval     = 600
    order        = 5
    key          = "disk_usage"
    script       = "coder stat disk $HOME"
  }

  metadata {
    display_name = "Word of the Day"
    interval     = 86400
    order        = 6
    key          = "word_of_the_day"
    script       = <<EOT
      #!/bin/bash
      curl -o - --silent https://www.merriam-webster.com/word-of-the-day 2>&1 | awk ' $0 ~ "Word of the Day: [A-z]+" { print $5; exit }'
    EOT
  }

}

locals {
  registry_name = "matifali/dockerdl"
}

data "docker_registry_image" "deeplearning" {
  name = "${local.registry_name}:${data.coder_parameter.framework.value}"
}

resource "docker_image" "deeplearning" {
  name          = "${local.registry_name}@${data.docker_registry_image.deeplearning.sha256_digest}"
  pull_triggers = [data.docker_registry_image.deeplearning.sha256_digest]
  keep_locally  = true
}

#Volumes Resources
#home_volume
resource "docker_volume" "home_volume" {
  name = "coder-${data.coder_workspace_owner.me.name}-${lower(data.coder_workspace.me.name)}-home"
}

#usr_volume
resource "docker_volume" "usr_volume" {
  name = "coder-${data.coder_workspace_owner.me.name}-${lower(data.coder_workspace.me.name)}-usr"
}

#etc_volume
resource "docker_volume" "etc_volume" {
  name = "coder-${data.coder_workspace_owner.me.name}-${lower(data.coder_workspace.me.name)}-etc"
}

#opt_volume
resource "docker_volume" "opt_volume" {
  name = "coder-${data.coder_workspace_owner.me.name}-${lower(data.coder_workspace.me.name)}-opt"
}

resource "docker_container" "workspace" {
  count    = data.coder_workspace.me.start_count
  image    = docker_image.deeplearning.image_id
  memory   = data.coder_parameter.ram.value * 1024
  gpus     = "device=${data.coder_parameter.gpu_index.value}" # GPU Index
  name     = "coder-${data.coder_workspace_owner.me.name}-${lower(data.coder_workspace.me.name)}"
  hostname = lower(data.coder_workspace.me.name)
  dns      = ["1.1.1.1"]
  command  = ["sh", "-c", replace(coder_agent.main.init_script, "127.0.0.1", "host.docker.internal")]
  env      = ["CODER_AGENT_TOKEN=${coder_agent.main.token}"]
  restart  = "unless-stopped"

  # devices {
  #   host_path = "/dev/nvidia0"
  # }
  # devices {
  #   host_path = "/dev/nvidiactl"
  # }
  # devices {
  #   host_path = "/dev/nvidia-uvm-tools"
  # }
  # devices {
  #   host_path = "/dev/nvidia-uvm"
  # }
  # devices {
  #   host_path = "/dev/nvidia-modeset"
  # }

  host {
    host = "host.docker.internal"
    ip   = "host-gateway"
  }

  ipc_mode = "host"

  # users home directory
  volumes {
    container_path = "/home/coder"
    volume_name    = docker_volume.home_volume.name
    read_only      = false
  }
  volumes {
    container_path = "/usr/"
    volume_name    = docker_volume.usr_volume.name
    read_only      = false
  }
  volumes {
    container_path = "/etc/"
    volume_name    = docker_volume.etc_volume.name
    read_only      = false
  }
  volumes {
    container_path = "/opt/"
    volume_name    = docker_volume.opt_volume.name
    read_only      = false
  }
  # users data directory
  volumes {
    container_path = "/home/coder/data/"
    host_path      = "/sda1/vuhuy/NTQ_Hub/coder/deployment/coder-data/user_data/${data.coder_parameter.folder_loc.value}/"
    read_only      = false
  }

  # Add labels in Docker to keep track of orphan resources.
  labels {
    label = "coder.owner"
    value = data.coder_workspace_owner.me.name
  }
  labels {
    label = "coder.owner_id"
    value = data.coder_workspace_owner.me.id
  }
  labels {
    label = "coder.workspace_id"
    value = data.coder_workspace.me.id
  }
  labels {
    label = "coder.workspace_name"
    value = data.coder_workspace.me.name
  }
}

resource "coder_metadata" "workspace" {
  count       = data.coder_workspace.me.start_count
  resource_id = docker_container.workspace[count.index].id
  daily_cost  = 50
  item {
    key   = "Framework"
    value = data.coder_parameter.framework.option[index(data.coder_parameter.framework.option.*.value, data.coder_parameter.framework.value)].name
  }
}