import std/[strutils, strformat, sequtils]

const input = readFile("day_07/input.txt")

type
  FileKind = enum
    fkFile,
    fkDirectory
  File = ref object
    name: string
    parent: File
    case kind: FileKind
    of fkFile:
      size: uint64
    of fkDirectory:
      children: seq[File]
  
proc `$`(file: FIle): string =
  result.add fmt"File(name: {file.name}, kind: {$file.kind}, "
  if file.kind == fkDirectory:
    let childrenAsStr = file.children.mapIt($it).join(",")
    result.add fmt"children: [{childrenAsStr}]"
  elif file.kind == fkFile:
    result.add fmt"size: {file.size}"
  result.add ")"

proc root(file: File): File =
  if file.parent.isNil:
    return file
  return file.parent.root

proc calcSize(file: File): uint64 =
  assert file.kind == fkDirectory
  return file.children.mapIt(if it.kind == fkFIle: it.size else: it.calcSize).foldl(a + b, 0.uint64)

proc processCommand(wd: File, command: string, output: string): File =
  case command.split(" ")[0]:
  of "cd":
    let path = command.split(" ")[1]
    case path:
    of "/":
      return wd.root
    of "..":
      return wd.parent
    else:
      for child in wd.children:
        if child.name == path:
          return child
      let newFile = File(name: path, parent: wd, kind: fkDirectory)
      wd.children.add newFile
      return newFile
  of "ls":
    for file in output.splitLines:
      if file.isEmptyOrWhitespace: break
      let desc = file.split(" ")[0]
      let name = file.split(" ")[1]
      if desc == "dir":
        wd.children.add(File(name: name, parent: wd, kind: fkDirectory))
      else:
        wd.children.add(File(name: name, parent: wd, kind: fkFile, size: desc.parseBiggestUInt))
  return wd

proc traverseDirectoryLog(data: string): File =
  var wd = File(name: "/", parent: nil, kind: fkDirectory, children: newSeq[File]())
  var command: string
  var output: string
  for line in data.splitLines:
    if line.startsWith "$":
      if not command.isEmptyOrWhitespace:
        wd = processCommand(wd, command, output)
      command = line[2..line.len-1]
      output = ""
      continue
    output.add line & "\n"
  return processCommand(wd, command, output).root

iterator traverse(dir: File): File {.closure.} =
  assert dir.kind == fkDirectory
  for child in dir.children:
    if child.kind != fkDirectory: continue
    yield child
    for subchild in child.traverse:
      yield subchild

proc solvePart1() =
  let root = traverseDirectoryLog(input)
  var sum: uint64 = 0
  for dir in traverse(root):
    if (dir.calcSize <= 100000):
      sum += dir.calcSize
  echo "Sum of dirs < 100000: " & $sum

proc solvePart2() =
  let root = traverseDirectoryLog(input)
  var minimumSize = 30000000 - (70000000 - root.calcSize)
  var smallestDir = root.calcSize
  for dir in traverse(root):
    let size = dir.calcSize
    if size < minimumSize: continue
    smallestDir = min(smallestDir, size)
  echo "Smalles dir for update: " & $smallestDir

if isMainModule:
  solvePart1()
  solvePart2()