/**
 * @file Codegen.cpp
 * @brief 寄存器码生成（12.8，切片 2：+ 全局 / I/Q/M）
 * @author 
 * @date 2026-08-21
 *
 * @details 设计说明（详见 Doc/compiler/寄存器码.md）：
 * - 切片 1：PROGRAM 标量变量帧（r0 起按声明序）、字面量 LOADK、
 *   变量间 MOVE、RET；空 MAIN 出映像。
 * - 切片 2：+ 全局数据区（声明序、对齐 BOOL 1B / INT 2B / TIME 8B、
 *   初值入数据段）；读全局 LOAD_GLOBAL / io.input 绑定用 LOAD_I；
 *   写全局 STORE_GLOBAL / io.output 绑定用 STORE_Q；写 io.input → 拒绝。
 *   其余构造报 codegen error（后续切片展开）
 * - 常量统一进常量表（BOOL/INT/TIME），const_id = 首次出现序
 * - 指令 slot = 数据区字节偏移；符号表 address 保持逻辑槽号（12.6 冻结）
 *
 * 函数清单：
 * - put_le32 / put_le64      小端写入映像缓冲
 * - align_up                按 2 的幂宽度向上对齐
 * - width_of                类型名 → 槽宽（bool 1 / int 2 / time 8）
 * - Builder::Builder        （构造）存工程/源文件/链接结果/输出，收集 io 绑定分类
 * - Builder::run            逐 POU 建函数 → 布局数据区 → 拼映像
 * - Builder::fail           组装 "codegen error: <msg>" 返回 false
 * - find_pou                按名查 POU AST
 * - build_function          编译一个 POU 的语句体（帧寄存器分配）
 * - compile_stmt            语句编译（切片 2：仅赋值，左值可帧寄存器或全局）
 * - compile_expr            表达式编译到寄存器（字面量 / 变量读，含全局 LOAD_*）
 * - store_target            赋值左值：帧寄存器直写或 STORE_* 到全局槽
 * - load_op / store_op      按 io 绑定选读取/写入操作码
 * - global_offset           逻辑槽号 → 数据区字节偏移
 * - layout_data             布局全局数据区（对齐 + 初值字节）
 * - init_of                 查全局声明的初值（AST）
 * - const_id                取常量表 id（无则追加）
 * - assemble_image          拼头 + 常量表 + 函数表 + 字节码 + 数据段
 * - codegen_project         对外入口
 */

#include "compiler/Codegen.h"

#include <cctype>
#include <cstdio>
#include <map>
#include <string>
#include <vector>

#include "isa/Encode.h"
#include "isa/Image.h"
#include "isa/Instr.h"
#include "isa/Types.h"

namespace compiler {
namespace {

    // ---- 小端写入（Image.cpp 内部实现不可见，这里自带最小版）----

    void put_le32(std::vector<uint8_t>& b, size_t off, uint32_t v) {
        b[off + 0] = static_cast<uint8_t>(v & 0xFFu);
        b[off + 1] = static_cast<uint8_t>((v >> 8) & 0xFFu);
        b[off + 2] = static_cast<uint8_t>((v >> 16) & 0xFFu);
        b[off + 3] = static_cast<uint8_t>((v >> 24) & 0xFFu);
    }

    void put_le64(std::vector<uint8_t>& b, size_t off, uint64_t v) {
        for (int i = 0; i < 8; ++i) {
            b[off + i] = static_cast<uint8_t>((v >> (8 * i)) & 0xFFu);
        }
    }

    /**
     * @brief 按 2 的幂宽度向上对齐
     * @param v 当前偏移
     * @param w 对齐宽度（1/2/8）
     * @return 对齐后的偏移
     */
    uint32_t align_up(uint32_t v, uint32_t w) {
        return (v + w - 1) & ~(w - 1);
    }

    /**
     * @brief 类型名 → 槽宽（对齐规则见 Doc/isa/指令与映像.md）
     * @param name "bool" / "int" / "time"
     * @return 1 / 2 / 8；未知返回 1
     */
    uint32_t width_of(const std::string& name) {
        if (name == "int") return 2;
        if (name == "time") return 8;
        return 1;  // bool 及未知
    }

