package prog8.compiler.target.c64.codegen

import prog8.ast.INameScope
import prog8.ast.Node
import prog8.ast.Program
import prog8.ast.antlr.escape
import prog8.ast.base.*
import prog8.ast.expressions.*
import prog8.ast.statements.*
import prog8.compiler.*
import prog8.compiler.target.CompilationTarget
import prog8.compiler.target.CpuType
import prog8.compiler.target.IAssemblyGenerator
import prog8.compiler.target.IAssemblyProgram
import prog8.compiler.target.c64.AssemblyProgram
import prog8.compiler.target.c64.Petscii
import prog8.compiler.target.c64.codegen.assignment.AsmAssignSource
import prog8.compiler.target.c64.codegen.assignment.AsmAssignTarget
import prog8.compiler.target.c64.codegen.assignment.AsmAssignment
import prog8.compiler.target.c64.codegen.assignment.AssignmentAsmGen
import prog8.compiler.target.c64.codegen.assignment.TargetStorageKind
import prog8.compiler.target.generatedLabelPrefix
import prog8.functions.BuiltinFunctions
import prog8.functions.FSignature
import java.nio.file.Path
import java.time.LocalDate
import java.time.LocalDateTime
import java.util.*
import kotlin.math.absoluteValue


internal class AsmGen(private val program: Program,
                      val errors: ErrorReporter,
                      val zeropage: Zeropage,
                      val options: CompilationOptions,
                      private val outputDir: Path): IAssemblyGenerator {

    // for expressions and augmented assignments:
    val optimizedByteMultiplications = setOf(3,5,6,7,9,10,11,12,13,14,15,20,25,40,50,80,100)
    val optimizedWordMultiplications = setOf(3,5,6,7,9,10,12,15,20,25,40,50,80,100)

    private val assemblyLines = mutableListOf<String>()
    private val globalFloatConsts = mutableMapOf<Double, String>()     // all float values in the entire program (value -> varname)
    private val allocatedZeropageVariables = mutableMapOf<String, Pair<Int, DataType>>()
    private val breakpointLabels = mutableListOf<String>()
    private val builtinFunctionsAsmGen = BuiltinFunctionsAsmGen(program, this)
    private val forloopsAsmGen = ForLoopsAsmGen(program, this)
    private val postincrdecrAsmGen = PostIncrDecrAsmGen(program, this)
    private val functioncallAsmGen = FunctionCallAsmGen(program, this)
    private val expressionsAsmGen = ExpressionsAsmGen(program, this)
    private val assignmentAsmGen = AssignmentAsmGen(program, this, expressionsAsmGen)
    internal val loopEndLabels = ArrayDeque<String>()
    private val blockLevelVarInits = mutableMapOf<Block, MutableSet<VarDecl>>()

    override fun compileToAssembly(optimize: Boolean): IAssemblyProgram {
        assemblyLines.clear()
        loopEndLabels.clear()

        println("Generating assembly code... ")

        header()
        val allBlocks = program.allBlocks()
        if(allBlocks.first().name != "main")
            throw AssemblyError("first block should be 'main'")
        for(b in program.allBlocks())
            block2asm(b)
        footer()

        if(optimize) {
            var optimizationsDone = 1
            while (optimizationsDone > 0) {
                optimizationsDone = optimizeAssembly(assemblyLines)
            }
        }

        val outputFile = outputDir.resolve("${program.name}.asm").toFile()
        outputFile.printWriter().use {
            for (line in assemblyLines) { it.println(line) }
        }

        return AssemblyProgram(program.name, outputDir)
    }


    private fun header() {
        val ourName = this.javaClass.name
        val cpu = when(CompilationTarget.instance.machine.cpu) {
            CpuType.CPU6502 -> "6502"
            CpuType.CPU65c02 -> "65c02"
            else -> "unsupported"
        }

        out("; $cpu assembly code for '${program.name}'")
        out("; generated by $ourName on ${LocalDateTime.now().withNano(0)}")
        out("; assembler syntax is for the 64tasm cross-assembler")
        out("; output options: output=${options.output} launcher=${options.launcher} zp=${options.zeropage}")
        out("\n.cpu  '$cpu'\n.enc  'none'\n")

        program.actualLoadAddress = program.definedLoadAddress
        if (program.actualLoadAddress == 0)   // fix load address
            program.actualLoadAddress = if (options.launcher == LauncherType.BASIC)
                CompilationTarget.instance.machine.BASIC_LOAD_ADDRESS else CompilationTarget.instance.machine.RAW_LOAD_ADDRESS

        // the global prog8 variables needed
        val zp = CompilationTarget.instance.machine.zeropage
        out("P8ZP_SCRATCH_B1 = ${zp.SCRATCH_B1}")
        out("P8ZP_SCRATCH_REG = ${zp.SCRATCH_REG}")
        out("P8ZP_SCRATCH_W1 = ${zp.SCRATCH_W1}    ; word")
        out("P8ZP_SCRATCH_W2 = ${zp.SCRATCH_W2}    ; word")
        out("P8ESTACK_LO = ${CompilationTarget.instance.machine.ESTACK_LO.toHex()}")
        out("P8ESTACK_HI = ${CompilationTarget.instance.machine.ESTACK_HI.toHex()}")

        when {
            options.launcher == LauncherType.BASIC -> {
                if (program.actualLoadAddress != 0x0801)
                    throw AssemblyError("BASIC output must have load address $0801")
                out("; ---- basic program with sys call ----")
                out("* = ${program.actualLoadAddress.toHex()}")
                val year = LocalDate.now().year
                out("  .word  (+), $year")
                out("  .null  $9e, format(' %d ', _prog8_entrypoint), $3a, $8f, ' prog8 by idj'")
                out("+\t.word  0")
                out("_prog8_entrypoint\t; assembly code starts here\n")
                if(!options.noSysInit)
                    out("  jsr  ${CompilationTarget.instance.name}.init_system")
            }
            options.output == OutputType.PRG -> {
                out("; ---- program without basic sys call ----")
                out("* = ${program.actualLoadAddress.toHex()}\n")
                if(!options.noSysInit)
                    out("  jsr  ${CompilationTarget.instance.name}.init_system")
            }
            options.output == OutputType.RAW -> {
                out("; ---- raw assembler program ----")
                out("* = ${program.actualLoadAddress.toHex()}\n")
            }
        }

        if(options.zeropage !in setOf(ZeropageType.BASICSAFE, ZeropageType.DONTUSE)) {
            out("""
                ; zeropage is clobbered so we need to reset the machine at exit
                lda  #>${CompilationTarget.instance.name}.reset_system
                pha
                lda  #<${CompilationTarget.instance.name}.reset_system
                pha""")
        }

        out("  jmp  main.start  ; start program / force start proc to be included")
    }

    private fun footer() {
        // the global list of all floating point constants for the whole program
        out("; global float constants")
        for (flt in globalFloatConsts) {
            val floatFill = CompilationTarget.instance.machine.getFloat(flt.key).makeFloatFillAsm()
            val floatvalue = flt.key
            out("${flt.value}\t.byte  $floatFill  ; float $floatvalue")
        }
    }

    private fun block2asm(block: Block) {
        out("\n\n; ---- block: '${block.name}' ----")
        out("${block.name}\t" + (if("force_output" in block.options()) ".block\n" else ".proc\n"))

        if(block.address!=null) {
            out(".cerror * > ${block.address.toHex()}, 'block address overlaps by ', *-${block.address.toHex()},' bytes'")
            out("* = ${block.address.toHex()}")
        }

        outputSourceLine(block)
        zeropagevars2asm(block.statements)
        memdefs2asm(block.statements)
        vardecls2asm(block.statements)
        out("\n; subroutines in this block")

        // first translate regular statements, and then put the subroutines at the end.
        val (subroutine, stmts) = block.statements.partition { it is Subroutine }
        stmts.forEach { translate(it) }
        subroutine.forEach { translateSubroutine(it as Subroutine) }

        // if any global vars need to be initialized, generate a subroutine that does this
        // it will be called from program init.
        if(block in blockLevelVarInits) {
            out("prog8_init_vars\t.proc\n")
            blockLevelVarInits.getValue(block).forEach { decl ->
                val scopedFullName = decl.makeScopedName(decl.name).split('.')
                require(scopedFullName.first()==block.name)
                assignInitialValueToVar(decl, scopedFullName.drop(1))
            }
            out("  rts\n  .pend")
        }

        out(if("force_output" in block.options()) "\n\t.bend\n" else "\n\t.pend\n")
    }

    private fun assignInitialValueToVar(decl: VarDecl, variableName: List<String>) {
        val variable = IdentifierReference(variableName, decl.position)
        variable.linkParents(decl.parent)
        val asgn = AsmAssignment(
                AsmAssignSource.fromAstSource(decl.value!!, program),
                AsmAssignTarget(TargetStorageKind.VARIABLE, program, this, decl.datatype, variable = variable),
                false, decl.position)
        assignmentAsmGen.translateNormalAssignment(asgn)
    }

    private var generatedLabelSequenceNumber: Int = 0

    internal fun makeLabel(postfix: String): String {
        generatedLabelSequenceNumber++
        return "${generatedLabelPrefix}${generatedLabelSequenceNumber}_$postfix"
    }

    private fun outputSourceLine(node: Node) {
        out(" ;\tsrc line: ${node.position.file}:${node.position.line}")
    }

    internal fun out(str: String, splitlines: Boolean = true) {
        val fragment = (if(" | " in str) str.replace("|", "\n") else str).trim('\n')

        if (splitlines) {
            for (line in fragment.split('\n')) {
                val trimmed = if (line.startsWith(' ')) "\t" + line.trim() else line.trim()
                // trimmed = trimmed.replace(Regex("^\\+\\s+"), "+\t")  // sanitize local label indentation
                assemblyLines.add(trimmed)
            }
        } else assemblyLines.add(fragment)
    }

    private fun encode(str: String, altEncoding: Boolean): List<Short> {
         val bytes = if(altEncoding) Petscii.encodeScreencode(str, true) else Petscii.encodePetscii(str, true)
         return bytes.plus(0)
    }

    private fun zeropagevars2asm(statements: List<Statement>) {
        out("; vars allocated on zeropage")
        val variables = statements.filterIsInstance<VarDecl>().filter { it.type==VarDeclType.VAR }
        for(variable in variables) {
            val fullName = variable.makeScopedName(variable.name)
            val zpVar = allocatedZeropageVariables[fullName]
            if(zpVar==null) {
                // This var is not on the ZP yet. Attempt to move it there (if it's not a float, those take up too much space)
                if(variable.zeropage != ZeropageWish.NOT_IN_ZEROPAGE &&
                        variable.datatype in zeropage.allowedDatatypes
                        && variable.datatype != DataType.FLOAT
                        && options.zeropage != ZeropageType.DONTUSE) {
                    try {
                        val errors = ErrorReporter()
                        val address = zeropage.allocate(fullName, variable.datatype, null, errors)
                        errors.handle()
                        out("${variable.name} = $address\t; auto zp ${variable.datatype}")
                        // make sure we add the var to the set of zpvars for this block
                        allocatedZeropageVariables[fullName] = Pair(address, variable.datatype)
                    } catch (x: ZeropageDepletedError) {
                        // leave it as it is.
                    }
                }
            }
        }
    }

    private fun vardecl2asm(decl: VarDecl) {
        val name = decl.name
        when (decl.datatype) {
            DataType.UBYTE -> out("$name\t.byte  0")
            DataType.BYTE -> out("$name\t.char  0")
            DataType.UWORD -> out("$name\t.word  0")
            DataType.WORD -> out("$name\t.sint  0")
            DataType.FLOAT -> out("$name\t.byte  0,0,0,0,0  ; float")
            DataType.STRUCT -> {}       // is flattened
            DataType.STR -> {
                val str = decl.value as StringLiteralValue
                outputStringvar(decl, encode(str.value, str.altEncoding))
            }
            DataType.ARRAY_UB -> {
                val data = makeArrayFillDataUnsigned(decl)
                if (data.size <= 16)
                    out("$name\t.byte  ${data.joinToString()}")
                else {
                    out(name)
                    for (chunk in data.chunked(16))
                        out("  .byte  " + chunk.joinToString())
                }
            }
            DataType.ARRAY_B -> {
                val data = makeArrayFillDataSigned(decl)
                if (data.size <= 16)
                    out("$name\t.char  ${data.joinToString()}")
                else {
                    out(name)
                    for (chunk in data.chunked(16))
                        out("  .char  " + chunk.joinToString())
                }
            }
            DataType.ARRAY_UW -> {
                val data = makeArrayFillDataUnsigned(decl)
                if (data.size <= 16)
                    out("$name\t.word  ${data.joinToString()}")
                else {
                    out(name)
                    for (chunk in data.chunked(16))
                        out("  .word  " + chunk.joinToString())
                }
            }
            DataType.ARRAY_W -> {
                val data = makeArrayFillDataSigned(decl)
                if (data.size <= 16)
                    out("$name\t.sint  ${data.joinToString()}")
                else {
                    out(name)
                    for (chunk in data.chunked(16))
                        out("  .sint  " + chunk.joinToString())
                }
            }
            DataType.ARRAY_F -> {
                val array =
                        if(decl.value!=null)
                            (decl.value as ArrayLiteralValue).value
                        else {
                            // no init value, use zeros
                            val zero = decl.zeroElementValue()
                            Array(decl.arraysize!!.constIndex()!!) { zero }
                        }
                val floatFills = array.map {
                    val number = (it as NumericLiteralValue).number
                    CompilationTarget.instance.machine.getFloat(number).makeFloatFillAsm()
                }
                out(name)
                for (f in array.zip(floatFills))
                    out("  .byte  ${f.second}  ; float ${f.first}")
            }
        }
    }

    private fun memdefs2asm(statements: List<Statement>) {
        out("\n; memdefs and kernel subroutines")
        val memvars = statements.filterIsInstance<VarDecl>().filter { it.type==VarDeclType.MEMORY || it.type==VarDeclType.CONST }
        for(m in memvars) {
            if(m.value is NumericLiteralValue)
                out("  ${m.name} = ${(m.value as NumericLiteralValue).number.toHex()}")
            else
                out("  ${m.name} = ${asmVariableName((m.value as AddressOf).identifier)}")
        }
        val asmSubs = statements.filterIsInstance<Subroutine>().filter { it.isAsmSubroutine }
        for(sub in asmSubs) {
            if(sub.asmAddress!=null) {
                if(sub.statements.isNotEmpty())
                    throw AssemblyError("kernel subroutine cannot have statements")
                out("  ${sub.name} = ${sub.asmAddress.toHex()}")
            }
        }
    }

    private fun vardecls2asm(statements: List<Statement>) {
        out("\n; non-zeropage variables")
        val vars = statements.filterIsInstance<VarDecl>().filter { it.type==VarDeclType.VAR }

        // first output the flattened struct member variables *in order*
        // after that, the other variables sorted by their datatype

        val (structMembers, normalVars) = vars.partition { it.struct!=null }
        structMembers.forEach { vardecl2asm(it) }

        // special treatment for string types: merge strings that are identical
        val encodedstringVars = normalVars
                .filter {it.datatype == DataType.STR }
                .map {
                    val str = it.value as StringLiteralValue
                    it to encode(str.value, str.altEncoding)
                }
                .groupBy({it.second}, {it.first})
        for((encoded, variables) in encodedstringVars) {
            variables.dropLast(1).forEach { out(it.name) }
            val lastvar = variables.last()
            outputStringvar(lastvar, encoded)
        }

        // non-string variables
        normalVars.filter{ it.datatype != DataType.STR }.sortedBy { it.datatype }.forEach {
            if(it.makeScopedName(it.name) !in allocatedZeropageVariables)
                vardecl2asm(it)
        }
    }

    private fun outputStringvar(lastvar: VarDecl, encoded: List<Short>) {
        val string = (lastvar.value as StringLiteralValue).value
        out("${lastvar.name}\t; ${lastvar.datatype} \"${escape(string).replace("\u0000", "<NULL>")}\"")
        val outputBytes = encoded.map { "$" + it.toString(16).padStart(2, '0') }
        for (chunk in outputBytes.chunked(16))
            out("  .byte  " + chunk.joinToString())
    }

    private fun makeArrayFillDataUnsigned(decl: VarDecl): List<String> {
        val array =
                if(decl.value!=null)
                    (decl.value as ArrayLiteralValue).value
                else {
                    // no array init value specified, use a list of zeros
                    val zero = decl.zeroElementValue()
                    Array(decl.arraysize!!.constIndex()!!) { zero }
                }
        return when (decl.datatype) {
            DataType.ARRAY_UB ->
                // byte array can never contain pointer-to types, so treat values as all integers
                array.map {
                    val number = (it as NumericLiteralValue).number.toInt()
                    "$"+number.toString(16).padStart(2, '0')
                }
            DataType.ARRAY_UW -> array.map {
                when (it) {
                    is NumericLiteralValue -> {
                        "$" + it.number.toInt().toString(16).padStart(4, '0')
                    }
                    is AddressOf -> {
                        it.identifier.firstStructVarName(program.namespace) ?: asmSymbolName(it.identifier)
                    }
                    is IdentifierReference -> {
                        it.firstStructVarName(program.namespace) ?: asmSymbolName(it)
                    }
                    else -> throw AssemblyError("weird array elt dt")
                }
            }
            else -> throw AssemblyError("invalid arraysize type")
        }
    }

    private fun makeArrayFillDataSigned(decl: VarDecl): List<String> {
        val array =
                if(decl.value!=null)
                    (decl.value as ArrayLiteralValue).value
                else {
                    // no array init value specified, use a list of zeros
                    val zero = decl.zeroElementValue()
                    Array(decl.arraysize!!.constIndex()!!) { zero }
                }
        return when (decl.datatype) {
            DataType.ARRAY_UB ->
                // byte array can never contain pointer-to types, so treat values as all integers
                array.map {
                    val number = (it as NumericLiteralValue).number.toInt()
                    "$"+number.toString(16).padStart(2, '0')
                }
            DataType.ARRAY_B ->
                // byte array can never contain pointer-to types, so treat values as all integers
                array.map {
                    val number = (it as NumericLiteralValue).number.toInt()
                    val hexnum = number.absoluteValue.toString(16).padStart(2, '0')
                    if(number>=0)
                        "$$hexnum"
                    else
                        "-$$hexnum"
                }
            DataType.ARRAY_UW -> array.map {
                val number = (it as NumericLiteralValue).number.toInt()
                "$" + number.toString(16).padStart(4, '0')
            }
            DataType.ARRAY_W -> array.map {
                val number = (it as NumericLiteralValue).number.toInt()
                val hexnum = number.absoluteValue.toString(16).padStart(4, '0')
                if(number>=0)
                    "$$hexnum"
                else
                    "-$$hexnum"
            }
            else -> throw AssemblyError("invalid arraysize type ${decl.datatype}")
        }
    }

    internal fun getFloatAsmConst(number: Double): String {
        var asmName = CompilationTarget.instance.machine.getFloatRomConst(number)
        if(asmName.isNullOrEmpty()) {
            // no ROM float const for this value, create our own
            asmName = globalFloatConsts[number]
            if(asmName==null) {
                asmName = "prog8_float_const_${globalFloatConsts.size}"
                globalFloatConsts[number] = asmName
            }
        }
        return asmName
    }

    internal fun asmSymbolName(identifier: IdentifierReference): String {
        return if(identifier.memberOfStruct(program.namespace)!=null) {
            val name = identifier.targetVarDecl(program.namespace)!!.name
            fixNameSymbols(name)
        } else {
            fixNameSymbols(identifier.nameInSource.joinToString("."))
        }
    }

    internal fun asmVariableName(identifier: IdentifierReference): String {
        return if(identifier.memberOfStruct(program.namespace)!=null) {
            val name = identifier.targetVarDecl(program.namespace)!!.name
            fixNameSymbols(name)
        } else {
            fixNameSymbols(identifier.nameInSource.joinToString("."))
        }
    }

    internal fun loadByteFromPointerIntoA(pointervar: IdentifierReference): Pair<Boolean, String> {
        // returns if the pointer is already on the ZP itself or not (in which case SCRATCH_W1 is used as intermediary)
        val sourceName = asmVariableName(pointervar)
        val vardecl = pointervar.targetVarDecl(program.namespace)!!
        val scopedName = vardecl.makeScopedName(vardecl.name)
        return if(scopedName in allocatedZeropageVariables) {
            // pointervar is already in the zero page, no need to copy
            out(" ldy  #0 |  lda  ($sourceName),y")
            Pair(true, sourceName)
        } else {
            out("""
                lda  $sourceName
                ldy  $sourceName+1
                sta  P8ZP_SCRATCH_W1
                sty  P8ZP_SCRATCH_W1+1
                ldy  #0
                lda  (P8ZP_SCRATCH_W1),y""")
            Pair(false, sourceName)
        }
    }

    fun storeByteIntoPointer(pointervar: IdentifierReference, ldaInstructionArg: String?) {
        val sourceName = asmVariableName(pointervar)
        val vardecl = pointervar.targetVarDecl(program.namespace)!!
        val scopedName = vardecl.makeScopedName(vardecl.name)
        if(scopedName in allocatedZeropageVariables) {
            // pointervar is already in the zero page, no need to copy
            if(ldaInstructionArg!=null)
                out("  lda  $ldaInstructionArg")
            out(" ldy  #0 |  sta  ($sourceName),y")
        } else {
            out("""
                ldy  $sourceName
                sty  P8ZP_SCRATCH_W2
                ldy  $sourceName+1
                sty  P8ZP_SCRATCH_W2+1
                ${if(ldaInstructionArg==null) "" else "lda  $ldaInstructionArg"}
                ldy  #0
                sta  (P8ZP_SCRATCH_W2),y""")
        }
    }

    internal fun fixNameSymbols(name: String) = name.replace("<", "prog8_").replace(">", "")     // take care of the autogenerated invalid (anon) label names


    private val saveRegisterLabels = Stack<String>();

    internal fun saveRegister(register: CpuRegister) {
        when(register) {
            CpuRegister.A -> out("  pha")
            CpuRegister.X -> {
                if (CompilationTarget.instance.machine.cpu == CpuType.CPU65c02) out("  phx")
                else out("  stx  _prog8_regsave${register.name}")
            }
            CpuRegister.Y -> {
                if (CompilationTarget.instance.machine.cpu == CpuType.CPU65c02) out("  phy")
                else out("  sty  _prog8_regsave${register.name}")
            }
        }
    }

    internal fun restoreRegister(register: CpuRegister) {
        when(register) {
            CpuRegister.A -> out("  pla")
            CpuRegister.X -> {
                if (CompilationTarget.instance.machine.cpu == CpuType.CPU65c02) out("  plx")
                else out("  ldx  _prog8_regsave${register.name}")
            }
            CpuRegister.Y -> {
                if (CompilationTarget.instance.machine.cpu == CpuType.CPU65c02) out("  ply")
                else out("  ldy  _prog8_regsave${register.name}")
            }
        }
    }


    internal fun translate(stmt: Statement) {
        outputSourceLine(stmt)
        when(stmt) {
            is ParameterVarDecl -> { /* subroutine parameter vardecls don't get any special treatment here */ }
            is VarDecl -> translate(stmt)
            is StructDecl, is NopStatement -> {}
            is Directive -> translate(stmt)
            is Return -> translate(stmt)
            is Subroutine -> translateSubroutine(stmt)
            is InlineAssembly -> translate(stmt)
            is FunctionCallStatement -> {
                val functionName = stmt.target.nameInSource.last()
                val builtinFunc = BuiltinFunctions[functionName]
                if (builtinFunc != null) {
                    builtinFunctionsAsmGen.translateFunctioncallStatement(stmt, builtinFunc)
                } else {
                    functioncallAsmGen.translateFunctionCall(stmt)
                    // discard any results from the stack:
                    val sub = stmt.target.targetSubroutine(program.namespace)!!
                    val returns = sub.returntypes.zip(sub.asmReturnvaluesRegisters)
                    for ((t, reg) in returns) {
                        if (reg.stack) {
                            if (t in IntegerDatatypes || t in PassByReferenceDatatypes) out("  inx")
                            else if (t == DataType.FLOAT) out("  inx |  inx |  inx")
                        }
                    }
                }
            }
            is Assignment -> assignmentAsmGen.translate(stmt)
            is Jump -> translate(stmt)
            is PostIncrDecr -> postincrdecrAsmGen.translate(stmt)
            is Label -> translate(stmt)
            is BranchStatement -> translate(stmt)
            is IfStatement -> translate(stmt)
            is ForLoop -> forloopsAsmGen.translate(stmt)
            is Break -> {
                if(loopEndLabels.isEmpty())
                    throw AssemblyError("break statement out of context  ${stmt.position}")
                out("  jmp  ${loopEndLabels.peek()}")
            }
            is WhileLoop -> translate(stmt)
            is RepeatLoop -> translate(stmt)
            is UntilLoop -> translate(stmt)
            is WhenStatement -> translate(stmt)
            is BuiltinFunctionStatementPlaceholder -> throw AssemblyError("builtin function should not have placeholder anymore?")
            is AnonymousScope -> translate(stmt)
            is Block -> throw AssemblyError("block should have been handled elsewhere")
            else -> throw AssemblyError("missing asm translation for $stmt")
        }
    }

    internal fun loadScaledArrayIndexIntoRegister(expr: ArrayIndexedExpression,
                                                  elementDt: DataType,
                                                  register: CpuRegister,
                                                  addOneExtra: Boolean=false) {
        val reg = register.toString().toLowerCase()
        val index = expr.arrayspec.index
        if(index is NumericLiteralValue) {
            val indexValue = index.number.toInt() * elementDt.memorySize() + if(addOneExtra) 1 else 0
            out("  ld$reg  #$indexValue")
            return
        }

        if(addOneExtra) {
            // add 1 to the result
            if (index is IdentifierReference) {
                val indexName = asmVariableName(index)
                when(elementDt) {
                    in ByteDatatypes -> {
                        out("  ldy  $indexName |  iny")
                        when(register) {
                            CpuRegister.A -> out(" tya")
                            CpuRegister.X -> out(" tyx")
                            CpuRegister.Y -> {}
                        }
                    }
                    in WordDatatypes -> {
                        out("  lda  $indexName |  sec |  rol a")
                        when(register) {
                            CpuRegister.A -> {}
                            CpuRegister.X -> out(" tax")
                            CpuRegister.Y -> out(" tay")
                        }
                    }
                    DataType.FLOAT -> {
                        require(DataType.FLOAT.memorySize()==5)
                        out("""
                                    lda  $indexName
                                    asl  a
                                    asl  a
                                    sec
                                    adc  $indexName""")
                        when(register) {
                            CpuRegister.A -> {}
                            CpuRegister.X -> out(" tax")
                            CpuRegister.Y -> out(" tay")
                        }
                    }
                    else -> throw AssemblyError("weird dt")
                }
            }
            else {
                expressionsAsmGen.translateExpression(index)
                out("""
                    inc  P8ESTACK_LO,x
                    bne  +
                    inc  P8ESTACK_HI,x
+""")
                when(register) {
                    CpuRegister.A -> out("  inx |  lda  P8ESTACK_LO,x")
                    CpuRegister.X -> out("  inx |  lda  P8ESTACK_LO,x |  tax")
                    CpuRegister.Y -> out("  inx |  ldy  P8ESTACK_LO,x")
                }
            }
        } else {
            if (index is IdentifierReference) {
                val indexName = asmVariableName(index)
                when(elementDt) {
                    in ByteDatatypes -> out("  ld$reg  $indexName")
                    in WordDatatypes -> {
                        out("  lda  $indexName |  asl a")
                        when(register) {
                            CpuRegister.A -> {}
                            CpuRegister.X -> out(" tax")
                            CpuRegister.Y -> out(" tay")
                        }
                    }
                    DataType.FLOAT -> {
                        require(DataType.FLOAT.memorySize()==5)
                        out("""
                                    lda  $indexName
                                    asl  a
                                    asl  a
                                    clc
                                    adc  $indexName""")
                        when(register) {
                            CpuRegister.A -> {}
                            CpuRegister.X -> out(" tax")
                            CpuRegister.Y -> out(" tay")
                        }
                    }
                    else -> throw AssemblyError("weird dt")
                }
            }
            else {
                expressionsAsmGen.translateExpression(index)
                when(elementDt) {
                    in ByteDatatypes -> {
                        when (register) {
                            CpuRegister.A -> out("  inx |  lda  P8ESTACK_LO,x")
                            CpuRegister.X -> out("  inx |  lda  P8ESTACK_LO,x |  tax")
                            CpuRegister.Y -> out("  inx |  ldy  P8ESTACK_LO,x")
                        }
                    }
                    in WordDatatypes -> {
                        out("""
                            inx
                            lda  P8ESTACK_LO,x
                            asl  a""")
                        when (register) {
                            CpuRegister.A -> {}
                            CpuRegister.X -> out("  tax")
                            CpuRegister.Y -> out("  tay")
                        }
                    }
                    DataType.FLOAT -> {
                        require(DataType.FLOAT.memorySize()==5)
                        out("""
                            inx
                            lda  P8ESTACK_LO,x
                            asl  a
                            asl  a
                            clc
                            adc  P8ESTACK_LO,x""")
                        when (register) {
                            CpuRegister.A -> {}
                            CpuRegister.X -> out("  tax")
                            CpuRegister.Y -> out("  tay")
                        }
                    }
                }
            }
        }
    }

    internal fun translateExpression(expression: Expression) =
            expressionsAsmGen.translateExpression(expression)

    internal fun translateFunctioncallExpression(functionCall: FunctionCall, signature: FSignature) =
            builtinFunctionsAsmGen.translateFunctioncallExpression(functionCall, signature)

    internal fun translateFunctionCall(functionCall: FunctionCall) =
            functioncallAsmGen.translateFunctionCall(functionCall)

    internal fun translateNormalAssignment(assign: AsmAssignment) =
            assignmentAsmGen.translateNormalAssignment(assign)

    private fun translateSubroutine(sub: Subroutine) {
        out("")
        outputSourceLine(sub)

        if(sub.isAsmSubroutine) {
            if(sub.asmAddress!=null)
                return  // already done at the memvars section

            // asmsub with most likely just an inline asm in it
            out("${sub.name}\t.proc")
            sub.statements.forEach{ translate(it) }
            out("  .pend\n")
        } else {
            // regular subroutine
            out("${sub.name}\t.proc")
            zeropagevars2asm(sub.statements)
            memdefs2asm(sub.statements)

            // the main.start subroutine is the program's entrypoint and should perform some initialization logic
            if(sub.name=="start" && sub.definingBlock().name=="main") {
                out("; program startup initialization")
                out("  cld")
                program.allBlocks().forEach {
                    if(it.statements.filterIsInstance<VarDecl>().any { vd->vd.value!=null && vd.type==VarDeclType.VAR && vd.datatype in NumericDatatypes})
                        out("  jsr  ${it.name}.prog8_init_vars")
                }
                out("""
                    tsx
                    stx  prog8_lib.orig_stackpointer    ; required for func_exit                    
                    ldx  #255       ; init estack ptr
                    clv
                    clc""")
            }

            out("; statements")
            sub.statements.forEach{ translate(it) }
            out("; variables")
            out("""
                ; register saves
_prog8_regsaveX     .byte  0                
_prog8_regsaveY     .byte  0""")        // TODO only generate these bytes if they're actually used by saveRegister()
            vardecls2asm(sub.statements)
            out("  .pend\n")
        }
    }

    private fun branchInstruction(condition: BranchCondition, complement: Boolean) =
            if(complement) {
                when (condition) {
                    BranchCondition.CS -> "bcc"
                    BranchCondition.CC -> "bcs"
                    BranchCondition.EQ, BranchCondition.Z -> "beq"
                    BranchCondition.NE, BranchCondition.NZ -> "bne"
                    BranchCondition.VS -> "bvc"
                    BranchCondition.VC -> "bvs"
                    BranchCondition.MI, BranchCondition.NEG -> "bmi"
                    BranchCondition.PL, BranchCondition.POS -> "bpl"
                }
            } else {
                when (condition) {
                    BranchCondition.CS -> "bcs"
                    BranchCondition.CC -> "bcc"
                    BranchCondition.EQ, BranchCondition.Z -> "beq"
                    BranchCondition.NE, BranchCondition.NZ -> "bne"
                    BranchCondition.VS -> "bvs"
                    BranchCondition.VC -> "bvc"
                    BranchCondition.MI, BranchCondition.NEG -> "bmi"
                    BranchCondition.PL, BranchCondition.POS -> "bpl"
                }
            }

    private fun translate(stmt: IfStatement) {
        checkBooleanExpression(stmt.condition)  // we require the condition to be of the form  'x <comparison> <value>'
        val booleanCondition = stmt.condition as BinaryExpression

        if (stmt.elsepart.containsNoCodeNorVars()) {
            // empty else
            val endLabel = makeLabel("if_end")
            expressionsAsmGen.translateComparisonExpressionWithJumpIfFalse(booleanCondition, endLabel)
            translate(stmt.truepart)
            out(endLabel)
        }
        else {
            // both true and else parts
            val elseLabel = makeLabel("if_else")
            val endLabel = makeLabel("if_end")
            expressionsAsmGen.translateComparisonExpressionWithJumpIfFalse(booleanCondition, elseLabel)
            translate(stmt.truepart)
            out("  jmp  $endLabel")
            out(elseLabel)
            translate(stmt.elsepart)
            out(endLabel)
        }
    }

    private fun checkBooleanExpression(condition: Expression) {
        if(condition !is BinaryExpression || condition.operator !in comparisonOperators)
            throw AssemblyError("expected boolean expression $condition")
    }

    private fun translate(stmt: RepeatLoop) {
        val repeatLabel = makeLabel("repeat")
        val endLabel = makeLabel("repeatend")
        loopEndLabels.push(endLabel)

        when (stmt.iterations) {
            null -> {
                // endless loop
                out(repeatLabel)
                translate(stmt.body)
                out("  jmp  $repeatLabel")
                out(endLabel)
            }
            is NumericLiteralValue -> {
                val iterations = (stmt.iterations as NumericLiteralValue).number.toInt()
                if(iterations<0 || iterations > 65536)
                    throw AssemblyError("invalid number of iterations")
                when {
                    iterations == 0 -> {}
                    iterations <= 256 -> {
                        out("  lda  #${iterations and 255}")
                        repeatByteCountInA(iterations, repeatLabel, endLabel, stmt.body)
                    }
                    else -> {
                        out("  lda  #<${iterations} |  ldy  #>${iterations}")
                        repeatWordCountInAY(iterations, repeatLabel, endLabel, stmt.body)
                    }
                }
            }
            is IdentifierReference -> {
                val vardecl = (stmt.iterations as IdentifierReference).targetStatement(program.namespace) as VarDecl
                val name = asmVariableName(stmt.iterations as IdentifierReference)
                when(vardecl.datatype) {
                    DataType.UBYTE, DataType.BYTE -> {
                        out("  lda  $name")
                        repeatByteCountInA(null, repeatLabel, endLabel, stmt.body)
                    }
                    DataType.UWORD, DataType.WORD -> {
                        out("  lda  $name |  ldy  $name+1")
                        repeatWordCountInAY(null, repeatLabel, endLabel, stmt.body)
                    }
                    else -> throw AssemblyError("invalid loop variable datatype $vardecl")
                }
            }
            else -> {
                translateExpression(stmt.iterations!!)
                val dt = stmt.iterations!!.inferType(program).typeOrElse(DataType.STRUCT)
                when (dt) {
                    in ByteDatatypes -> {
                        out("  inx |  lda  P8ESTACK_LO,x")
                        repeatByteCountInA(null, repeatLabel, endLabel, stmt.body)
                    }
                    in WordDatatypes -> {
                        out("  inx |  lda  P8ESTACK_LO,x |  ldy  P8ESTACK_HI,x")
                        repeatWordCountInAY(null, repeatLabel, endLabel, stmt.body)
                    }
                    else -> throw AssemblyError("invalid loop expression datatype $dt")
                }
            }
        }

        loopEndLabels.pop()
    }

    private fun repeatWordCountInAY(constIterations: Int?, repeatLabel: String, endLabel: String, body: AnonymousScope) {
        // note: A/Y must have been loaded with the number of iterations already!
        val counterVar = makeLabel("repeatcounter")
        out("""
                sta  $counterVar
                sty  $counterVar+1
$repeatLabel    lda  $counterVar
                bne  +
                lda  $counterVar+1
                beq  $endLabel
+               lda  $counterVar
                bne  +
                dec  $counterVar+1
+               dec  $counterVar
""")
        translate(body)
        out("  jmp  $repeatLabel")
        if(constIterations!=null && constIterations>=16 && zeropage.available() > 1) {
            // allocate count var on ZP
            val zpAddr = zeropage.allocate(counterVar, DataType.UWORD, body.position, errors)
            out("""$counterVar = $zpAddr  ; auto zp UWORD""")
        } else {
            out("""
$counterVar    .word  0""")
        }
        out(endLabel)

    }

    private fun repeatByteCountInA(constIterations: Int?, repeatLabel: String, endLabel: String, body: AnonymousScope) {
        // note: A must have been loaded with the number of iterations already!
        val counterVar = makeLabel("repeatcounter")
        out("""
            sta  $counterVar
$repeatLabel""")
        translate(body)
        out("""
             dec  $counterVar
             bne  $repeatLabel
             beq  $endLabel""")
        if(constIterations!=null && constIterations>=16 && zeropage.available() > 0) {
            // allocate count var on ZP
            val zpAddr = zeropage.allocate(counterVar, DataType.UBYTE, body.position, errors)
            out("""$counterVar = $zpAddr  ; auto zp UBYTE""")
        } else {
            out("""
$counterVar    .byte  0""")
        }
        out(endLabel)
    }

    private fun translate(stmt: WhileLoop) {
        checkBooleanExpression(stmt.condition)  // we require the condition to be of the form  'x <comparison> <value>'
        val booleanCondition = stmt.condition as BinaryExpression
        val whileLabel = makeLabel("while")
        val endLabel = makeLabel("whileend")
        loopEndLabels.push(endLabel)
        out(whileLabel)
        expressionsAsmGen.translateComparisonExpressionWithJumpIfFalse(booleanCondition, endLabel)
        translate(stmt.body)
        out("  jmp  $whileLabel")
        out(endLabel)
        loopEndLabels.pop()
    }

    private fun translate(stmt: UntilLoop) {
        checkBooleanExpression(stmt.condition)  // we require the condition to be of the form  'x <comparison> <value>'
        val booleanCondition = stmt.condition as BinaryExpression
        val repeatLabel = makeLabel("repeat")
        val endLabel = makeLabel("repeatend")
        loopEndLabels.push(endLabel)
        out(repeatLabel)
        translate(stmt.body)
        expressionsAsmGen.translateComparisonExpressionWithJumpIfFalse(booleanCondition, repeatLabel)
        out(endLabel)
        loopEndLabels.pop()
    }

    private fun translate(stmt: WhenStatement) {
        expressionsAsmGen.translateExpression(stmt.condition)
        val endLabel = makeLabel("choice_end")
        val choiceBlocks = mutableListOf<Pair<String, AnonymousScope>>()
        val conditionDt = stmt.condition.inferType(program)
        if(!conditionDt.isKnown)
            throw AssemblyError("unknown condition dt")
        if(conditionDt.typeOrElse(DataType.BYTE) in ByteDatatypes)
            out("  inx |  lda  P8ESTACK_LO,x")
        else
            out("  inx |  lda  P8ESTACK_LO,x |  ldy  P8ESTACK_HI,x")
        for(choice in stmt.choices) {
            val choiceLabel = makeLabel("choice")
            if(choice.values==null) {
                // the else choice
                translate(choice.statements)
                out("  jmp  $endLabel")
            } else {
                choiceBlocks.add(Pair(choiceLabel, choice.statements))
                for (cv in choice.values!!) {
                    val value = (cv as NumericLiteralValue).number.toInt()
                    if(conditionDt.typeOrElse(DataType.BYTE) in ByteDatatypes) {
                        out("  cmp  #${value.toHex()} |  beq  $choiceLabel")
                    } else {
                        out("""
                            cmp  #<${value.toHex()}
                            bne  +
                            cpy  #>${value.toHex()}
                            beq  $choiceLabel
+
                            """)
                    }
                }
            }
        }
        out("  jmp  $endLabel")
        for(choiceBlock in choiceBlocks) {
            out(choiceBlock.first)
            translate(choiceBlock.second)
            out("  jmp  $endLabel")
        }
        out(endLabel)
    }

    private fun translate(stmt: Label) {
        out("_${stmt.name}")        // underscore prefix to make sure it's a local label
    }

    private fun translate(scope: AnonymousScope) {
        // note: the variables defined in an anonymous scope have been moved to their defining subroutine's scope
        scope.statements.forEach{ translate(it) }
    }

    private fun translate(stmt: BranchStatement) {
        if(stmt.truepart.containsNoCodeNorVars() && stmt.elsepart.containsCodeOrVars())
            throw AssemblyError("only else part contains code, shoud have been switched already")

        val jump = stmt.truepart.statements.first() as? Jump
        if(jump!=null) {
            // branch with only a jump
            val instruction = branchInstruction(stmt.condition, false)
            out("  $instruction  ${getJumpTarget(jump)}")
            translate(stmt.elsepart)
        } else {
            if(stmt.elsepart.containsNoCodeNorVars()) {
                val instruction = branchInstruction(stmt.condition, true)
                val elseLabel = makeLabel("branch_else")
                out("  $instruction  $elseLabel")
                translate(stmt.truepart)
                out(elseLabel)
            } else {
                val instruction = branchInstruction(stmt.condition, false)
                val trueLabel = makeLabel("branch_true")
                val endLabel = makeLabel("branch_end")
                out("  $instruction  $trueLabel")
                translate(stmt.elsepart)
                out("  jmp  $endLabel")
                out(trueLabel)
                translate(stmt.truepart)
                out(endLabel)
            }
        }
    }

    private fun translate(stmt: VarDecl) {
        if(stmt.value!=null && stmt.type==VarDeclType.VAR && stmt.datatype in NumericDatatypes) {
            // generate an assignment statement to (re)initialize the variable's value.
            // if the vardecl is not in a subroutine however, we have to initialize it globally.
            if(stmt.definingSubroutine()==null) {
                val block = stmt.definingBlock()
                var inits = blockLevelVarInits[block]
                if(inits==null) {
                    inits = mutableSetOf()
                    blockLevelVarInits[block] = inits
                }
                inits.add(stmt)
            } else {
                val next = (stmt.parent as INameScope).nextSibling(stmt)
                if (next !is ForLoop || next.loopVar.nameInSource.single() != stmt.name) {
                    assignInitialValueToVar(stmt, listOf(stmt.name))
                }
            }
        }
    }

    private fun translate(stmt: Directive) {
        when(stmt.directive) {
            "%asminclude" -> {
                val sourcecode = loadAsmIncludeFile(stmt.args[0].str!!, stmt.definingModule().source)
                val scopeprefix = stmt.args[1].str ?: ""
                if(!scopeprefix.isBlank())
                    out("$scopeprefix\t.proc")
                assemblyLines.add(sourcecode.trimEnd().trimStart('\n'))
                if(!scopeprefix.isBlank())
                    out("  .pend\n")
            }
            "%asmbinary" -> {
                val offset = if(stmt.args.size>1) ", ${stmt.args[1].int}" else ""
                val length = if(stmt.args.size>2) ", ${stmt.args[2].int}" else ""
                out("  .binary \"${stmt.args[0].str}\" $offset $length")
            }
            "%breakpoint" -> {
                val label = "_prog8_breakpoint_${breakpointLabels.size+1}"
                breakpointLabels.add(label)
                out("$label\tnop")
            }
        }
    }

    private fun translate(jmp: Jump) {
        out("  jmp  ${getJumpTarget(jmp)}")
    }

    private fun getJumpTarget(jmp: Jump): String {
        return when {
            jmp.identifier!=null -> {
                val target = jmp.identifier.targetStatement(program.namespace)
                val asmName = asmSymbolName(jmp.identifier)
                if(target is Label)
                    "_$asmName"  // prefix with underscore to jump to local label
                else
                    asmName
            }
            jmp.generatedLabel!=null -> jmp.generatedLabel
            jmp.address!=null -> jmp.address.toHex()
            else -> "????"
        }
    }

    private fun translate(ret: Return) {
        ret.value?.let { expressionsAsmGen.translateExpression(it) }
        out("  rts")
    }

    private fun translate(asm: InlineAssembly) {
        val assembly = asm.assembly.trimEnd().trimStart('\n')
        assemblyLines.add(assembly)
    }
}
