import { convertIR } from "./ir/convert"
import { prettyPrintBlock } from "./ir/pretty"
import { parse } from "./parser"
import { allocateRegisters } from "./regalloc"
import { generateBlock } from "./sm83/codegen"
import { R8 } from "./sm83/cpu"

import type { Attr, Decl, Stmt, Type, VarDecl } from "./ast"

type SymbolDefn = {
    attrs: Array<Attr>,
    name: string,
    type: Type
}

type SymbolMap = {
    [name: string]: SymbolDefn
}

// TODO: Support more than one TU
export const compile = (fileName: string, source: string): string => {
    // 1. Parse
    const ast = parse(source)

    // 2. Partition declarations and statements
    const decls = ast.filter((x): x is Decl => x.type == "decl")
    const stmts = ast.filter((x): x is Stmt => x.type == "stmt")

    // 3. Create top-level symbol map
    const symbols = processDecls(decls)

    // TODO: Support declaring types other than U8/S8
    const asmDecls = Object.values(symbols).map(symbol => `${symbol.name}:: DB`)

    // TODO: Some form of type-checking

    // 4. Generate IR
    const ir = convertIR(stmts)

    console.log("=== IR === ")
    console.log(prettyPrintBlock(ir))

    // 5. Generate code
    const alloc = allocateRegisters(ir, Object.values(R8))
    const insns = generateBlock(alloc, ir)

    let output = ''
    if (asmDecls.length > 0) {
        output += `SECTION "${fileName} Data", WRAM0\n\n`
        output += asmDecls.join("\n")
        output += "\n\n"
    }

    if (insns.length > 0) {
        output += `SECTION "${fileName} Code", ROM0\n\n`
        output += insns.join("\n")
    }

    console.log("=== ASM ===")
    return output
}

export const processDecls = (decls: Array<Decl>): SymbolMap => decls.reduce(processDecl, {})

export const processDecl = (symbols: SymbolMap, decl: Decl): SymbolMap => {
    const symbol = processVarDecl(symbols, decl)

    // Don't declare extern symbols
    if (hasAttr(symbol, { name: 'extern', args: [] })) {
        return symbols
    }

    return {
        ...symbols,
        [symbol.name]: symbol,
    }
}

export const processVarDecl = (symbols: SymbolMap, decl: VarDecl): SymbolDefn => {
    if (typeof symbols[decl.name] !== 'undefined') {
        throw new Error(`a variable named \`${decl.name}' is already defined`)
    }

    return {
        attrs: decl.attrs,
        name: decl.name,
        type: decl.args.type,
    }
}

export const hasAttr = (symbol: SymbolDefn, attr: Attr) => {
    let has = false

    for (let i = 0; i < symbol.attrs.length; ++i) {
        has = attrEquals(symbol.attrs[i], attr)
        if (has) {
            break
        }
    }

    return has
}

const attrEquals = (attr1: Attr, attr2: Attr): boolean => {
    if (attr1.name !== attr2.name || attr1.args.length !== attr2.args.length) {
        return false
    }

    for (let i = 0; i < attr1.args.length; ++i) {
        if (attr1.args[i] !== attr2.args[i]) {
            return false
        }
    }

    return true
}