# SPDX-FileCopyrightText: 2021 ☭ Emery Hemingway
# SPDX-License-Identifier: Unlicense

import std/[os, strutils, tables]

import compiler/[ast, idents, renderer, lineinfos]

import ../../preserves, ../schemas

type TypeTable = OrderedTable[string, PNode]

proc add(parent, child: PNode): PNode {.discardable.} =
  parent.sons.add child
  parent

proc add(parent: PNode; children: varargs[PNode]): PNode {.discardable.} =
  parent.sons.add children
  parent

proc child(sn: SchemaNode): SchemaNode =
  assert(sn.nodes.len == 1)
  sn.nodes[0]

proc nn(kind: TNodeKind; children: varargs[PNode]): PNode =
  result = newNode(kind)
  result.sons.add(children)

proc nn(kind: TNodeKind; child: PNode): PNode =
  result = newNode(kind)
  result.sons.add(child)

proc ident(s: string): PNode =
  newIdentNode(PIdent(s: s), TLineInfo())

proc accQuote(n: PNode): Pnode =
  nkAccQuoted.newNode.add(n)

proc ident(sn: SchemaNode): PNode =
  var s: string
  case sn.kind
  of snkAlt:
    s = sn.altLabel.toLower.nimIdentNormalize
  of snkLiteral: s = $sn.value
  of snkRecord:
    s = $sn.nodes[0]
  of snkNamed:
    s = sn.name
  of snkDictionary, snkVariableTuple:
    s = "data"
  else:
    raiseAssert("no ident for " & $sn.kind & " " & $sn)
  ident(s)

proc typeIdent(sn: SchemaNode): PNode =
  case sn.kind
  of snkAtom:
    case sn.atom
    of akBool: ident"bool"
    of akFloat: ident"float32"
    of akDouble: ident"float64"
    of akInt: ident"BiggestInt"
    of akString: ident"string"
    of akBytes: nn(nkBracketExpr, ident"seq", ident"byte")
    of akSymbol: ident"string" # TODO: distinct string type for symbols?
  of snkNamed:
    sn.pattern.typeIdent
  of snkRef:
    var id = ident sn.refPath[sn.refPath.high]
    for i in countDown(sn.refPath.high.pred, 0):
      id = nn(nkDotExpr, ident(sn.refPath[i].toLowerAscii), id)
    id
  else:
    stderr.writeLine("no typeIdent for " & $sn.kind & " " & $sn)
    ident"Preserve"

proc toExport(n: sink PNode): PNode =
  nkPostFix.newNode.add(ident"*", n)

proc newEmpty(): PNode = newNode(nkEmpty)

proc isConst(scm: Schema; sn: SchemaNode): bool =
  case sn.kind
  of snkLiteral: result = true
  of snkRef:
    if sn.refPath.len == 1:
      result = isConst(scm, scm.definitions[sn.refPath[0]])
  else: discard

proc isSymbolEnum(sn: SchemaNode): bool =
  if sn.kind == snkOr:
    for bn in sn.nodes:
      if bn.altBranch.kind != snkLiteral or
          bn.altBranch.value.kind != pkSymbol:
        return false
    result = true

proc toEnumTy(sn: SchemaNode): PNode =
  result = nkEnumTy.newNode.add newEmpty()
  for bn in sn.nodes:
    result.add bn.altLabel.nimIdentNormalize.ident.accQuote

proc toEnumDef(name: string; sn: SchemaNode): PNode =
  nkTypeDef.newNode.add(
    nkPragmaExpr.newNode.add(
      name.ident.toExport,
      nkPragma.newNode.add(ident"pure")),
      newEmpty(),
      sn.toEnumTy)

