#import "@preview/cetz:0.2.2": canvas, draw
#import "utils.typ": get-participants-i, get-style
#import "group.typ"
#import "participant.typ"
#import "sequence.typ"
#import "separator.typ"
#import "consts.typ": *


#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 = ()

  // Compute max lifeline levels
  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
      }
    }
  }

  // Compute column widths
  // Compute minimum widths for participant names
  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)
  }

  // Compute minimum width for simple sequences (spanning 1 column)
  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
    )
  }

  // Compute remaining widths for longer sequences (spanning multiple columns)
  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()
    )
  }

  // Add lifeline widths
  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 render(participants, elements) = context canvas(length: 1pt, {
  let shapes = ()
  let pars-i = get-participants-i(participants)

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

  // Compute each column's X position
  let x-pos = (0,)
  for width in widths {
    x-pos.push(x-pos.last() + width)
  }

  let draw-seq = sequence.render.with(pars-i, x-pos, participants)
  let draw-group = group.render.with()
  let draw-sep = separator.render.with(x-pos)
  let draw-par = participant.render.with(x-pos)
  
  // Draw participants (start)
  for p in participants {
    if p.from-start {
      shapes += draw-par(p)
    }
  }

  let y = -Y-SPACE
  let groups = ()
  let lifelines = participants.map(_ => (
    level: 0,
    lines: ()
  ))

  // Draw elemnts
  for elmt in elements {
    // Sequences
    if elmt.type == "seq" {
      let shps
      (y, lifelines, shps) = draw-seq(elmt, y, lifelines)
      shapes += shps

    // Groups (start) -> reserve space for labels + store position
    } else if elmt.type == "grp" {
      let m = measure(
        box(
          elmt.name,
          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
    
    // Groups (end) -> actual drawing
    } 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
      shapes += draw-group(x0, x1, start-y, y, group)

      y -= Y-SPACE

    // Separator
    } else if elmt.type == "sep" {
      let shps
      (y, shps) = draw-sep(elmt, y)
      shapes += shps
    
    // Gap
    } else if elmt.type == "gap" {
      y -= elmt.size
    }
  }

  // Draw vertical lines + lifelines + end participants
  shapes += draw.on-layer(-1, {
    for (i, p) in participants.enumerate() {
      let x = x-pos.at(i)

      // Draw vertical line
      let last-y = 0

      let rects = ()
      let destructions = ()
      let lines = ()

      // Compute lifeline rectangles + destruction positions
      for line in lifelines.at(i).lines {
        let event = line.first()
        if event == "create" {
          last-y = line.at(1)

        } else if event == "enable" {
          if lines.len() == 0 {
            draw.line(
              (x, last-y),
              (x, line.at(1)),
              stroke: (dash: "dashed", paint: gray.darken(40%))
            )
          }
          lines.push(line)
        
        } else if event == "disable" or event == "destroy" {
          let lvl = 0
          if lines.len() != 0 {
            let l = lines.pop()
            lvl = lines.len()
            rects.push((
              x + lvl * LIFELINE-W / 2,
              l.at(1),
              line.at(1),
              l.at(2)
            ))
            last-y = line.at(1)
          }

          if event == "destroy" {
            destructions.push((x + lvl * LIFELINE-W / 2, line.at(1)))
          }
        }
      }

      draw.line(
        (x, last-y),
        (x, y),
        stroke: (dash: "dashed", paint: gray.darken(40%))
      )

      // Draw lifeline rectangles (reverse for bottom to top)
      for rect in rects.rev() {
        let (cx, y0, y1, style) = rect
        let style = get-style("lifeline", style)
        draw.rect(
          (cx - LIFELINE-W / 2, y0),
          (cx + LIFELINE-W / 2, y1),
          ..style
        )
      }

      // Draw lifeline destructions
      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 participants (end)
      draw.content(
        (x, y),
        p.display-name,
        name: p.name,
        frame: "rect",
        padding: PAR-PAD,
        anchor: "north"
      )
    }
  })

  shapes
})