#import "/src/cetz.typ": draw, styles

#import "/src/consts.typ": *
#import "/src/core/utils.typ": get-ctx, get-style, is-elmt, set-ctx

#let shapes = {
  let from-module(path) = {
    import path as mod
    return (mod.name: (
      get-size: mod.get-size,
      render: mod.render,
      default-style: mod.default-style
    ))
  }
  from-module("participant/default.typ")
  from-module("participant/actor.typ")
  from-module("participant/boundary.typ")
  from-module("participant/control.typ")
  from-module("participant/entity.typ")
  from-module("participant/database.typ")
  from-module("participant/collections.typ")
  from-module("participant/queue.typ")
  from-module("participant/custom.typ")
}

#let participant-default-style = (
  fill: rgb("#E2E2F0"),
  stroke: black + .5pt,
  from-start: true,
  show-bottom: true,
  show-top: true,
  shape: "participant",
  track: (
    dash: "dashed",
    paint: gray.darken(40%),
    thickness: .5pt
  )
)

#let resolve-style(ctx, par) = {
  let style = styles.resolve(
    ctx.style,
    merge: par.style,
    root: "participant",
    base: participant-default-style
  )
  let shape-style = shapes.at(style.shape, default: (:))
                          .at("default-style", default: (:))
  style = styles.resolve(
    ctx.style,
    merge: style,
    base: shape-style
  )
  return style
}

#let pre-resolve-styles() = get-ctx(ctx => {
  let idx = (:)
  let elements = ctx.setup.elements
  let participants = ctx.setup.participants
  for (i, par) in participants.enumerate() {
    par.insert("resolved-style", resolve-style(ctx, par))
    participants.at(i) = par
    idx.insert(par.name, i)
  }
  for (i, elmt) in elements.enumerate() {
    if type(elmt) == function {
      ctx = elmt(ctx).ctx
    } else if is-elmt(elmt) {
      if elmt.type == "par" {
        let style = resolve-style(ctx, elmt)
        elements.at(i).insert("resolved-style", style)
        let i = idx.at(elmt.name)
        participants.at(i).resolved-style = style
      }
    }
  }
  
  set-ctx(c => {
    c.setup.elements = elements
    c.setup.participants = participants
    return c
  })
})

#let get-size(par) = {
  if par.invisible {
    return (width: 0, height: 0)
  }
  let style = par.resolved-style
  let func = shapes.at(style.shape).get-size
  return func(par)
}

#let render(par, y: 0, bottom: false) = get-ctx(ctx => {
  let style = resolve-style(ctx, par)
  let func = shapes.at(style.shape).render
  let par = par
  par.resolved-style = style
  func(ctx.x-pos.at(par.i), y, par, bottom)
},)

#let render-lifelines() = get-ctx(ctx => {
  let participants = ctx.participants
  for p in participants.filter(p => not p.invisible) {
    let style = p.resolved-style
    let x = ctx.x-pos.at(p.i)

    // Draw vertical line
    let last-y = 0

    let rects = ()
    let destructions = ()
    let stack = ()

    // Compute lifeline rectangles + destruction positions
    for event in ctx.lifelines.at(p.i).events {
      if event.type == "create" {
        last-y = line.at(1)

      } else if event.type == "enable" {
        if stack.len() == 0 {
          draw.line(
            (x, last-y),
            (x, event.y),
            stroke: style.track
          )
        }
        stack.push(event)
      
      } else if event.type == "disable" or event.type == "destroy" {
        let lvl = 0
        if stack.len() != 0 {
          let e = stack.pop()
          lvl = stack.len()
          rects.push((
            x + lvl * LIFELINE-W / 2,
            e.y,
            event.y,
            e.style
          ))
          last-y = event.y
        }

        if event.type == "destroy" {
          destructions.push((x + lvl * LIFELINE-W / 2, event.y))
        }
      } else if event.type == "delay-start" {
        draw.line(
          (x, last-y),
          (x, event.y),
          stroke: style.track
        )
        last-y = event.y
      } else if event.type == "delay-end" {
        draw.line(
          (x, last-y),
          (x, event.y),
          stroke: event.stroke
        )
        last-y = event.y
      }
    }

    draw.line(
      (x, last-y),
      (x, ctx.y),
      stroke: style.track
    )

    // 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: COL-DESTRUCTION + 2pt)
      draw.line((cx - 8, cy + 8), (cx + 8, cy - 8), stroke: COL-DESTRUCTION + 2pt)
    }

    // Draw participants (end)
    if style.show-bottom {
      (p.draw)(p, y: ctx.y, bottom: true)
    }
  }
},)