#import "@preview/cetz:0.2.2": canvas, draw

#let class-counter = counter("uml-class")

#let class-colors = (
  rgb("#C02B22"),
  rgb("#D4721C"),
  rgb("D3BC2F"),
  rgb("60CB12"),
  rgb("#2ECC40"),
  rgb("#39CCA5"),
  rgb("#38B4DA"),
  rgb("#0074d9"),
  rgb("#450DC9"),
  rgb("#890DC9"),
)

#let visibility(vis, symbol) = (type: "visibility", visibility: vis, symbol: symbol)

#let public = visibility("public", "+")
#let private = visibility("private", "-")
#let protected = visibility("protected", "#")
#let package = visibility("package", "~")

#let class-family(family, prefix: none) = (
  type: "class-family",
  family: family,
  prefix: prefix
)
#let abstract = class-family("abstract")
#let interface = class-family("interface", prefix: "interface")
#let trait = class-family("trait")
#let enum = class-family("enum")
#let utility = class-family("utility", prefix: "utility")

#let staticity(static) = (type: "staticity", static: static)
#let static = staticity(true)
#let non-static = staticity(false)

#let class(name, ..args, grouped: none) = {
  let attributes = ()
  let methods = ()
  let visibility = none
  let family = none

  let args = args.pos()
  if args.len() != 0 {
    for arg in args {
      if type(arg) == dictionary {
        if arg.type == "visibility" {
          visibility = arg
        } else if arg.type == "class-family" {
          family = arg
        }
      } else if type(arg) == array {
        for elmt in arg {
          if type(elmt) == dictionary {
            if elmt.type == "visibility" {
              visibility = elmt
            } else if elmt.type == "attribute" {
              attributes.push(elmt)
            } else if elmt.type == "method" {
              methods.push(elmt)
            }
          }
        }
      }
    }
  }
  return ((
    type: "class",
    name: name,
    visibility: visibility,
    family: family,
    attributes: attributes,
    methods: methods,
    color: rgb(184, 84, 80),
    grouped: grouped
  ),)
}

#let attr(name, default: none, ..args) = {
  let visibility = none
  let staticity = none
  let val-type = none
  for arg in args.pos() {
    if type(arg) == dictionary {
      if arg.type == "visibility" {
        visibility = arg
      } else if arg.type == "type" {
        val-type = arg
      } else if arg.type == "staticity" {
        staticity = staticity
      }
    } else if type(arg) == str {
      val-type = arg
    }
  }
  return ((
    type: "attribute",
    name: name,
    visibility: visibility,
    staticity: staticity,
    val-type: val-type,
    default: default
  ),)
}

#let method(name, ..args) = {
  let visibility = none
  let staticity = none
  let parameters = none
  let return-type = none
  
  for arg in args.pos() {
    if type(arg) == dictionary {
      if "type" in arg {
        if arg.type == "visibility" {
          visibility = arg
        } else if arg.type == "staticity" {
          staticity = arg
        }
      } else {
        parameters = arg
      }
    } else if type(arg) == str {
      return-type = arg
    } else if type(arg) == array {
      parameters = arg
    }
  }
  return ((
    type: "method",
    name: name,
    visibility: visibility,
    staticity: staticity,
    parameters: parameters,
    return-type: return-type
  ),)
}

#let render-attr(attr, defaults, show-visibility: true) = {
  let visibility = if attr.visibility == none {
    defaults.attr-visibility
  } else {
    attr.visibility
  }
  let staticity = if attr.staticity == none {
    defaults.attr-staticity
  } else {
    attr.staticity
  }

  let name = attr.name
  if staticity.static {
    name = underline(name)
  }
  let txt = name
  if show-visibility {
    txt = visibility.symbol + txt
  }
  
  if attr.val-type != none {
    txt += ": " + attr.val-type
  }
  if attr.default != none {
    txt += " = " + str(attr.default)
  }
  return txt
}

#let render-method(method, defaults, show-visibility: true) = {
  let visibility = if method.visibility == none {
    defaults.method-visibility
  } else {
    method.visibility
  }
  let staticity = if method.staticity == none {
    defaults.method-staticity
  } else {
    method.staticity
  }
  let params = ()

  if type(method.parameters) == array {
    params = method.parameters
  } else if type(method.parameters) == dictionary {
    params = method.parameters
                   .pairs()
                   .map(((n, t)) => n + ": " + t)
  }

  let name = method.name
  if staticity.static {
    name = underline(name)
  }
  let txt = name
  if show-visibility {
    txt = visibility.symbol + txt
  }
  txt += "(" + params.join(", ") + ")"
  if method.return-type != none {
    txt += ": " + method.return-type
  }
  
  return txt
}