    /**
     * @brief 代码生成器（切片 1）
     */
    class Builder {
    public:
        /**
         * @brief 构造生成器
         * @param proj 工程定义（cycle_limit / dt_ms / 哈希用）
         * @param units 全部源文件的 AST
         * @param link 链接结果（POU 顺序 / 符号）
         * @param image 输出映像字节
         * @param err 错误输出；可为 nullptr（静默）
         */
        Builder(const Project& proj, const std::vector<SourceUnit>& units,
                const LinkResult& link, std::vector<uint8_t>* image, std::string* err)
            : proj_(proj), units_(units), link_(link), image_(image), err_(err) {
            // io 绑定分类（名已折小写；不创造变量，只影响操作码选择）
            for (const isa::IoBinding& b : proj_.io) {
                std::string key = b.var;
                for (char& ch : key) {
                    ch = static_cast<char>(std::tolower(static_cast<unsigned char>(ch)));
                }
                if (b.is_input) {
                    io_input_[key] = true;
                } else {
                    io_output_[key] = true;
                }
            }
        }

        /**
         * @brief 布局数据区 → 逐 POU 建函数 → 拼映像
         * @details 数据区偏移先定（函数编译引用 global_offset）
         * @return true 成功；false（err 已写，前缀 "codegen error"）
         */
        bool run() {
            if (!layout_data()) {
                return false;
            }
            for (const LinkResult::PouScope& sc : link_.scopes) {
                const POU* pou = find_pou(sc.name);
                if (pou == nullptr) {
                    continue;
                }
                FuncCtx f;
                f.name = sc.name;
                if (!build_function(*pou, &f)) {
                    return false;
                }
                funcs_.push_back(std::move(f));
            }
            return assemble_image();
        }

    private:
        /**
         * @brief 组装错误消息并返回 false
         * @param msg 错误描述（不含前缀）
         * @return 恒 false
         */
        bool fail(const std::string& msg) {
            if (err_) {
                *err_ = "codegen error: " + msg;
            }
            return false;
        }

        /**
         * @brief 按名查 POU AST
         * @param name POU 名（小写）
         * @return POU 指针；未找到返回 nullptr
         */
        const POU* find_pou(const std::string& name) const {
            for (const SourceUnit& u : units_) {
                for (const POU& p : u.ast.pous) {
                    if (p.name == name) {
                        return &p;
                    }
                }
            }
            return nullptr;
        }

        // 每函数的编译态
        struct FuncCtx {
            std::string name;
            std::vector<isa::Instr> code;        // 字节码（函数表 code_offset 相对此段）
            std::map<std::string, uint8_t> regs;  // 变量名 → 帧寄存器
            uint8_t nlocals = 0;                 // 变量数 = 临时寄存器起始
            uint8_t nregs = 0;                   // 峰值（变量 + 临时）
            uint8_t temp_used = 0;               // 本语句已用临时数（语句结束清零）
        };

        /**
         * @brief 语句开始：临时寄存器基址复用
         * @param f 当前函数
         */
        void begin_stmt(FuncCtx& f) { f.temp_used = 0; }

        /**
         * @brief 分配一个临时寄存器（语句内递增）
         * @param f 当前函数
         * @return 临时寄存器号（可能更新 nregs 峰值）
         */
        uint8_t alloc_temp(FuncCtx& f) {
            const uint8_t r = f.nlocals + f.temp_used;
            ++f.temp_used;
            if (static_cast<uint16_t>(f.nlocals) + f.temp_used > f.nregs) {
                f.nregs = f.nlocals + f.temp_used;
            }
            return r;
        }

