const HEADER = [
    0x00, 0x61, 0x73, 0x6d,
    0x01, 0x00, 0x00, 0x00
];

const LINE_END = "\n"

class Tokenizer {
	constructor() {
		this.delims = new Set([" ", "\r", "\n", "\t"]);
		this.skips = new Set([" ", "\r", "\t"]);
        this.comment_start = ";"
		this.buffer = [];
        this.comment = false;
	}

	skip() {
		const idx = this.buffer.findIndex((cp) => !this.skips.has(cp));
        this.buffer = idx == -1 ? [] : this.buffer.slice(idx);
    }

	next() {
        this.skip();
        if (this.buffer[0] == LINE_END)
            return this.buffer.shift();

		const idx = this.buffer.findIndex((cp) => this.delims.has(cp));
        if (idx != -1) {
            const token = this.buffer.slice(0, idx).join("");
            this.buffer = this.buffer.slice(idx);
            return token;
        }
	}

    *handle(src) {
        this.buffer.push(...src);
        let token;
        while (token = this.next()) {
            if (token == this.comment_start)
                this.comment = true;
            else if (this.comment && token == LINE_END)
                this.comment = false;
            else if (!this.comment)
                yield token;
        }
	}
}

const State = Object.freeze({
    TOP:    0,
    EXPORT: 1,
    FUNC:   2,
    RESULT: 3,
});

const Action = Object.freeze({
    APPEND: 0,
    EXPORT: 1,
    FUNC:   2,
    RESULT: 3,
});

const types = {
    "func": 0x60,
    "i32":  0x7f,
};

const opcodes = {
    "end":       0x0b,
    "i32.const": 0x41,
    "i32.mul":   0x6c,
};

class Parser {
	constructor() {
		this.tokens = [];
		this.tokenizer = new Tokenizer();
        this.state = State.TOP;
        this.directives = {
            ".export": State.EXPORT,
            ".func":   State.FUNC,
            ".result": State.RESULT,
        };
        this.handlers = {
            [State.TOP]:    (token) => this.token_top(token),
            [State.EXPORT]: (token) => this.token_export(token),
            [State.FUNC]:   (token) => this.token_func(token),
            [State.RESULT]: (token) => this.token_result(token),
        };

        this.results = [];
	}

    translate_code(token) {
        return opcodes[token] ?? parseInt(token);
    }

    translate_type(token) {
        return types[token];
    }

    token_top(token) {
        if (token == LINE_END)
            return;
        let state;
        if (state = this.directives[token]) {
            this.state = state;
            return;
        }
        const code = this.translate_code(token);
        return { type: Action.APPEND, code };
    }

    token_export(token) {
        this.state = State.TOP;
        return { type: Action.EXPORT, name: token };
    }

    token_func(token) {
        this.state = State.TOP;
        return { type: Action.FUNC, name: token };
    }

    token_result(token) {
        if (token == LINE_END) {
            const action = { type: Action.RESULT, types: this.results };
            this.state = State.TOP;
            this.results = [];
            return action;
        } else {
            this.results.push(this.translate_type(token));
        }
    }

	*handle(src) {
        let action;
		for (const token of this.tokenizer.handle(src)) {
            if (action = this.handlers[this.state](token))
                yield action;
        }
	}
}

const Section = Object.freeze({
    TYPE:   0x01,
    FUNC:   0x03,
    EXPORT: 0x07,
    CODE:   0x0a
});

export class Assembler {
	constructor() {
        this.encoder = new TextEncoder("utf-8");
		this.decoder = new TextDecoder("utf-8");
		this.parser = new Parser();
        this.handlers = {
            [Action.APPEND]: (action) => this.action_append(action),
            [Action.EXPORT]: (action) => this.action_export(action),
            [Action.FUNC]:   (action) => this.action_func(action),
            [Action.RESULT]: (action) => this.action_result(action),
        };

        this.exports = [];
        this.funcs = {}
	}

    action_append(action) {
        this.funcs[this.current_func].body.push(action.code);
    }

    action_export(action) {
        const index = Object.keys(this.funcs).indexOf(action.name);
        this.exports[action.name] = { kind: 0x00, index };
    }

    action_func(action) {
        this.funcs[action.name] = {
            params: {},
            results: [],
            locals: {},
            body: [],
        }
        this.current_func = action.name;
    }

    action_result(action) {
        this.funcs[this.current_func].results.push(...action.types);
    }

	push(chunk) {
		const text = this.decoder.decode(chunk, { stream: true });
		for (const action of this.parser.handle(text))
            this.handlers[action.type](action);
	}

    wasm_section_type() {
        const funcs = Object.values(this.funcs);
        const contents = funcs.map(({ params, results }) => {
            const param_types = Object.values(params);
            return [
                types["func"],
                param_types.length,
                ...param_types,
                results.length,
                ...results,
            ];
        });
        return [ contents.length ].concat(...contents);
    }

    wasm_section_func() {
        const func_count = Object.entries(this.funcs).length;
        return [ func_count, ...Array(func_count).keys() ];
    }

    wasm_section_export() {
        const exports = Object.entries(this.exports);
        const contents = exports.map(([ name, { kind, index }]) => {
            const name_utf8 = this.encoder.encode(name);
            return [
                name_utf8.length,
                ...name_utf8,
                kind,
                index,
            ];
        });
        return [ exports.length ].concat(...contents);
    }

    wasm_section_code() {
        const funcs = Object.values(this.funcs);
        const contents = funcs.map(({ body, locals }) => {
            const local_count = Object.entries(locals).length;
            return [
                body.length + 2,
                local_count,
                ...body,
                opcodes["end"]
            ]
        });
        return [ contents.length ].concat(...contents);
    }

	wasm() {
        const template = [
            [ Section.TYPE,   () => this.wasm_section_type() ],
            [ Section.FUNC,   () => this.wasm_section_func() ],
            [ Section.EXPORT, () => this.wasm_section_export() ],
            [ Section.CODE,   () => this.wasm_section_code() ],
        ];
        const sections = template.map(([ code, generator ]) => {
            const body = generator();
            return [ code, body.length, ...body ];
        });

		return new Uint8Array(HEADER.concat(...sections));
	}
}