#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 COMMENT-PAD = 8
#let LIFELINE-W = 10


#let get-columns-width(participants, elements) = {
  participants = participants.map(p => {
    p.insert("lifeline-lvl", 0)
    p.insert("max-lifelines", 0)
    p
  })
  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)
        )
      )

      if elmt.disable-src or elmt.destroy-src {
        let p = participants.at(i1)
        p.lifeline-lvl -= 1
        participants.at(i1) = p
      }
      if elmt.disable-dst {
        let p = participants.at(i2)
        p.lifeline-lvl -= 1
        participants.at(i2) = p
      }
      if elmt.enable-dst {
        let p = participants.at(i2)
        p.lifeline-lvl += 1
        p.max-lifelines = calc.max(p.max-lifelines, p.lifeline-lvl)
        participants.at(i2) = p
      }
    }
  }

  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 + COMMENT-PAD
    )
  }

  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()
    )
  }
  for (i, w) in widths.enumerate() {
    let p1 = participants.at(i)
    let p2 = participants.at(i + 1)
    let w = w + p1.max-lifelines * LIFELINE-W / 2
    if p2.max-lifelines != 0 {
      w += LIFELINE-W / 2
    }
    widths.at(i) = w
  }
  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 = ()
  let lifelines = participants.map(_ => (
    level: 0,
    lines: ()
  ))

  // Draw sequences
  for elmt in elements {
    if elmt.type == "seq" {
      let i1 = pars-i.at(elmt.p1)
      let i2 = pars-i.at(elmt.p2)

      if elmt.comment != none {
        y -= measure(box(elmt.comment)).height / 1pt + 6
      }

      if elmt.disable-src {
        let src-line = lifelines.at(i1)
        src-line.level -= 1
        src-line.lines.push(("disable", y, auto))
        lifelines.at(i1) = src-line
      }
      if elmt.destroy-src {
        let src-line = lifelines.at(i1)
        src-line.level -= 1
        src-line.lines.push(("destroy", y, auto))
        lifelines.at(i1) = src-line
      }

      let ll-lvl1 = lifelines.at(i1).level * LIFELINE-W / 2

      if elmt.disable-dst {
        let dst-line = lifelines.at(i2)
        dst-line.level -= 1
        dst-line.lines.push(("disable", y, auto))
        lifelines.at(i2) = dst-line
      }
      if elmt.destroy-dst {
        let dst-line = lifelines.at(i2)
        dst-line.level -= 1
        dst-line.lines.push(("destroy", y, auto))
        lifelines.at(i2) = dst-line
      }
      if elmt.enable-dst {
        let dst-line = lifelines.at(i2)
        dst-line.level += 1
        lifelines.at(i2) = dst-line
      }

      let x1 = x-pos.at(i1)
      let x2 = x-pos.at(i2)

      let ll-lvl2 = lifelines.at(i2).level * LIFELINE-W / 2

      let f = if elmt.flip {-1} else {1}
      if i1 <= i2 {
        x1 += ll-lvl1 * f
        x2 -= ll-lvl2 * f
      } else {
        x1 -= ll-lvl1 * f
        x2 += ll-lvl2 * f
      }

      let style = (
        mark: (end: "straight"),
        stroke: (
          dash: if elmt.dashed {"dashed"} else {"solid"},
          paint: elmt.color
        )
      )

      if elmt.p1 == elmt.p2 {
        let x3 = x1 - ll-lvl1 + ll-lvl2

        x2 = if elmt.flip {x1 - 20} else {x1 + 20}

        if elmt.comment != none {
          draw.content(
            (x1, y),
            elmt.comment,
            anchor: if elmt.flip {"south-east"} else {"south-west"},
            padding: 3pt
          )
        }

        draw.line(
          (x1, y),
          (x2, y),
          (x2, y - 10),
          (x3, y - 10),
          ..style
        )
        y -= 10

      } else {
        if elmt.comment != none {
          let x = calc.min(x1, x2)
          if x2 < x1 {
            x += COMMENT-PAD
          }
          draw.content(
            (x, y),
            elmt.comment,
            anchor: "south-west",
            padding: 3pt
          )
        }

        draw.line(
          (x1, y),
          (x2, y),
          ..style
        )
      }
      if elmt.enable-dst {
        let dst-line = lifelines.at(i2)
        dst-line.lines.push(("enable", y, auto))
        lifelines.at(i2) = dst-line
      }
      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

    } else if elmt.type == "sep" {
      let x0 = x-pos.first() - 20
      let x1 = x-pos.last() + 20
      let m = measure(
        box(
          elmt.name,
          inset: (left: 3pt, right: 3pt, top: 5pt, bottom: 5pt)
        )
      )
      let w = m.width / 1pt
      let h = m.height / 1pt
      let cx = (x0 + x1) / 2
      let xl = cx - w / 2
      let xr = cx + w / 2

      y -= h / 2
      draw.line((x0, y), (xl, y))
      draw.line((xr, y), (x1, y))
      y -= 3
      draw.line((x0, y), (xl, y))
      draw.line((xr, y), (x1, y))
      draw.content(
        ((x0 + x1) / 2, y + 1.5),
        elmt.name,
        anchor: "center",
        padding: (5pt, 3pt),
        frame: "rect"
      )
      y -= h / 2
      y -= Y-SPACE
    } else if elmt.type == "gap" {
      y -= elmt.size
    }
  }

  // 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%))
      )

      let rects = ()
      let destructions = ()
      let lines = ()
      for line in lifelines.at(i).lines {
        let event = line.first()
        if event == "enable" {
          lines.push(line)
        } else if event == "disable" or event == "destroy" {
          let l = lines.pop()
          let lvl = lines.len()
          rects.push((x + lvl * LIFELINE-W / 2, l.at(1), line.at(1)))
          if event == "destroy" {
            destructions.push((x + lvl * LIFELINE-W / 2, line.at(1)))
          }
        }
      }

      for rect in rects.rev() {
        let (cx, y0, y1) = rect
        draw.rect(
          (cx - LIFELINE-W / 2, y0),
          (cx + LIFELINE-W / 2, y1),
        )
      }
      for dest in destructions {
        let (cx, cy) = dest
        draw.line((cx - 8, cy - 8), (cx + 8, cy + 8), stroke: red + 2pt)
        draw.line((cx - 8, cy + 8), (cx + 8, cy - 8), stroke: red + 2pt)
      }

      draw.content(
        (x, y),
        p.display-name,
        name: p.name,
        frame: "rect",
        padding: PAR-PAD,
        anchor: "north"
      )
    }
  })
})