#include <pproc_expand.h>
#include <scc_pproc.h>

static scc_lexer_tok_t stringify_argument(scc_lexer_tok_vec_t *arg_tokens) {
    scc_cstring_t str = scc_cstring_create();
    scc_cstring_append_ch(&str, '\"'); // 左引号

    int need_space = 0; // 是否需要插入空格
    for (usize i = 0; i < arg_tokens->size; i++) {
        scc_lexer_tok_t *tok = &scc_vec_at(*arg_tokens, i);
        if (tok->type == SCC_TOK_BLANK) {
            need_space = 1; // 标记遇到空白
            continue;
        }

        // 需要空格且当前不是第一个有效token，插入一个空格
        if (need_space && i > 0) {
            scc_cstring_append_ch(&str, ' ');
            need_space = 0;
        }

        // 对字符串/字符常量内的 " 和 \ 进行转义
        if (tok->type == SCC_TOK_STRING_LITERAL ||
            tok->type == SCC_TOK_CHAR_LITERAL) {
            // 注意：lex包含两端的引号，需要跳过首尾，转义内部字符
            // 简化：暂不处理内部转义，直接追加
        }
        scc_cstring_append(&str, &tok->lexeme);
        need_space = 0;
    }
    scc_cstring_append_ch(&str, '\"'); // 右引号

    scc_lexer_tok_t result;
    result.type = SCC_TOK_STRING_LITERAL;
    result.lexeme = str;
    return result;
}

static inline void scc_copy_expand(scc_pproc_expand_t *expand_ctx,
                                   scc_pproc_expand_t *copyed_ctx,
                                   scc_lexer_tok_ring_t *ring) {

    copyed_ctx->input = ring;
    copyed_ctx->expanded_set = expand_ctx->expanded_set;
    copyed_ctx->macro_table = expand_ctx->macro_table;
    copyed_ctx->need_rescan = false;
    scc_vec_init(copyed_ctx->output);
}

void scc_pproc_expand_by_src(scc_pproc_t *pp, const scc_pproc_macro_t *macro) {
    scc_pproc_expand_t ctx;
    scc_lexer_tok_vec_t expaned_buffer;
    scc_vec_init(expaned_buffer);

    int ok;
    scc_lexer_tok_t tok;
    scc_ring_next_consume(*pp->cur_ring, tok, ok);
    if (macro->type == SCC_PP_MACRO_NONE || ok == false) {
        UNREACHABLE();
    } else if (macro->type == SCC_PP_MACRO_OBJECT) {
        scc_vec_push(expaned_buffer, tok);
    } else if (macro->type == SCC_PP_MACRO_FUNCTION) {
        scc_vec_push(expaned_buffer, tok);
        scc_pproc_parse_macro_arguments(pp->cur_ring, &expaned_buffer, true);
    }

    scc_lexer_tok_ring_t ring = scc_lexer_array_to_ring(&expaned_buffer);
    ctx.input = &ring;
    ctx.macro_table = &pp->macro_table;
    ctx.need_rescan = false;
    scc_vec_init(ctx.output);

    scc_pproc_marco_table_init(&ctx.expanded_set);
    scc_pproc_expand_macro(&ctx);
    pp->expanded_ring = scc_lexer_array_to_ring(&ctx.output);
    scc_pproc_macro_table_drop(&ctx.expanded_set);
}

static inline void
split_arguments(scc_pproc_macro_extened_params_t *splited_params,
                scc_lexer_tok_vec_t *raw_args) {
    scc_lexer_tok_vec_t arg;
    scc_vec_init(arg);
    scc_vec_foreach(*raw_args, i) {
        scc_lexer_tok_t *raw_arg = &scc_vec_at(*raw_args, i);
        if (raw_arg->type == SCC_TOK_COMMA) {
            scc_lexer_tok_drop(raw_arg);
            scc_vec_push(*splited_params, arg);
            scc_vec_init(arg);
        } else {
            if (raw_arg->type == SCC_TOK_BLANK ||
                scc_get_tok_subtype(raw_arg->type) ==
                    SCC_TOK_SUBTYPE_EMPTYSPACE) {
                scc_lexer_tok_drop(raw_arg);
            }
            scc_vec_push(arg, *raw_arg);
        }
    }
    scc_vec_push(*splited_params, arg);
}

static inline void
expand_arguments(scc_pproc_macro_extened_params_t *expanded_params,
                 scc_pproc_macro_extened_params_t *splited_params,
                 scc_pproc_expand_t *expand_ctx) {
    scc_vec_foreach(*splited_params, i) {
        scc_pproc_expand_t ctx;
        scc_lexer_tok_vec_t splite_param = scc_vec_at(*splited_params, i);
        scc_lexer_tok_vec_t expanded_param;
        scc_vec_init(expanded_param);
        scc_vec_foreach(splite_param, j) {
            scc_lexer_tok_t tok = scc_vec_at(splite_param, j);
            tok.lexeme = scc_cstring_copy(&tok.lexeme);
            scc_vec_push(expanded_param, tok);
        }
        scc_lexer_tok_ring_t ring = scc_lexer_array_to_ring(&expanded_param);
        scc_copy_expand(expand_ctx, &ctx, &ring);
        scc_pproc_expand_macro(&ctx);
        scc_vec_push(*expanded_params, ctx.output);
    }
}