#let get-attributes-by-visibility(class, defaults) = {
  let by-vis = (:)
  for attr in class.attributes {
    let visibility = if attr.visibility == none {
      defaults.attr-visibility
    } else {
      attr.visibility
    }
    if not visibility.visibility in by-vis {
      by-vis.insert(visibility.visibility, ())
    }
    by-vis.at(visibility.visibility).push(attr)
  }
  return by-vis
}

#let get-methods-by-visibility(class, defaults) = {
  let by-vis = (:)
  for method in class.methods {
    let visibility = if method.visibility == none {
      defaults.method-visibility
    } else {
      method.visibility
    }
    if not visibility.visibility in by-vis {
      by-vis.insert(visibility.visibility, ())
    }
    by-vis.at(visibility.visibility).push(method)
  }
  return by-vis
}

#let render-class-attributes(class, defaults) = {
  let grouped = if class.grouped == none {
    defaults.class-grouped
  } else {
    class.grouped
  }
  let cells = ()

  if not grouped {
    for attr in class.attributes {
      cells.push(render-attr(attr, defaults, show-visibility: true))
    }
  } else {
    let by-vis = get-attributes-by-visibility(class, defaults)
    for (vis, attrs) in by-vis {
      cells.push(table.cell(vis + ":"))
      cells += attrs.map(attr => {
        let txt = render-attr(attr, defaults, show-visibility: false)
        pad(left: 1em, txt)
      })
    }
  }
  return cells
}

#let render-class-methods(class, defaults) = {
  let grouped = if class.grouped == none {
    defaults.class-grouped
  } else {
    class.grouped
  }
  let cells = ()

  if not grouped {
    for method in class.methods {
      cells.push(render-method(method, defaults, show-visibility: true))
    }
  } else {
    let by-vis = get-methods-by-visibility(class, defaults)
    for (vis, attrs) in by-vis {
      cells.push(table.cell(vis + ":"))
      cells += attrs.map(method => {
        let txt = render-method(method, defaults, show-visibility: false)
        pad(left: 1em, txt)
      })
    }
  }
  return cells
}

#let render-class(class, defaults) = context {
  let family = if class.family == none {
    defaults.class-family
  } else {
    class.family
  }
  
  let cells = ()
  cells += render-class-attributes(class, defaults)
  
  cells.push(table.hline())

  cells += render-class-methods(class, defaults)

  let name = [*#class.name*]
  let family = class.family
  if family == abstract {
    name = emph(name)
  }
  if family != none and family.prefix != none {
    name = [*«#family.prefix»*\ #name]
  }
  let class-i = class-counter.get().first()
  let col = class-colors.at(calc.rem(class-i, class-colors.len()))
  
  set table.hline(stroke: col)
  table(
    inset: (top: 0.5em, bottom: 0.5em, left: 0.5em, right: 0.5em),
    stroke: (x, y) => {
      let s = (left: col, right: col)
      if y == 0 {
        s.insert("top", col)
        s.insert("bottom", col)
      }
      return s
    },
    align: (x, y) => {
      if y == 0 {
        return center + horizon
      }
      return left + horizon
    },
    fill: (x, y) => if y == 0 {
      col.lighten(75%)
    } else {
      none
    },
    table.header(name),
    ..cells,
    table.hline()
  )
  class-counter.step()
}

#let diagram(..args) = {
  class-counter.update(0)
  
  let args = args.pos()

  let defaults = (
    class-family: none,
    class-grouped: false,
    attr-visibility: public,
    attr-staticity: non-static,
    method-visibility: public,
    method-staticity: non-static,
  )
  let elements = ()
  if args.len() != 0 {
    for arg in args {
      if type(arg) == dictionary {
        defaults += arg
      } else if type(arg) == array {
        elements = arg
      }
    }
  }

  let cells = ()
  for element in elements {
    if element.type == "class" {
      cells.push(render-class(element, defaults))
    }
  }
  box(
    stroke: black,
    inset: 1em,
    grid(
      columns: 2,
      column-gutter: 1em,
      align: center + horizon,
      row-gutter: 1em,
      ..cells
    )
  )
}