        /**
         * @brief 编译一个 POU
         * @details 切片 1：仅 PROGRAM + 标量变量（FB 实例/函数留后续切片）
         * @param pou POU AST
         * @param f 输出函数编译态
         * @return true 成功；false（err 已写）
         */
        bool build_function(const POU& pou, FuncCtx* f) {
            if (pou.kind != PouKind::Program) {
                return fail("only PROGRAM supported in slice 1 ('" + pou.name + "')");
            }
            for (const VarBlock& b : pou.blocks) {
                // External / Global 走数据区槽，不占帧寄存器
                if (b.section == VarSection::External || b.section == VarSection::Global) {
                    continue;
                }
                for (const VarDecl& d : b.vars) {
                    if (d.type.kind == TypeKind::FbUser || d.type.kind == TypeKind::FbBuiltin) {
                        return fail("FB instance not supported in slice 1 ('" + d.name + "')");
                    }
                    if (f->regs.size() > 255) {
                        return fail("register overflow in function '" + pou.name + "'");
                    }
                    f->regs[d.name] = f->nregs++;
                }
            }
            f->nlocals = f->nregs;
            for (const Stmt& st : pou.body) {
                begin_stmt(*f);
                if (!compile_stmt(*f, st)) {
                    return false;
                }
            }
            f->code.push_back(isa::enc_ret());
            return true;
        }

        /**
         * @brief 编译一条语句（切片 2：仅赋值，左值可为帧寄存器或全局）
         * @param f 当前函数
         * @param st 语句 AST
         * @return true 成功；false（err 已写）
         */
        bool compile_stmt(FuncCtx& f, const Stmt& st) {
            if (st.kind != StmtKind::Assign) {
                return fail("statement not supported in slice 2");
            }
            const auto it = f.regs.find(st.target);
            if (it == f.regs.end()) {
                // 全局 / 外部左值 → STORE_* 到数据区
                return store_target(f, st.target, *st.value);
            }
            const uint8_t rd = it->second;
            return compile_expr(f, *st.value, rd);
        }

        /**
         * @brief 赋值左值：帧寄存器直写或 STORE_* 到全局槽
         * @param f 当前函数
         * @param target 左值名（全局 / 外部）
         * @param value 右值表达式
         * @return true 成功；false（err 已写）
         */
        bool store_target(FuncCtx& f, const std::string& target, const Expr& value) {
            const auto git = link_.global_index.find(target);
            if (git == link_.global_index.end()) {
                return fail("no storage for '" + target + "'");
            }
            if (io_input_.count(target)) {
                return fail("cannot write to input '" + target + "'");
            }
            const uint8_t tmp = alloc_temp(f);
            if (!compile_expr(f, value, tmp)) {
                return false;
            }
            const uint16_t slot = global_offset(git->second);
            f.code.push_back(isa::enc_slot(store_op(target), tmp, slot));
            return true;
        }

        /**
         * @brief 表达式编译到目标寄存器
         * @details 字面量 → LOADK；帧变量 → MOVE；
         * 全局/外部 → LOAD_GLOBAL（io.input 绑定用 LOAD_I）
         * @param f 当前函数
         * @param e 表达式 AST
         * @param rd 目标寄存器
         * @return true 成功；false（err 已写）
         */
        bool compile_expr(FuncCtx& f, const Expr& e, uint8_t rd) {
            if (e.kind == ExprKind::LitBool || e.kind == ExprKind::LitInt ||
                e.kind == ExprKind::LitTime) {
                const isa::types::TypeTag tag =
                    e.kind == ExprKind::LitBool ? isa::types::Bool
                    : e.kind == ExprKind::LitInt ? isa::types::Int
                                                 : isa::types::Time;
                f.code.push_back(
                    isa::enc_imm(isa::Op::LOADK, rd, const_id(tag, e.int_value)));
                return true;
            }
            if (e.kind == ExprKind::VarRef) {
                const auto sit = f.regs.find(e.name);
                if (sit != f.regs.end()) {
                    f.code.push_back(isa::enc_rr(isa::Op::MOVE, rd, sit->second));
                    return true;
                }
                const auto git = link_.global_index.find(e.name);
                if (git != link_.global_index.end()) {
                    const uint16_t slot = global_offset(git->second);
                    f.code.push_back(isa::enc_slot(load_op(e.name), rd, slot));
                    return true;
                }
                return fail("no register or slot for '" + e.name + "'");
            }
            return fail("expression not supported in slice 2");
        }

        /**
         * @brief 按 io 绑定选读取操作码
         * @param name 变量名（小写）
         * @return io.input 绑定 → LOAD_I；否则 LOAD_GLOBAL
         */
        isa::Op load_op(const std::string& name) const {
            return io_input_.count(name) ? isa::Op::LOAD_I : isa::Op::LOAD_GLOBAL;
        }

