/**
 * @file Codegen.cpp
 * @brief 寄存器码生成（12.8，切片 1：MOVE / LOADK / RET）
 * @author 
 * @date 2026-08-21
 *
 * @details 设计说明（详见 Doc/compiler/寄存器码.md）：
 * - 切片 1 范围：PROGRAM 标量变量帧（r0 起按声明序）、字面量 LOADK、
 *   变量间 MOVE、RET；空 MAIN 出映像。其余构造报 codegen error（后续切片展开）
 * - 常量统一进常量表（BOOL/INT/TIME），const_id = 首次出现序
 * - 映像：头（cycle_limit/dt_ms/工程哈希/入口）+ 常量表 + 函数表 + 字节码，
 *   数据段为空（无全局/实例）
 *
 * 函数清单：
 * - put_le32 / put_le64      小端写入映像缓冲
 * - Builder::Builder        （构造）存工程/源文件/链接结果/输出
 * - Builder::run             逐 POU 建函数 → 拼映像
 * - Builder::fail            组装 "codegen error: <msg>" 返回 false
 * - build_frame              给 POU 标量变量分配帧寄存器（0 起按声明序）
 * - build_function           编译一个 POU 的语句体
 * - compile_stmt             语句编译（切片 1：仅赋值）
 * - compile_expr             表达式编译（切片 1：仅字面量 / 变量读）
 * - const_id                 取常量表 id（无则追加）
 * - assemble_image           拼头 + 常量表 + 函数表 + 字节码
 * - codegen_project          对外入口
 */

#include "compiler/Codegen.h"

#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 代码生成器（切片 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) {}

        /**
         * @brief 逐 POU 建函数后拼映像
         * @return true 成功；false（err 已写，前缀 "codegen error"）
         */
        bool run() {
            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 nregs = 0;                   // 变量数（切片 1 无临时）
        };

        /**
         * @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) {
                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++;
                }
            }
            for (const Stmt& st : pou.body) {
                if (!compile_stmt(*f, st)) {
                    return false;
                }
            }
            f->code.push_back(isa::enc_ret());
            return true;
        }

        /**
         * @brief 编译一条语句（切片 1：仅赋值）
         * @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 1");
            }
            const auto it = f.regs.find(st.target);
            if (it == f.regs.end()) {
                return fail("no register for '" + st.target + "'");
            }
            const uint8_t rd = it->second;
            uint8_t rs = 0;
            const ExprKind k = st.value->kind;
            if (k == ExprKind::LitBool || k == ExprKind::LitInt ||
                k == ExprKind::LitTime) {
                const isa::types::TypeTag tag =
                    k == ExprKind::LitBool ? isa::types::Bool
                    : k == ExprKind::LitInt ? isa::types::Int
                                            : isa::types::Time;
                f.code.push_back(
                    isa::enc_imm(isa::Op::LOADK, rd, const_id(tag, st.value->int_value)));
                return true;
            }
            if (k == ExprKind::VarRef) {
                const auto sit = f.regs.find(st.value->name);
                if (sit == f.regs.end()) {
                    return fail("no register for '" + st.value->name + "'");
                }
                f.code.push_back(isa::enc_rr(isa::Op::MOVE, rd, sit->second));
                return true;
            }
            return fail("expression not supported in slice 1");
        }

        /**
         * @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 拼映像：头 + 常量表 + 函数表 + 字节码（数据段空）
         * @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_end = off_code + code_total;

            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, 0);  // n_globals（切片 1 无全局）
            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_end);  // offset_fb（空）
            put_le32(b, 68, off_end);  // 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 + 4, 0);  // code_offset 相对段起点 = 前序函数长度和
                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;
            }
            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_;
    };

}  // 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