Setting terrain in current cell

For testing purposes, I used a simple function to generate terrain as either true or false. Now, I want to configure the TileMapLayer for a terrain named ‘dirt’. Specifically, my goal is to set the terrain to dirt (with connected textures) when the state is false, and to dirt_walls when the state is true. How can I achieve this functionality? Is there a function to set the current terrain on a cell?

extends Node2D

@export var tileset : TileSet
@export var height = 3
@export var width = 3

var grid = []

func _ready() -> void:
	generate()
	draw()

func generate():
	for x in range(height):
		grid.append([])
		for y in range(width):
			grid[x].append(randi() % 2)

func draw():
	for y in range(height):
		print(grid[y])

func draw_tiles():
	var tile_map = TileMap.new()
	tile_map.tile_set = tileset

	for x in range(width):
		for y in range(height):
			if grid[x][y] == 1:
				tile_map.set_cell(0, Vector2i(x, y), 0)
	add_child(tile_map)

Yes I know, i could use else for place terrain on walls

If you look at the documentation of TileMapLayer you can scroll down to the methods table and there you also have a set_cell method which you can use. I havent tried it, but I am sure it will work the same way as the old TileMap node