import strformat
import strutils
import sequtils
import sugar
import tables
import options

import scanner
import chunk
import value
import config

type
  Local = object
    name: string # name of this local
    index: int # what is its index in the stack (0 is the stack bottom - nil in the main function)
    depth: int # depth of this local
    # if depth is -1, the variable cannot be referenced yet
    # its depth will be set once its first ever value is determined
    owning: bool # whether this reference owns the local on the stack
    # always true except when it's the "result" value
  
  Scope = ref object
    labels: seq[string]
    depth: int
    goalStackIndex: int # the stack count it started with plus 1
    jumps: seq[int] # jumps to be patched that jump to the end

  Compiler = ref object
    #
    # input
    #
    scanner: Scanner
    source: string
    #
    # state
    #
    previous: Token
    current: Token
    canAssign: bool
    locals: seq[Local]
    scopes: seq[Scope]
    stackIndex: int # how large the stack is
    # when there's an error both are set
    # panic mode can be turned off e.g. at block boundaries
    panicMode: bool
    #
    # output
    #
    chunk*: Chunk
    hadError*: bool
  
  Precedence = enum
    pcNone, pcAssignment, pcOr, pcAnd, pcEquality, pcComparison,
    pcTerm, pcFactor, pcUnary, pcCall, pcPrimary
    # pcUnary applies to all prefix operators regardless of this enum's value
    # changing pcUnary's position can change the priority of all unary ops
    # 
    # Note: unary only rules should have precedence pcNone!!!
  
  ParseRule = object
    name: string # debug purposes only
    prefix: (Compiler) -> void
    infix: (Compiler) -> void
    prec: Precedence # only relevant to infix, prefix always has pcUnary

proc newScope(comp: Compiler): Scope =
  result.new()
  result.depth = comp.scopes.len + 1
  result.goalStackIndex = comp.stackIndex + 1
  comp.scopes.add(result)

# HELPERS FOR THE COMPILER TYPE

proc newCompiler*(source: string): Compiler =
  result = new(Compiler)
  result.source = source
  result.hadError = false
  result.panicMode = false
  result.canAssign = true
  result.locals = @[]
  result.scopes = @[]

proc errorAt(comp: Compiler, line: int, msg: string, at: string = "") =
  if comp.panicMode:
    return
  write stderr, &"[line {line}] Error "
  if at.len > 0:
    write stderr, &"at {at} "
  write stderr, msg
  write stderr, "\n"
  comp.hadError = true
  comp.panicMode = true

proc error(comp: Compiler, msg: string) =
  ## create a simple error message
  comp.errorAt(comp.previous.line, msg)

proc errorAtCurrent(comp: Compiler, msg: string) =
  comp.errorAt(comp.current.line, msg)

proc advance(comp: Compiler) =
  comp.previous = comp.current
  while true:
    comp.current = comp.scanner.scanToken()
    when debugScanner:
      comp.current.debugPrint()
    if (comp.current.tokenType != tkError):
      break
    comp.errorAtCurrent(comp.current.text)

proc match(comp: Compiler, tokenType: TokenType): bool =
  if comp.current.tokenType == tokenType:
    comp.advance()
    true
  else:
    false

proc consume(comp: Compiler, tokenType: TokenType, msg: string) =
  if comp.current.tokenType == tokenType:
    comp.advance()
  else:
    comp.errorAtCurrent(msg)

proc synchronize(comp: Compiler) =
  comp.panicMode = false
  while comp.current.tokenType != tkEof:
    if comp.previous.tokenType in {tkSemicolon, tkRightBrace}:
      return
    if comp.current.tokenType in {tkFun, tkVar, tkFor, tkIf, tkWhile}:
      return
    comp.advance()

proc writeChunk(comp: Compiler, dStackIndex: int, ch: OpCode | Triple) =
  comp.stackIndex += dStackIndex
  when debugCompiler:
    debugEcho &"new stackindex: {comp.stackIndex}, delta: {dStackIndex} due to {ch.repr}"
  comp.chunk.writeChunk(ch, comp.previous.line)