static inline void
expanded_params_free(scc_pproc_macro_extened_params_t *expanded_params) {
    scc_vec_foreach(*expanded_params, i) {
        scc_lexer_tok_vec_t expanded_param = scc_vec_at(*expanded_params, i);
        scc_vec_foreach(expanded_param, j) {
            scc_lexer_tok_t tok = scc_vec_at(expanded_param, j);
            scc_lexer_tok_drop(&tok);
        }
        scc_vec_free(expanded_param);
    }
    scc_vec_free(*expanded_params);
}

static inline void expand_function_macro(scc_pproc_expand_t *expand_ctx,
                                         const scc_pproc_macro_t *macro) {
    Assert(macro->type == SCC_PP_MACRO_FUNCTION);
    scc_lexer_tok_vec_t raw_args;
    scc_vec_init(raw_args);
    scc_pproc_parse_macro_arguments(expand_ctx->input, &raw_args, false);

    // collect, fill and expand arg
    scc_pproc_macro_extened_params_t splited_params;
    scc_vec_init(splited_params);
    split_arguments(&splited_params, &raw_args);

    scc_pproc_macro_extened_params_t expanded_params;
    scc_vec_init(expanded_params);
    expand_arguments(&expanded_params, &splited_params, expand_ctx);

    // replace
    scc_vec_foreach(macro->replaces, i) {
        scc_lexer_tok_t tok = scc_vec_at(macro->replaces, i);
        scc_lexer_tok_t prev_tok = {0};
        if (i >= 1) {
            prev_tok = scc_vec_at(macro->replaces, i - 1);
        }
        if (tok.type == SCC_TOK_BLANK) {
            tok.lexeme = scc_cstring_from_cstr(" ");
            scc_vec_push(expand_ctx->output, tok);
            continue;
        }
        scc_vec_foreach(macro->params, j) {
            if (scc_cstring_cmp(&tok.lexeme,
                                &(scc_vec_at(macro->params, j).lexeme)) == 0) {
                if (j >= scc_vec_size(expanded_params)) {
                    LOG_ERROR("Invalid macro parameter");
                    goto CONTINUE;
                }

                if (prev_tok.type == SCC_TOK_SHARP) {
                    // # stringify
                    scc_lexer_tok_t out =
                        stringify_argument(&scc_vec_at(splited_params, j));
                    scc_vec_pop(expand_ctx->output);
                    scc_vec_push(expand_ctx->output, out);
                } else {
                    scc_lexer_tok_vec_t expanded_param =
                        scc_vec_at(expanded_params, j);
                    scc_vec_foreach(expanded_param, k) {
                        tok = scc_vec_at(expanded_param, k);
                        tok.lexeme = scc_cstring_copy(&tok.lexeme);
                        scc_vec_push(expand_ctx->output, tok);
                    }
                }
                goto CONTINUE;
            }
        }
        tok.lexeme = scc_cstring_copy(&tok.lexeme);
        scc_vec_push(expand_ctx->output, tok);
    CONTINUE:
        continue;
    }
    expanded_params_free(&splited_params);
    expanded_params_free(&expanded_params);
}

static inline void expand_object_macro(scc_pproc_expand_t *expand_ctx,
                                       const scc_pproc_macro_t *macro) {
    scc_vec_foreach(macro->replaces, i) {
        scc_lexer_tok_t tok = scc_vec_at(macro->replaces, i);
        if (tok.type == SCC_TOK_BLANK) {
            tok.lexeme = scc_cstring_from_cstr(" ");
        } else {
            tok.lexeme = scc_cstring_copy(&tok.lexeme);
        }
        scc_vec_push(expand_ctx->output, tok);
    }
}

void scc_pproc_expand_macro(scc_pproc_expand_t *expand_ctx) {
    int ok;
    scc_lexer_tok_t tok;
    while (1) {
        scc_ring_next_consume(*expand_ctx->input, tok, ok);
        if (!ok) {
            return;
        }
        if (tok.type != SCC_TOK_IDENT) {
            scc_vec_push(expand_ctx->output, tok);
            continue;
        }
        // maybe expanded
        scc_pproc_macro_t *macro =
            scc_pproc_macro_table_get(expand_ctx->macro_table, &tok.lexeme);

        if (macro == null || scc_pproc_macro_table_get(
                                 &expand_ctx->expanded_set, &macro->name)) {
            scc_vec_push(expand_ctx->output, tok);
            continue;
        }
        expand_ctx->need_rescan = true;

        scc_pproc_macro_t *expanded_macro =
            scc_pproc_macro_new(&macro->name, macro->type);
        if (expanded_macro == null) {
            LOG_FATAL("Out of memory");
        }
        scc_pproc_macro_table_set(&expand_ctx->expanded_set, expanded_macro);
        if (macro->type == SCC_PP_MACRO_OBJECT) {
            expand_object_macro(expand_ctx, macro);
        } else if (macro->type == SCC_PP_MACRO_FUNCTION) {
            expand_function_macro(expand_ctx, macro);
        } else {
            UNREACHABLE();
        }
    RESCAN:
        scc_pproc_macro_table_remove(&expand_ctx->expanded_set, &macro->name);
        // TODO expand # and ##
        continue;
    }
    if (expand_ctx->need_rescan) {
        expand_ctx->need_rescan = false;
        scc_pproc_expand_t rescan_ctx;
        scc_lexer_tok_ring_t ring =
            scc_lexer_array_to_ring(&expand_ctx->output);
        scc_copy_expand(expand_ctx, &rescan_ctx, &ring);
        scc_pproc_expand_macro(&rescan_ctx);
        scc_ring_free(ring);
        expand_ctx->output = rescan_ctx.output;
    }
}