#import "@preview/cetz:0.2.2": canvas, draw
#import "utils.typ": get-participants-i

#let Y-SPACE = 10
#let PAR-PAD = (5pt, 3pt)
#let PAR-SPACE = 10


#let get-columns-width(participants, elements) = {
  let pars-i = get-participants-i(participants)
  let cells = ()
  for elmt in elements {
    if elmt.type == "seq" {
      let com = if elmt.comment == none {""} else {elmt.comment}
      let i1 = pars-i.at(elmt.p1)
      let i2 = pars-i.at(elmt.p2)
      cells.push(
        (
          elmt: elmt,
          i1: calc.min(i1, i2),
          i2: calc.max(i1, i2),
          cell: box(com, inset: 3pt)
        )
      )
    }
  }

  let widths = ()
  for i in range(participants.len() - 1) {
    let p1 = participants.at(i)
    let p2 = participants.at(i + 1)
    let w1 = measure(box(p1.display-name)).width + PAR-PAD.last() * 2
    let w2 = measure(box(p2.display-name)).width + PAR-PAD.last() * 2
    widths.push(w1 / 2pt + w2 / 2pt + PAR-SPACE)
  }

  for cell in cells.filter(c => c.i2 - c.i1 == 1) {
    let m = measure(cell.cell)
    widths.at(cell.i1) = calc.max(
      widths.at(cell.i1),
      m.width / 1pt
    )
  }

  let multicol-cells = cells.filter(c => c.i2 - c.i1 > 1)
  multicol-cells = multicol-cells.sorted(key: c => {
    c.i1 * 1000 + c.i2
  })
  for cell in multicol-cells {
    let m = measure(cell.cell)
    widths.at(cell.i2 - 1) = calc.max(
      widths.at(cell.i2 - 1),
      m.width / 1pt - widths.slice(0, cell.i2 - 1).sum()
    )
  }
  return widths
}

#let draw-group(x0, x1, y0, y1, group) = {
  let m = measure(box(group.name))
  let w = m.width / 1pt + 15
  let h = m.height / 1pt + 6
  draw.rect(
    (x0, y0),
    (x1, y1)
  )
  draw.merge-path(
    fill: gray.lighten(20%),
    close: true,
    {
      draw.line(
        (x0, y0),
        (x0 + w, y0),
        (x0 + w, y0 - h / 2),
        (x0 + w - 5, y0 - h),
        (x0, y0 - h)
      )
    }
  )
  draw.content(
    (x0, y0),
    group.name,
    anchor: "north-west",
    padding: (left: 5pt, right: 10pt, top: 3pt, bottom: 3pt)
  )

  if group.desc != none {
    draw.content(
      (x0 + w, y0),
      text([\[#group.desc\]], weight: "bold"),
      anchor: "north-west",
      padding: 3pt
    )
  }
}

#let render(participants, elements) = context canvas(length: 1pt, {
  let pars-i = get-participants-i(participants)

  let widths = get-columns-width(participants, elements)

  let x-pos = (0,)
  for width in widths {
    x-pos.push(x-pos.last() + width)
  }
  
  // Draw participants
  for (i, p) in participants.enumerate() {
    draw.content(
      (x-pos.at(i), 0),
      p.display-name,
      name: p.name,
      frame: "rect",
      padding: PAR-PAD,
      anchor: "south"
    )
  }

  let y = -Y-SPACE
  let groups = ()

  // Draw sequences
  for elmt in elements {
    if elmt.type == "seq" {
      let x1 = x-pos.at(pars-i.at(elmt.p1))
      let x2 = x-pos.at(pars-i.at(elmt.p2))
      let style = (
        mark: (end: "straight"),
        stroke: (
          dash: if elmt.dashed {"dashed"} else {"solid"},
          paint: elmt.color
        )
      )

      if elmt.comment != none {
        y -= measure(box(elmt.comment)).height / 1pt + 6
        draw.content(
          (calc.min(x1, x2), y),
          elmt.comment,
          anchor: "south-west",
          padding: 3pt
        )
      }

      draw.line(
        (x1, y),
        (x2, y),
        ..style
      )
      y -= Y-SPACE
    } else if elmt.type == "grp" {
      let m = measure(
        box(
          inset: (left: 5pt, right: 5pt, top: 3pt, bottom: 3pt),
        )
      )
      groups = groups.map(g => {
        if g.at(1).min-i == elmt.min-i { g.at(2) += 1 }
        if g.at(1).max-i == elmt.max-i { g.at(3) += 1 }
        g
      })
      groups.push((y, elmt, 0, 0))
      y -= m.height / 1pt + Y-SPACE
    
    } else if elmt.type == "grp-end" {
      let (start-y, group, start-lvl, end-lvl) = groups.pop()
      let x0 = x-pos.at(group.min-i) - start-lvl * 10 - 20
      let x1 = x-pos.at(group.max-i) + end-lvl * 10 + 20
      draw-group(x0, x1, start-y, y, group)

      y -= Y-SPACE
    }
  }

  // Draw vertical lines + end participants
  draw.on-layer(-1, {
    for (i, p) in participants.enumerate() {
      let x = x-pos.at(i)
      draw.line(
        (x, 0),
        (x, y),
        stroke: (dash: "dashed", paint: gray.darken(40%))
      )
      draw.content(
        (x, y),
        p.display-name,
        name: p.name,
        frame: "rect",
        padding: PAR-PAD,
        anchor: "north"
      )
    }
  })
})