#let render-attr2(pos, id, attr, defaults, show-visibility: true) = {
  let visibility = if attr.visibility == none {
    defaults.attr-visibility
  } else {
    attr.visibility
  }
  let staticity = if attr.staticity == none {
    defaults.attr-staticity
  } else {
    attr.staticity
  }

  let name = attr.name
  if staticity.static {
    name = underline(name)
  }
  let txt = name
  if show-visibility {
    txt = visibility.symbol + txt
  }
  
  if attr.val-type != none {
    txt += ": " + attr.val-type
  }
  if attr.default != none {
    txt += " = " + str(attr.default)
  }
  
  draw.content(pos, txt, name: id)
}

#let render-method2(method, defaults, show-visibility: true) = {
  let visibility = if method.visibility == none {
    defaults.method-visibility
  } else {
    method.visibility
  }
  let staticity = if method.staticity == none {
    defaults.method-staticity
  } else {
    method.staticity
  }
  let params = ()

  if type(method.parameters) == array {
    params = method.parameters
  } else if type(method.parameters) == dictionary {
    params = method.parameters
                   .pairs()
                   .map(((n, t)) => n + ": " + t)
  }

  let name = method.name
  if staticity.static {
    name = underline(name)
  }
  let txt = name
  if show-visibility {
    txt = visibility.symbol + txt
  }
  txt += "(" + params.join(", ") + ")"
  if method.return-type != none {
    txt += ": " + method.return-type
  }
  
  draw.content((), txt)
}

#let render-class-attributes2(class, defaults) = {
  let grouped = if class.grouped == none {
    defaults.class-grouped
  } else {
    class.grouped
  }
  let cells = ()
  let prev = ()

  if not grouped {
    for (i, attr) in class.attributes.enumerate() {
      let id = "class-"+class.name+"-attr-"+str(i)
      let txt = render-attr(attr, defaults, show-visibility: true)
      draw.content((rel: (0, -0.4em), to: prev), name: id, anchor: "north-west", txt)
      prev = id + ".south-west"
    }
  } else {
    let by-vis = get-attributes-by-visibility(class, defaults)
    for (vis, attrs) in by-vis {
      cells.push(table.cell(vis + ":"))
      let vis-id = "class-"+class.name+"-attrs-"+vis
      draw.content((rel: (0, -0.4em), to: prev), name: vis-id, anchor: "north-west", vis + ":")
      prev = vis-id + ".south-west"
      
      for (i, attr) in attrs.enumerate() {
        let txt = render-attr(attr, defaults, show-visibility: false)
        txt = pad(left: 1em, txt)
        let id = vis-id + "-" + str(i)
        draw.content((rel: (0, -0.4em), to: prev), name: id, anchor: "north-west", txt)
        prev = id + ".south-west"
      }
    }
  }
}

#let render-class-methods2(class, defaults) = {
  
}

#let render-class2(class, defaults) = {
  let attr-group = render-class-attributes2(class, defaults)
  let method-group = render-class-methods2(class, defaults)

  let family = if class.family == none {
    defaults.class-family
  } else {
    class.family
  }
  let name = [*#class.name*]
  let family = class.family
  if family == abstract {
    name = emph(name)
  }
  if family != none and family.prefix != none {
    name = [*«#family.prefix»*\ #name]
  }
  let class-i = class-counter.get().first()
  let col = class-colors.at(calc.rem(class-i, class-colors.len()))

  draw.group(name: "class-"+class.name, {
    draw.group(name: "class-"+class.name+"-name", padding: 5pt, {
      draw.content((), name)
    })
    draw.on-layer(-1, {
      draw.rect("class-"+class.name+"-name.north-west", "class-"+class.name+"-name.south-east", fill: col.lighten(75%), stroke: col, layer: -1)
    })
    draw.group(name: "class-"+class.name+"-attrs", padding: 0.2em, {
      attr-group
    })
    method-group
  })
  draw.rect("class-"+class.name+".north-west", "class-"+class.name+".south-east", stroke: col)
}

#let render-class3(class, defaults) = {
  let id = "class-" + class.name
  return draw.content((rel: (3, 0), to: ()), render-class(class, defaults), name: id)
}


#let diagram2(..args) = context {
  class-counter.update(0)
  
  let args = args.pos()

  let defaults = (
    class-family: none,
    class-grouped: false,
    attr-visibility: public,
    attr-staticity: non-static,
    method-visibility: public,
    method-staticity: non-static,
  )
  let elements = ()
  if args.len() != 0 {
    for arg in args {
      if type(arg) == dictionary {
        defaults += arg
      } else if type(arg) == array {
        elements = arg
      }
    }
  }

  let cells = ()
  canvas({
    for element in elements {
      if element.type == "class" {
        render-class3(element, defaults)
      }
    }
    draw.line("class-Test", "class-Essai", mark: (end: "diamond", scale: 2))
  })
}