proc nimTypeOf(scm: Schema; known: var TypeTable; sn: SchemaNode; name = ""): PNode =
  case sn.kind
  of snkOr:
    if sn.isSymbolEnum:
      result = sn.toEnumTy
    else:
      let
        enumName = name.nimIdentNormalize & "Kind"
        enumIdent = ident(enumName)
      if enumName notin known:
        known[enumName] = toEnumDef(enumName, sn)
      let recCase = nkRecCase.newNode.add(
        nkIdentDefs.newNode.add(
          "kind".ident.toExport,
          enumName.ident,
          newEmpty()))
      for bn in sn.nodes:
        assert(bn.kind == snkAlt, $bn.kind)
        var recList = nkRecList.newNode
        case bn.altBranch.kind
        of snkRecord:
          case bn.altBranch.nodes.len
          of 0, 1: discard
          of 2:
            if not isConst(scm, bn.altBranch.nodes[1]):
              let label = bn.ident
              recList.add nkIdentDefs.newNode.add(
                label.accQuote.toExport,
                nimTypeOf(scm, known, bn.altBranch.nodes[1], $label),
                newEmpty())
          else:
            for i, field in bn.altBranch.nodes:
              if i > 0 and (not isConst(scm, field)):
                let label = field.ident
                recList.add nkIdentDefs.newNode.add(
                  label.accQuote.toExport,
                  nimTypeOf(scm, known, field, $label),
                  newEmpty())
        else:
          if isConst(scm, bn.altBranch):
            recList.add nkDiscardStmt.newNode.add(newEmpty())
          else:
            let label = bn.ident
            recList.add(nkIdentDefs.newNode.add(
              label.accQuote,
              nimTypeOf(scm, known, bn.altBranch, $label),
              newEmpty()))
        let disc = nkDotExpr.newNode.add(
          enumIdent, bn.altLabel.nimIdentNormalize.ident.accQuote)
        recCase.add nkOfBranch.newNode.add(disc, recList)
      result = nn(nkObjectTy,
        newEmpty(),
        newEmpty(),
        nn(nkRecList, recCase))
  of snkAny:
    result = ident"Preserve"
  of snkAtom:
    result = typeIdent(sn)
  of snkEmbedded:
    result = nimTypeOf(scm, known, sn.embed)
  of snkLiteral:
    result = case sn.value.kind # nearly verbatim from ../../preserves/src/preserves.nim
    of pkBoolean: ident"bool"
    of pkFloat: ident"float32"
    of pkDouble: ident"float64"
    of pkSignedInteger: ident"BiggestInt"
    of pkBigInteger: ident"BigInt"
    of pkString: ident"string"
    of pkByteString: nn(
      nkBracketExpr, ident"seq", ident"byte")
    of pkSymbol: ident"string"
    of pkRecord: ident"Preserve"
    of pkSequence: nn(
      nkBracketExpr, ident"seq", ident"Preserve")
    of pkSet: nn(
      nkBracketExpr, ident"HashSet", ident"Preserve")
    of pkDictionary: nn(
      nkBracketExpr, ident"Table", ident"Preserve", ident"Preserve")
    of pkEmbedded: nn(nkDiscardStmt, newEmpty()) # TODO: is this possible?
  of snkSequenceOf:
    result = nkBracketExpr.newNode.add(
        ident"seq",
        nimTypeOf(scm, known, sn.child))
  of snkSetOf:
    result = nkBracketExpr.newNode.add(
        ident"HashedSet",
        nimTypeOf(scm, known, sn.child))
  of snkDictOf:
    result = nkBracketExpr.newNode.add(
        ident"Table",
        nimTypeOf(scm, known, sn.nodes[0]),
        nimTypeOf(scm, known, sn.nodes[1]))
  of snkRecord:
    case sn.nodes.len
    of 0, 1:
      result = nn(nkObjectTy,
        newEmpty(),
        newEmpty(),
        nn(nkRecList, nn(nkDiscardStmt, newEmpty())))
    else:
      let recList = nkRecList.newNode()
      for i, field in sn.nodes:
        if i > 0:
          let id = field.ident
          recList.add nkIdentDefs.newNode.add(
            id.accQuote.toExport,
            nimTypeOf(scm, known, field, $id),
            newEmpty())
      result = nn(nkObjectTy,
        newEmpty(),
        newEmpty(),
        recList)
  of snkTuple, snkVariableTuple:
    # TODO: the variable part
    result = nkTupleTy.newNode
    for tn in sn.nodes:
      result.add nkIdentDefs.newNode.add(
        tn.ident.accQuote, nimTypeOf(scm, known, tn), newEmpty())
  of snkDictionary:
    result = nkTupleTy.newNode
    for i in countup(0, sn.nodes.high, 2):
      let id = ident(sn.nodes[i+0])
      result.add nkIdentDefs.newNode.add(
        id.accQuote,
        nimTypeOf(scm, known, sn.nodes[i+1], $id),
        newEmpty())
  of snkNamed:
    result = nimTypeOf(scm, known, sn.pattern, name)
  of snkRef:
    result = typeIdent(sn)
  else:
    result = nkCommentStmt.newNode
    result.comment.add("Missing type generator for " & $sn.kind & " " & $sn)