        /**
         * @brief 按 io 绑定选写入操作码
         * @param name 变量名（小写）
         * @return io.output 绑定 → STORE_Q；否则 STORE_GLOBAL
         */
        isa::Op store_op(const std::string& name) const {
            return io_output_.count(name) ? isa::Op::STORE_Q : isa::Op::STORE_GLOBAL;
        }

        /**
         * @brief 逻辑槽号 → 数据区字节偏移（layout_data 先行）
         * @param slot 全局逻辑槽号（符号表 address）
         * @return 数据区字节偏移（u16 可容）
         */
        uint16_t global_offset(uint32_t slot) const {
            return static_cast<uint16_t>(global_offsets_[slot]);
        }

        /**
         * @brief 布局全局数据区：对齐 + 初值字节
         * @details 槽序 = 声明序（12.6）；对齐规则 BOOL 1B / INT 2B / TIME 8B，
         * 起点对齐自身宽度；初值取自 GVL 声明的 has_init（无则 0）
         * @return true 成功；false（槽溢出等，err 已写）
         */
        bool layout_data() {
            uint32_t cur = 0;
            for (const Symbol& s : link_.globals) {
                const uint32_t w = width_of(s.type_name);
                cur = align_up(cur, w);
                if (cur + w > 0xFFFF) {
                    return fail("data area exceeds slot range");
                }
                global_offsets_.push_back(cur);
                bool has_init = false;
                int64_t init = 0;
                init_of(s.name, &has_init, &init);
                data_.resize(cur + w, 0);
                if (s.type_name == "bool") {
                    data_[cur] = static_cast<uint8_t>(has_init ? init : 0);
                } else if (s.type_name == "int") {
                    const uint16_t v = static_cast<uint16_t>(has_init ? init : 0);
                    data_[cur] = static_cast<uint8_t>(v & 0xFFu);
                    data_[cur + 1] = static_cast<uint8_t>((v >> 8) & 0xFFu);
                } else if (s.type_name == "time") {
                    const uint64_t v = static_cast<uint64_t>(has_init ? init : 0);
                    for (int i = 0; i < 8; ++i) {
                        data_[cur + i] = static_cast<uint8_t>((v >> (8 * i)) & 0xFFu);
                    }
                }
                cur += w;
            }
            return true;
        }

        /**
         * @brief 查全局声明的初值（AST；gvl 文件顶层段，声明序 = 槽号序）
         * @param name 全局名（小写）
         * @param has_init 输出是否有初值
         * @param init 输出初值
         */
        void init_of(const std::string& name, bool* has_init, int64_t* init) const {
            for (const SourceUnit& u : units_) {
                for (const VarBlock& b : u.ast.globals) {
                    for (const VarDecl& d : b.vars) {
                        if (d.name == name) {
                            *has_init = d.has_init;
                            *init = d.init_value;
                            return;
                        }
                    }
                }
            }
            *has_init = false;
            *init = 0;
        }

        /**
         * @brief 取常量表 id（无则追加）
         * @param tag 类型标记（BOOL/INT/TIME）
         * @param value 常量值
         * @return const_id（u16）
         */
        uint16_t const_id(isa::types::TypeTag tag, int64_t value) {
            for (size_t i = 0; i < consts_.size(); ++i) {
                if (consts_[i].tag == tag && consts_[i].value == static_cast<uint64_t>(value)) {
                    return static_cast<uint16_t>(i);
                }
            }
            consts_.push_back({tag, static_cast<uint64_t>(value)});
            return static_cast<uint16_t>(consts_.size() - 1);
        }

