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

locals {
  username = data.coder_workspace_owner.me.name
}

variable "docker_socket" {
  default     = ""
  description = "(Optional) Docker socket URI"
  type        = string
}

provider "docker" {
  host = var.docker_socket != "" ? var.docker_socket : null
}

data "coder_provisioner" "me" {}
data "coder_workspace" "me" {}
data "coder_workspace_owner" "me" {}

data "coder_parameter" "gpu" {
  name         = "gpu"
  display_name = "GPU Assignment"
  description  = "Which GPU(s) to attach. GPU 1 is shared with Ollama and Immich, so GPU 0 is the default to avoid contention."
  default      = "0"
  mutable      = true
  icon         = "/icon/memory.svg"

  option {
    name  = "GPU 0 only"
    value = "0"
  }
  option {
    name  = "GPU 1 only (shared with Ollama and Immich)"
    value = "1"
  }
  option {
    name  = "Both GPUs"
    value = "all"
  }
}

resource "coder_agent" "main" {
  arch             = data.coder_provisioner.me.arch
  os               = "linux"
  startup_script = <<-EOT
    set -e

    # 1. Allow user to use the Docker socket (Docker-out-of-Docker)
    # We set strict permissions on the socket file to ensure access
    if [ -e /var/run/docker.sock ]; then
      sudo chmod 666 /var/run/docker.sock
    fi

    # 2. Initialize Home Directory
    if [ ! -f ~/.init_done ]; then
      cp -rT /etc/skel ~
      touch ~/.init_done
    fi
  EOT

  env = {
    GIT_AUTHOR_NAME     = coalesce(data.coder_workspace_owner.me.full_name, data.coder_workspace_owner.me.name)
    GIT_AUTHOR_EMAIL    = "${data.coder_workspace_owner.me.email}"
    GIT_COMMITTER_NAME  = coalesce(data.coder_workspace_owner.me.full_name, data.coder_workspace_owner.me.name)
    GIT_COMMITTER_EMAIL = "${data.coder_workspace_owner.me.email}"
    OLLAMA_HOST         = "http://host.docker.internal:11434"
  }

  metadata {
    display_name = "CPU Usage"
    key          = "0_cpu_usage"
    script       = "coder stat cpu"
    interval     = 10
    timeout      = 1
  }

  metadata {
    display_name = "RAM Usage"
    key          = "1_ram_usage"
    script       = "coder stat mem"
    interval     = 10
    timeout      = 1
  }

  metadata {
    display_name = "Home Disk"
    key          = "3_home_disk"
    script       = "coder stat disk --path $${HOME}"
    interval     = 60
    timeout      = 1
  }

  metadata {
    display_name = "GPU Usage"
    key          = "8_gpu_usage"
    script       = "nvidia-smi --query-gpu=utilization.gpu --format=csv,noheader,nounits | awk '{sum+=$1} END {print sum/NR \"%\"}'"
    interval     = 10
    timeout      = 1
  }

  metadata {
    display_name = "CPU Usage (Host)"
    key          = "4_cpu_usage_host"
    script       = "coder stat cpu --host"
    interval     = 10
    timeout      = 1
  }

  metadata {
    display_name = "Memory Usage (Host)"
    key          = "5_mem_usage_host"
    script       = "coder stat mem --host"
    interval     = 10
    timeout      = 1
  }

  metadata {
    display_name = "Load Average (Host)"
    key          = "6_load_host"
    script   = <<EOT
      echo "`cat /proc/loadavg | awk '{ print $1 }'` `nproc`" | awk '{ printf "%0.2f", $1/$2 }'
    EOT
    interval = 60
    timeout  = 1
  }

  metadata {
    display_name = "Swap Usage (Host)"
    key          = "7_swap_host"
    script       = <<EOT
      free -b | awk '/^Swap/ { printf("%.1f/%.1f", $3/1024.0/1024.0/1024.0, $2/1024.0/1024.0/1024.0) }'
    EOT
    interval     = 10
    timeout      = 1
  }
}

module "code-server" {
  count    = data.coder_workspace.me.start_count
  source   = "registry.coder.com/coder/code-server/coder"
  version  = "~> 1.0"
  agent_id = coder_agent.main.id
  order    = 1
}

module "jetbrains" {
  count      = data.coder_workspace.me.start_count
  source     = "registry.coder.com/coder/jetbrains/coder"
  version    = "~> 1.1"
  agent_id   = coder_agent.main.id
  agent_name = "main"
  folder     = "/home/coder"
  tooltip    = "You need to [install JetBrains Toolbox](https://coder.com/docs/user-guides/workspace-access/jetbrains/toolbox) to use this app."
}

resource "docker_image" "workspace" {
  name = "coder-${data.coder_workspace.me.id}"
  build {
    context    = "."
    dockerfile = "Dockerfile"
  }
  keep_locally = true
}

resource "docker_container" "workspace" {
  count    = data.coder_workspace.me.start_count
  image    = docker_image.workspace.name
  name     = "coder-${data.coder_workspace_owner.me.name}-${lower(data.coder_workspace.me.name)}"
  hostname = data.coder_workspace.me.name
  
  entrypoint = [
    "sh", "-c",
    <<EOT
    sudo chown -R coder:coder /home/coder
    ${replace(coder_agent.main.init_script, "/localhost|127\\.0\\.0\\.1/", "host.docker.internal")}
    EOT
  ]
  
  gpus = data.coder_parameter.gpu.value == "all" ? "all" : "device=${data.coder_parameter.gpu.value}"

  env = [
    "CODER_AGENT_TOKEN=${coder_agent.main.token}",
    "NVIDIA_VISIBLE_DEVICES=${data.coder_parameter.gpu.value}",
    "NVIDIA_DRIVER_CAPABILITIES=all",
    "OLLAMA_HOST=http://host.docker.internal:11434"
  ]

  host {
    host = "host.docker.internal"
    ip   = "host-gateway"
  }
  
  # 1. Existing Home Directory Volume
  volumes {
    container_path = "/home/coder"
    host_path      = "/media/blob/ChaiNet/coder/workspaces/coder-${data.coder_workspace.me.id}-home"
    read_only      = false
  }

  # 2. NEW: Docker Socket Volume (Enables Docker-out-of-Docker)
  volumes {
    container_path = "/var/run/docker.sock"
    host_path      = "/var/run/docker.sock"
    read_only      = false
  }

  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
  }
}