proc writeConstant(comp: Compiler, constant: KonValue) =
  comp.stackIndex.inc
  let index = comp.chunk.writeConstant(constant, comp.previous.line)
  if index >= tripleMax:
    comp.error("Too many constants in one chunk.")


proc addLocal(comp: Compiler, name: string, delta: int, owning: bool = true) =
  if comp.locals.len >= tripleMax:
    comp.error("Too many local variables in function.")

  # if delta is 0 or negative, it means that it is already on the stack when addLocal is called  
  # if delta is positive, the first ever value of the local is to the right
  comp.locals.add(Local(name: name, depth: if delta > 0: -1 else: comp.scopes.high, index: comp.stackIndex + delta, owning: owning))

proc markInitialized(comp: Compiler) =
  comp.locals[comp.locals.high].depth = comp.scopes.high

# PARSE RULE/PRECEDENCE MISC
proc nop(comp: Compiler) = discard

var rules: Table[TokenType, ParseRule]
template genRule(ttype: TokenType, tprefix: (Compiler) -> void, tinfix: (Compiler) -> void, tprec: Precedence) =
  if tprec == pcUnary:
    raise newException(Exception, "pcUnary cannot be used as a rule precedence! Use pcNone for unary-only rules!")
  elif tprec == pcPrimary:
    raise newException(Exception, "Invalid rule: pcPrimary cannot be used for binary operators, if this rule is for a primary value, use pcNone!")
  elif tprec == pcNone and tinfix != nop:
    raise newException(Exception, "Invalid rule: pcNone only allowed for unary operators and primary values, not for infix ones!")
  rules[ttype] = ParseRule(name: $ttype, prefix: tprefix, infix: tinfix, prec: tprec)

proc getRule(opType: TokenType): Option[ParseRule] =
  if rules.hasKey(opType):
    some(rules[opType])
  else:
    none(ParseRule)

proc applyRule(rule: ParseRule): Precedence =
  # returns the rule's precedence
  rule.prec

proc increment(prec: Precedence): Precedence =
  # increases precedence by one
  if prec == pcPrimary:
    raise newException(Exception, "Invalid ruletable, pcPrimary precedence increment attempted.")
  else:
    Precedence(int(prec) + 1)

proc `<=`(a, b: Precedence): bool =
  int(a) <= int(b)

# JUMP HELPERS

proc emitJump(comp: Compiler, delta: int, op: OpCode): int =
  # delta -> 0 if the jump does not pop
  # delta -> -1 if the jump pops the condition from the stack
  comp.writeChunk(delta, op)
  comp.writeChunk(0, 0xffffff.toTriple)
  comp.chunk.len - 3

proc patchJump(comp: Compiler, offset: int) =
  let jump = (comp.chunk.len - offset - 3)

  if (jump > tripleMax):
    comp.error("Too much code to jump over.")
  
  let jumpt = jump.toTriple

  comp.chunk.code[offset] = jumpt[0]
  comp.chunk.code[offset + 1] = jumpt[1]
  comp.chunk.code[offset + 2] = jumpt[2]

proc emitLoop(comp: Compiler, loopstart: int, delta: int, op: OpCode) =
  comp.writeChunk(delta, op)

  let offset = comp.chunk.len - loopstart + 3
  if offset > tripleMax:
    comp.error("Loop body too large.")
  
  comp.writeChunk(0, offset.toTriple)

# SCOPE HELPERS
proc beginScope(comp: Compiler, consumeLabels: bool = true) =
  # owning false because it is the return value
  let scope = comp.newScope()

  while (consumeLabels and comp.match(tkLabel)):
    let label = comp.previous.text[1..^1]
    scope.labels.add(label)

  comp.writeChunk(1, opNil)
  comp.addLocal("^", delta = 0, owning = false)

proc restore(comp: Compiler, scope: Scope) =
  let delta = comp.stackIndex - scope.goalStackIndex
  for i in countup(1, delta):
    comp.writeChunk(0, opPop)
  comp.stackIndex = scope.goalStackIndex

proc jumpToEnd(comp: Compiler, scope: Scope) =
  let delta = comp.stackIndex - scope.goalStackIndex
  for i in countup(1, delta):
    comp.writeChunk(0, opPop)
  let jmp = comp.emitJump(0, opJump)
  scope.jumps.add(jmp)
  