proc toConst(name: string; def: SchemaNode): Pnode =
  case def.kind
  of snkLiteral:
    result = nkConstDef.newNode.add(
      name.ident.accQuote, newEmpty())
    case def.value.kind
    of pkSignedInteger:
      discard result.add newIntNode(nkIntLit, def.value.int)
    of pkSymbol:
      discard result.add nn(nkCall,
        ident"symbol",
        PNode(kind: nkStrLit, strVal: def.value.symbol),
        ident"EmbeddedType")
    else:
      raiseAssert("cannot convert " & $def & " to a Nim literal")
  else: discard

proc toNimLit(sn: SchemaNode): PNode =
  assert(sn.kind == snkLiteral, $sn)
  case sn.value.kind
  of pkSymbol:
    nkCall.newNode.add(
      ident"symbol",
      PNode(kind: nkStrLit, strVal: sn.value.symbol),
      ident"EmbeddedType")
  else:
    raiseAssert("no Nim literal for " & $sn)

proc preserveTypeOf(scm: Schema; known: var TypeTable; sn: SchemaNode; name = ""): PNode =
  case sn.kind
  of snkOr:
    if sn.isSymbolEnum:
      result = sn.toEnumTy
    else:
      let enumName = name.nimIdentNormalize & "Kind"
      if enumName notin known:
        known[enumName] = toEnumDef(enumName, sn)
      result = nkObjectTy.newNode.add(
        newEmpty(),
        newEmpty(),
        nkRecList.newNode.add(
          nkIdentDefs.newNode.add(
            ident"value", ident"Preserve", newEmpty()),
          nkIdentDefs.newNode.add(
            "schemaKind".ident,
            enumName.ident,
            newEmpty())))
  of snkAny:
    result = ident"Preserve"
  of snkAtom:
    result = case sn.atom
      of akBool: ident"bool"
      of akFloat: ident"float32"
      of akDouble: ident"float64"
      of akInt: ident"BiggestInt"
      of akString: ident"string"
      of akBytes: nkBracketExpr.newNode.add(
        ident"seq",
        ident"byte")
      of akSymbol: ident"string" # TODO: distinct string type for symbols?
  of snkRecord:
    case sn.nodes.len
    of 0, 1:
      result = nn(nkObjectTy,
        newEmpty(),
        newEmpty(),
        nn(nkDiscardStmt, newEmpty()))
    else:
      let recList = nkRecList.newNode()
      for i, field in sn.nodes:
        if i > 0:
          let id = field.ident
          recList.add nkIdentDefs.newNode.add(
            id.accQuote.toExport,
            nimTypeOf(scm, known, field, $id),
            newEmpty())
      result = nn(nkObjectTy,
        newEmpty(),
        newEmpty(),
        recList)
  of snkTuple:
    let recList = nkRecList.newNode()
    for i, field in sn.nodes:
      let id = field.ident
      recList.add nkIdentDefs.newNode.add(
        id.accQuote,
        nimTypeOf(scm, known, field, $id),
        newEmpty())
    result = nn(nkTupleTy,
      newEmpty(),
      newEmpty(),
      recList)
  else:
    result = nimTypeOf(scm, known, sn, name)

proc generateProcs(result: var seq[PNode]; name: string; sn: SchemaNode) =
  proc exportIdent(id: string): PNode = nn(nkPostFix, ident"*", ident(id))
  case sn.kind
  of snkRecord:
    var
      params = nn(nkFormalParams, ident"Preserve")
      initRecordCall = nn(nkCall,
        nn(nkBracketExpr, ident"initRecord", ident"EmbeddedType"),
        sn.nodes[0].toNimLit)
    for i, field in sn.nodes:
      if i > 0:
        let
          id = field.ident
        var fieldType  = field.typeIdent
        if fieldType.kind != nkIdent or fieldType.ident.s != "Preserve":
          fieldType =
            nn(nkInfix,
              ident"|",
              fieldType,
              ident"Preserve")
        params.add nn(nkIdentDefs,
          id,
          fieldType,
          newEmpty())
        initRecordCall.add(
          nn(nkCall, ident"toPreserve", id, ident"EmbeddedType"))
    var procId = name
    procId[0] = procId[0].toLowerAscii
    result.add nn(nkProcDef,
      exportIdent(procId),
      newEmpty(),
      newEmpty(),
      params,
      newEmpty(),
      newEmpty(),
      nn(nkStmtList,
        PNode(
          kind: nkCommentStmt,
          comment: "Preserves constructor for ``" & name & "``."),
        initRecordCall))
  else: discard

