// Helpers for code block displaying

#import "@preview/based:0.2.0": base64

#let code_font_scale = 0.5em

#let cell_matcher(cell, cell_tag: none, cell_id: none) = {
  // Matching function to check if a cell has a specific tag
  if cell.cell_type != "code" {
    return false
  }
  let metadata = cell.metadata
  if cell_tag != none {
    return metadata.keys().contains("tags") and cell.metadata.tags.contains(cell_tag)
  } else if cell_id != none {
    return metadata.keys().contains("id") and metadata.id == cell_id
  } else {
    return false
  }
}


#let get_cell(notebook, cell_tag: none, cell_id: none, cell_number: none) = {
  if cell_tag == none and cell_id == none and cell_number == none {
    error("One of cell_tag, cell_id, or cell_number should be provided.")
  }
  let cells = notebook.cells
  if cell_tag != none {
    let matching_cell = cells.find(x => cell_matcher(x, cell_tag: cell_tag))
    return matching_cell
  } else if cell_id != none {
    let matching_cell = cells.find(x => cell_matcher(x, cell_id: cell_id))
    return matching_cell
  } else if cell_number != none {
    if cell_number < 0 or cell_number >= cells.len() {
      error("cell_number out of range.")
    }
    return cells.at(cell_number)
  }
}

#let code_cell(notebook, ..identifiers) = {
  // Extract the content of a cell and display it as a code block
  let matching_cell = get_cell(notebook, ..identifiers)

  let cell_content = matching_cell.source
  let single_line = cell_content.fold("", (acc, x) => acc + x)

  text(
    raw(
      single_line,
      lang: "python",
      block: true
    ),
    size: code_font_scale
  )
}


// #let code_cell(notebook, cell_tag) = {
//   // Extract the content of a cell and display it as a code block
//   let cells = notebook.cells
//   let matching_cell = cells.find(x => cell_matcher(x, cell_tag))

//   let cell_content = matching_cell.source
//   let single_line = cell_content.fold("", (acc, x) => acc + x)

//   text(
//     raw(
//       single_line,
//       lang: "python",
//       block: true
//     ),
//     size: code_font_scale
//   )
// }

#let image_cell(notebook, ..identifiers) = {
  let matching_cell = get_cell(notebook, ..identifiers)
  let outputs = matching_cell.outputs
  for output in outputs {
    let image_data = output.at("data", default: (:))

    let png_data = image_data.at("image/png", default: none)
    if png_data != none {
      align(
        center,
        image.decode(
          base64.decode(png_data),
          // height: 70% // the height should be set by the caller. This gives the flexibility to adjust the height of the image
        )
      )
    }

    let svg_data = image_data.at("image/svg+xml", default: none)
    if svg_data != none {
      let svg_string = svg_data.join("")
      align(
        center,
        image(bytes(svg_string)),
      )
    }
  }
}


// #let image_cell(notebook, cell_tag) = {
//   // Extract the output (image) of a cell and display it as an image
//   let cells = notebook.cells
//   let matching_cell = cells.find(x => cell_matcher(x, cell_tag))

//   let outputs = matching_cell.outputs
//   for output in outputs {
//     let image_data = output.at("data", default: (:)).at("image/png", default: none)
//     if image_data != none {
//       align(
//         center,
//         image.decode(
//           base64.decode(image_data),
//           // height: 70% // the height should be set by the caller. This gives the flexibility to adjust the height of the image
//         )
//       )
//     }
//   }
// }


#let code_reference_cell(notebook, ..identifiers) = {
  // Extract the output (text) of a cell and display it as a code block
  // This is useful for showing the code of imported functions
  let matching_cell = get_cell(notebook, ..identifiers)
  let outputs = matching_cell.outputs
  for output in outputs {
    let cell_output = output.at("text", default: (:))
    if cell_output != none {
      let single_line = cell_output.join("")
      text(
        raw(
          single_line,
          lang: "python",
          block: true
        ),
        size: code_font_scale
      )
    }
  }
}


// #let code_reference_cell(notebook, cell_tag) = {
//   // Extract the output (text) of a cell and display it as a code block
//   // This is useful for showing the code of imported functions
//   let cells = notebook.cells
//   let matching_cell = cells.find(x => cell_matcher(x, cell_tag))

//   let outputs = matching_cell.outputs
//   for output in outputs {
//     let cell_output = output.at("text", default: (:))
//     if cell_output != none {
//       let single_line = cell_output.join("")
//       text(
//         raw(
//           single_line,
//           lang: "python",
//           block: true
//         ),
//         size: code_font_scale
//       )
//     }
//   }
// }