proc endScope(comp: Compiler) =
  # remove locals
  let popped = comp.scopes.pop()
  comp.restore(popped)
  # patch jumps to after the scope (such jumps from breaks emit the pops before jumping)
  for jump in popped.jumps:
    comp.patchJump(jump)

# EXPRESSIONS
proc parsePrecedence(comp: Compiler, prec: Precedence) =
  comp.advance()
  let rule = comp.previous.tokenType.getRule

  if rule.isSome and rule.get.prefix != nop:
    comp.canAssign = prec <= pcAssignment

    let rule = rule.get
    when debugCompiler:
      debugEcho &"parsePrecedence call, valid prefix op found, rule used: {rule.name}, precedence: {prec}"

    rule.prefix(comp)

    while comp.current.tokenType.getRule.isSome and
      prec <= comp.current.tokenType.getRule.get.prec:
      comp.advance()
      # checked for isSome in the loop
      # since advance moves current to previous
      if comp.previous.tokenType.getRule.get.infix == nop:
        # should never happen, as having a precedence set
        # means that it is a binary op
        comp.error("Invalid rule table.")
        return
      else:
        let infixRule = comp.previous.tokenType.getRule.get.infix
        infixRule(comp)
  else:
    comp.error("Expect expression.")

proc expression(comp: Compiler) =
  ## The lowest precedence among the pratt-parsed expressions

  # DO NOT ADD ANYTHING HERE
  # only use the pratt table for parsing expressions!
  comp.parsePrecedence(pcAssignment)


proc number(comp: Compiler) =
  # assume the number is already advanced through
  let value = comp.previous.text.parseFloat.toKonValue
  comp.writeConstant(value)
  when debugCompiler:
    debugEcho &"Written constant (type: {value.konType}, str repr: {$value}) to chunk"

tkNumber.genRule(number, nop, pcNone)

proc expFalse(comp: Compiler) =
  comp.writeChunk(1, opFalse)

tkFalse.genRule(expFalse, nop, pcNone)

proc expTrue(comp: Compiler) =
  comp.writeChunk(1, opTrue)

tkTrue.genRule(expTrue, nop, pcNone)

proc expNil(comp: Compiler) =
  comp.writeChunk(1, opNil)

tkNil.genRule(expNil, nop, pcNone)

proc expString(comp: Compiler) =
  let value = comp.previous.text[1..^2].toKonValue()
  comp.writeConstant(value)
  when debugCompiler:
    debugEcho &"Written constant (type: {value.konType}, str repr: {$value}) to chunk"

tkString.genRule(expString, nop, pcNone)

proc resolveLocal(comp: Compiler, name: string): int =
  ## returns the stack index of the local of the name
  var i = comp.locals.high
  while i >= 0:
    let local = comp.locals[i]
    if local.name == name:
      if local.depth == -1:
        comp.error("Can't read local variable in its own initializer.")
      return local.index
    i.dec
  return -1

proc variable(comp: Compiler) =
  # named variable
  var getOp = opGetGlobal
  var setOp = opSetGlobal

  let name = comp.previous.text

  # try resolving local, set arg to the index on the stack
  var arg = comp.resolveLocal(name)

  if arg != -1:
    # local
    getOp = opGetLocal
    setOp = opSetLocal
  else:
    # global
    arg = comp.chunk.addConstant(name.toKonValue)

  if comp.match(tkEqual):
    # assignment (global/local)
    if not comp.canAssign:
      comp.error("Invalid assignment target.")
      return
    comp.expression()
    comp.writeChunk(0, setOp)
  else:
    # get (global/local)
    comp.writeChunk(1, getOp)

  comp.writeChunk(0, arg.toTriple)

tkIdentifier.genRule(variable, nop, pcNone)

proc grouping(comp: Compiler) =
  # assume initial '(' is already consumed
  comp.expression()
  comp.consume(tkRightParen, "Expect ')' after expression.")  

tkLeftParen.genRule(grouping, nop, pcNone)

proc unary(comp: Compiler) =
  let opType = comp.previous.tokenType

  comp.parsePrecedence(pcUnary)

  case opType:
    of tkMinus:
      comp.writeChunk(0, opNegate)
    of tkBang:
      comp.writeChunk(0, opNot)
    else:
      discard # unreachable