proc collectRefImports(imports: PNode; sn: SchemaNode) =
  case sn.kind
  of snkRef:
    if sn.refPath.len > 1:
      imports.add ident(sn.refPath[0].toLowerAscii)
  else:
    for child in sn.items:
      collectRefImports(imports, child)

proc collectRefImports(imports: PNode; scm: Schema) =
  if scm.embeddedType.contains {'.'}:
    let m = split(scm.embeddedType, '.', 1 )[0]
    imports.add ident(m.toLowerAscii)
  for _, def in scm.definitions:
    collectRefImports(imports, def)

proc moduleScopedIdent(s: string): PNode =
  var id: string
  let items = split(s, {'.'})
  for i in 0..<items.high:
    id.add items[i].toLowerAscii
    id.add '.'
  id.add items[items.high]
  ident(id)

proc generateNimFile*(scm: Schema; path: string) =
  var
    knownTypes: TypeTable
    typeSection = newNode nkTypeSection
    constSection = newNode nkLetSection
    procs: seq[PNode]
  if scm.embeddedType == "":
    typeSection.add nn(nkTypeDef,
      ident"EmbeddedType",
      newEmpty(),
      ident"void")
  else:
    typeSection.add nn(nkTypeDef,
      ident"EmbeddedType",
      newEmpty(),
      scm.embeddedType.moduleScopedIdent)
    typeSection.add nn(nkTypeDef,
      ident"Preserve",
      newEmpty(),
      nn(nkBracketExpr,
        ident"PreserveGen",
        ident"EmbeddedType"))
  for name, def in scm.definitions.pairs:
    if isConst(scm, def):
      constSection.add toConst(name, def)
    else:
      var name = name
      name[0] = name[0].toUpperAscii
      let t = nimTypeOf(scm, knownTypes, def, name)
      t.comment = $def
      case def.kind
      of snkAtom:
        knownTypes[name] = nkTypeDef.newNode.add(
          name.ident.toExport, newEmpty(), t)
      else:
        if def.kind == snkRecord:
          knownTypes[name] = nn(nkTypeDef,
            nn(nkPragmaExpr,
              name.ident.toExport,
              nn(nkPragma,
                nn(nkExprColonExpr,
                  ident"record",
                  PNode(kind: nkStrLit, strVal: $def.nodes[0])))),
              newEmpty(),
              t)
        else:
          case t.kind
          of nkEnumTy:
            knownTypes[name] = nn(nkTypeDef,
              nn(nkPragmaExpr,
                name.ident.toExport,
                nn(nkPragma, ident"pure")),
                newEmpty(),
                t)
          of nkDistinctTy:
            knownTypes[name] = nn(nkTypeDef,
              nn(nkPragmaExpr,
                name.ident.toExport,
                nn(nkPragma,
                  nn(nkExprColonExpr,
                    ident"borrow",
                    accQuote(ident".")))),
                newEmpty(),
                t)
          else:
            knownTypes[name] = nn(nkTypeDef,
              name.ident.toExport, newEmpty(), t)
      generateProcs(procs, name, def)
  for typeDef in knownTypes.values:
    typeSection.add typeDef
  var imports = nkImportStmt.newNode.add(
    ident"std/typetraits",
    ident"preserves")
  collectRefImports(imports, scm)
  var module = newNode(nkStmtList).add(
      imports,
      typeSection,
      constSection
    ).add(procs)
  writeFile(path, renderTree(module, {renderNone, renderIr}))

when isMainModule:
  import std/parseopt
  var inputs: seq[string]
  for kind, key, val in getopt():
    case kind
    of cmdLongOption:
      case key
      else: quit("unhandled option " & key)
    of cmdShortOption:
      case key
      else: quit("unhandled option " & key)
    of cmdArgument:
      inputs.add absolutePath(key)
    of cmdEnd: discard
  for filepath in inputs:
    let
      scm = parsePreservesSchema(readFile filepath, filepath)
      (dir, name, _) = splitFile(filepath)
      outputPath = dir / name & ".nim"
    generateNimFile(scm, outputPath)
    stdout.writeLine(outputPath)