        /**
         * @brief 拼映像：头 + 常量表 + 函数表 + 字节码 + 数据段
         * @details 段序：const → funcs → code → fb（空）→ data；
         * 数据段放全局初值（layout_data 已生成字节）
         * @return true 成功；false（err 已写）
         */
        bool assemble_image() {
            const uint32_t off_const = isa::kHeaderSize;
            const uint32_t off_funcs = off_const +
                static_cast<uint32_t>(consts_.size()) * isa::kConstEntrySize;
            uint32_t off_code = off_funcs +
                static_cast<uint32_t>(funcs_.size()) * isa::kFuncRowSize;
            uint32_t code_total = 0;
            for (const FuncCtx& f : funcs_) {
                code_total += static_cast<uint32_t>(f.code.size()) * 4;
            }
            const uint32_t off_data = off_code + code_total;
            const uint32_t off_end = off_data + static_cast<uint32_t>(data_.size());

            std::vector<uint8_t>& b = *image_;
            b.assign(off_end, 0);
            put_le32(b, 0, isa::kMagic);
            put_le32(b, 4, isa::kVersion);
            put_le32(b, 8, proj_.cycle_limit);
            put_le32(b, 12, proj_.dt_ms);

            uint64_t hash = isa::kFnvBasis;
            if (!compute_project_hash(proj_, &hash, err_)) {
                return false;
            }
            put_le64(b, 16, hash);

            uint32_t entry = 0;
            for (size_t i = 0; i < funcs_.size(); ++i) {
                if (funcs_[i].name == "main") {
                    entry = static_cast<uint32_t>(i);
                }
            }
            put_le32(b, 24, entry);
            put_le32(b, 28, static_cast<uint32_t>(link_.globals.size()));  // n_globals
            put_le32(b, 32, 0);  // n_i
            put_le32(b, 36, 0);  // n_q
            put_le32(b, 40, 0);  // n_m
            put_le32(b, 44, static_cast<uint32_t>(consts_.size()));
            put_le32(b, 48, static_cast<uint32_t>(funcs_.size()));
            put_le32(b, 52, off_const);
            put_le32(b, 56, off_funcs);
            put_le32(b, 60, off_code);
            put_le32(b, 64, off_data);  // offset_fb（空，与数据段起点相同）
            put_le32(b, 68, off_data);  // offset_data

            for (size_t i = 0; i < consts_.size(); ++i) {
                const size_t o = off_const + i * isa::kConstEntrySize;
                put_le32(b, o, static_cast<uint32_t>(consts_[i].tag));
                put_le64(b, o + 4, consts_[i].value);
            }

            uint32_t c = off_code;
            for (size_t i = 0; i < funcs_.size(); ++i) {
                const FuncCtx& f = funcs_[i];
                const size_t o = off_funcs + i * isa::kFuncRowSize;
                put_le32(b, o, f.nregs);
                put_le32(b, o + 8, static_cast<uint32_t>(f.code.size()));
                for (const isa::Instr in : f.code) {
                    put_le32(b, c, in);
                    c += 4;
                }
            }
            // code_offset 回填（相对字节码段起点）
            uint32_t acc = 0;
            for (size_t i = 0; i < funcs_.size(); ++i) {
                const size_t o = off_funcs + i * isa::kFuncRowSize;
                put_le32(b, o + 4, acc);
                acc += static_cast<uint32_t>(funcs_[i].code.size()) * 4;
            }

            // 数据段：全局初值字节
            for (size_t i = 0; i < data_.size(); ++i) {
                b[off_data + i] = data_[i];
            }
            return true;
        }

        // ---- 成员 ----
        const Project& proj_;
        const std::vector<SourceUnit>& units_;
        const LinkResult& link_;
        std::vector<uint8_t>* image_;
        std::string* err_;
        std::vector<FuncCtx> funcs_;
        std::vector<isa::ConstEntry> consts_;
        std::map<std::string, bool> io_input_;    // io.input 绑定名（小写）
        std::map<std::string, bool> io_output_;   // io.output 绑定名（小写）
        std::vector<uint32_t> global_offsets_;    // 逻辑槽号 → 数据区字节偏移
        std::vector<uint8_t> data_;               // 数据区（全局初值字节）
    };

}  // namespace

/**
 * @brief 编译工程为 .stb 映像字节（对外入口）
 * @param proj 工程定义
 * @param units 全部源文件的 AST
 * @param link 链接结果（在类型检查成功后调用）
 * @param image 输出映像字节
 * @param err 错误输出；可为 nullptr（静默）
 * @return true 成功；false 失败（err 前缀 "codegen error"）
 */
bool codegen_project(const Project& proj, const std::vector<SourceUnit>& units,
                     const LinkResult& link, std::vector<uint8_t>* image,
                     std::string* err) {
    Builder b(proj, units, link, image, err);
    return b.run();
}

}  // namespace compiler