tkBang.genRule(unary, nop, pcNone)

proc binary(comp: Compiler) =
  let opType = comp.previous.tokenType
  # safety checked in parsePrecedence
  let rule = opType.getRule.get
  comp.parsePrecedence(rule.applyRule.increment)

  case opType:
    of tkPlus:
      comp.writeChunk(-1, opAdd)
    of tkMinus:
      comp.writeChunk(-1, opSubtract)
    of tkStar:
      comp.writeChunk(-1, opMultiply)
    of tkSlash:
      comp.writeChunk(-1, opDivide)
    of tkEqualEqual:
      comp.writeChunk(-1, opEqual)
    of tkBangEqual:
      comp.writeChunk(-1, opEqual)
      comp.writeChunk(0, opNot)
    of tkGreater:
      comp.writeChunk(-1, opGreater)
    of tkLess:
      comp.writeChunk(-1, opLess)
    of tkGreaterEqual:
      comp.writeChunk(-1, opLess)
      comp.writeChunk(0, opNot)
    of tkLessEqual:
      comp.writeChunk(-1, opGreater)
      comp.writeChunk(0, opNot)
    else:
      return # unreachable

tkMinus.genRule(unary, binary, pcTerm)
tkPlus.genRule(nop, binary, pcTerm)
tkSlash.genRule(nop, binary, pcFactor)
tkStar.genRule(nop, binary, pcFactor)

tkEqualEqual.genRule(nop, binary, pcEquality)
tkBangEqual.genRule(nop, binary, pcEquality)
tkGreater.genRule(nop, binary, pcComparison)
tkGreaterEqual.genRule(nop, binary, pcComparison)
tkLess.genRule(nop, binary, pcComparison)
tkLessEqual.genRule(nop, binary, pcComparison)

proc ifExpr(comp: Compiler) =
  # if expressions return the body if condition is truthy,
  # the else expression otherwise, unless there is no else:
  # if there is no else, it returns the condition if it is falsey
  comp.consume(tkLeftParen, "Expect '(' after 'if'.")
  comp.expression()
  comp.consume(tkRightParen, "Expect ')' after condition.")

  let thenJump = comp.emitJump(0, opJumpIfFalse)

  # conditional code that can be jumped over must leave the stack in tact!

  comp.writeChunk(-1, opPop)
  comp.expression()
  # net change to stack: -1 + 1 = 0
  
  let elseJump = comp.emitJump(0, opJump)
  comp.patchJump(thenJump)

  if comp.match(tkElse):
    comp.writeChunk(-1, opPop)
    comp.expression()
  
  comp.patchJump(elseJump)


tkIf.genRule(ifExpr, nop, pcNone)

proc andExpr(comp: Compiler) =
  let endJump = comp.emitJump(0, opJumpIfFalse)

  comp.writeChunk(-1, opPop)
  comp.parsePrecedence(pcAnd)
  # net effect on stack: -1 + 1 = 0

  comp.patchJump(endJump)

tkAnd.genRule(nop, andExpr, pcAnd)

proc orExpr(comp: Compiler) =
  let elseJump = comp.emitJump(0, opJumpIfFalse)
  let endJump = comp.emitJump(0, opJump)

  comp.patchJump(elseJump)

  comp.writeChunk(-1, opPop)
  comp.parsePrecedence(pcOr)
  # net effect on stack: -1 + 1 = 0

  comp.patchJump(endJump)

tkOr.genRule(nop, orExpr, pcOr)

proc debugExpr(comp: Compiler) =
  comp.expression()
  when debugCompiler:
    debugEcho &"debug expression, current stackindex: {comp.stackIndex}"
  comp.writeChunk(0, opPrint)

tkPrint.genRule(debugExpr, nop, pcNone)

proc parseWhile(comp: Compiler) =

  comp.writeChunk(1, opNil) # return value
  let loopStart = comp.chunk.len
  comp.consume(tkLeftParen, "Expect '(' after 'while'.")
  # condition
  comp.expression()
  comp.consume(tkRightParen, "Expect ')' after condition.")

  let exitJump = comp.emitJump(-1, opJumpIfFalsePop)
  # this cannot be handled with just opPop, since the net change in the
  # stack size inside code that is conditional must be 0!

  # body
  comp.writeChunk(-1, opPop) # pop the old result
  comp.expression()
  # net stack change: 1 + -1 = 0

  comp.emitLoop(loopstart = loopStart, delta = 0, op = opLoop)
  comp.patchJump(exitJump)
  

tkWhile.genRule(parseWhile, nop, pcNone)

# below are the expressions that can contain statements in some way
# the only expressions that can contain a statement are:
  # the block expression
proc statement(comp: Compiler)

proc parseBlock(comp: Compiler) =
  ## Despite the name, can be used for statements if the arg statement is true
  ## Also can be used for function bodies
  comp.beginScope()
  while comp.current.tokenType != tkRightBrace and comp.current.tokenType != tkEof:
    comp.statement()
  comp.endScope()
  
  comp.consume(tkRightBrace, "Expect '}' after block.")

tkLeftBrace.genRule(parseBlock, nop, pcNone)

# statements

proc parseVariable(comp: Compiler, msg: string): int =
  ## Parses variable declarations
  ## During manipulation with variables:
  ## if global:
  ## consume the identifier and return index to work with
  ## if local:
  ## register the name with the vm
  comp.consume(tkIdentifier, msg)
  let name = comp.previous.text

  if name[0] in {'^'}:
    comp.error("Illegal variable name.")

  if comp.scopes.len > 0: # declareVariable
    # local

    # check if name exists already within scope
    for i in countdown(comp.locals.high, 0):
      let local = comp.locals[i]
      if local.depth != -1 and local.depth < comp.scopes.len:
        break
        
      if name == local.name:
        comp.error("Already a variable with this name in this scope.")
        break

    comp.addLocal(name, 1)
    0 # index to the constant is irrelevant if the var is local
  else:
    # global
    comp.chunk.addConstant(name.toKonValue())

proc defineVariable(comp: Compiler, index: int) =
  ## Generate code that moves the variable on the stack
  ## to a variable at the right place in memory
  ## the right place is defined by the 3 following byte after the op
  ## the thing to move is the item below it

  if comp.scopes.len > 0:
    # local variable: it's already on the right place
    # but we need to mark initialized
    comp.markInitialized()
  else:
    comp.writeChunk(-1, opDefineGlobal)
    comp.writeChunk(0, index.toTriple)

proc varStatement(comp: Compiler) =
  let globalIndex = comp.parseVariable("Expect variable name.")

  if comp.match(tkEqual):
    comp.expression()
  else:
    comp.writeChunk(1, opNil)
  
  comp.defineVariable(globalIndex)


proc breakStatement(comp: Compiler) =
  if not comp.match(tkLabel):
    comp.error("Label expected after break.")
  
  let label = comp.previous.text[1..^1]
  for i in countdown(comp.scopes.high, 0):
    let scope = comp.scopes[i]
    if scope.labels.contains(label):
      comp.jumpToEnd(scope)
      break
  
  comp.consume(tkSemicolon, "Semicolon expected after break statement.")
  if comp.current.tokenType != tkRightBrace:
    comp.error("Break statement must be the last element inside the innermost block it is in.")

proc statement(comp: Compiler) =
  if comp.match(tkVar):
    comp.varStatement()
    comp.consume(tkSemicolon, "Semicolon expected after expression statement.")
  elif comp.match(tkBreak):
    comp.breakStatement()
  else:
    comp.expression()
    comp.writeChunk(-1, opPop)
    comp.consume(tkSemicolon, "Semicolon expected after expression statement.")

  if comp.panicMode:
    comp.synchronize()

proc compile*(comp: Compiler) =
  comp.scanner = newScanner(comp.source)
  comp.writeChunk(0, opNil) 
  # the starting stackIndex is 0, which points to this nil
  # it is correctly set to delta = 0!!!
  comp.advance()
  while comp.current.tokenType != tkEof:
    comp.statement()
  comp.writeChunk(-1, opReturn)
  when debugDumpChunk:
    if not comp.hadError:
      comp.chunk.disassembleChunk()