#include <cassert>
#include <map>
#include <string>
#include <vector>
#include "MNN_generated.h"
#include "PluginModule.hpp"


class SourceTarget {
public:
    SourceTarget() {}
    ~SourceTarget() {}
    void addIndent() { indent++; }
    void subIndent() { indent--; }
    std::string getIndent() {
        return std::string(4 * indent, ' ');
    }
private:
    int indent = 0;
};

class CTarget : public SourceTarget {
public:
    CTarget(std::string& name) {}
    ~CTarget() {}
};

namespace AST {

class Expr {
public:
    virtual ~Expr() = default;
    virtual std::string codegen(SourceTarget* target) = 0;
private:
    friend class PluginModule;
};

class NumberExpr : public Expr {
private:
    union Val {
        char chVal;
        float f32Val;
        double f64Val;
        int8_t  i8Val;
        int16_t i16Val;
        int32_t i32Val;
        int64_t i64Val;
        uint8_t  ui8Val;
        uint16_t ui16Val;
        uint32_t ui32Val;
        uint64_t ui64Val;
    } mVal;
    enum DataType {
        CHAR = 0,
        FP16,
        FP32,
        FP64,
        INT1,
        INT8,
        INT16,
        INT32,
        INT64,
        UINT1,
        UINT8,
        UINT16,
        UINT32,
        UINT64
    };
    DataType mType;
public:
    NumberExpr(float Val) : mType(FP32) { mVal.f32Val = Val; }
    NumberExpr(double Val) : mType(FP64) { mVal.f64Val = Val; }
    NumberExpr(int8_t Val) : mType(INT8) { mVal.i8Val = Val; }
    NumberExpr(int16_t Val) : mType(INT16) { mVal.i16Val = Val; }
    NumberExpr(int32_t Val) : mType(INT32) { mVal.i32Val = Val; }
    NumberExpr(int64_t Val) : mType(INT64) { mVal.i64Val = Val; }
    NumberExpr(uint8_t Val) : mType(UINT8) { mVal.ui8Val = Val; }
    NumberExpr(uint16_t Val) : mType(UINT16) { mVal.ui16Val = Val;}
    NumberExpr(uint32_t Val) : mType(UINT32) { mVal.ui32Val = Val;}
    NumberExpr(uint64_t Val) : mType(UINT64) { mVal.ui64Val = Val;}
    CODEGEN_FUNCS
};

class VariableExpr : public Expr {
    std::string Name;

public:
    VariableExpr() = default;
    VariableExpr(const std::string &Name) : Name(Name) {}
    const std::string &getName() const { return Name; }
    std::string codegen(SourceTarget* target) override;
};

class SubscriptExpr : public VariableExpr {
    std::unique_ptr<Expr> Base, Offset;
public:
    SubscriptExpr(std::unique_ptr<Expr> Base, std::unique_ptr<Expr> Offset)
        : Base(std::move(Base)), Offset(std::move(Offset)) {}
    SubscriptExpr(std::unique_ptr<Expr> Base, const std::string& Offset)
        : Base(std::move(Base)), Offset(std::make_unique<VariableExpr>(Offset)) {}
    SubscriptExpr(std::unique_ptr<Expr> Base, int Offset)
        : Base(std::move(Base)), Offset(std::make_unique<NumberExpr>(Offset)) {}
    SubscriptExpr(const std::string& Base, const std::string& Offset)
        : Base(std::make_unique<VariableExpr>(Base)), Offset(std::make_unique<VariableExpr>(Offset)) {}
    SubscriptExpr(const std::string& Base, int Offset)
        : Base(std::make_unique<VariableExpr>(Base)), Offset(std::make_unique<NumberExpr>(Offset)) {}
    SubscriptExpr(const std::string& Base, std::unique_ptr<Expr> Offset)
        : Base(std::make_unique<VariableExpr>(Base)), Offset(std::move(Offset)) {}
    std::string codegen(SourceTarget* target) override;
};

class UnaryExpr : public Expr {
    MNN::UnaryOpOperation Op;
    std::unique_ptr<Expr> Operand;
public:
    UnaryExpr(MNN::UnaryOpOperation Op, std::unique_ptr<Expr> Operand)
        : Op(Op), Operand(std::move(Operand)) {}
    std::string codegen(SourceTarget* target) override;
};

class ReluExpr : public Expr {
    float minVal, maxVal;
    std::unique_ptr<Expr> Operand;
public:
    ReluExpr(float minVal, float maxVal, std::unique_ptr<Expr> Operand)
        : minVal(minVal), maxVal(maxVal), Operand(std::move(Operand)) {}
    std::string codegen(SourceTarget* target) override;
};

class BinaryExpr : public Expr {
    MNN::BinaryOpOperation Op;
    std::unique_ptr<Expr> LHS, RHS;
public:
    BinaryExpr(MNN::BinaryOpOperation Op, std::unique_ptr<Expr> LHS,
                    std::unique_ptr<Expr> RHS)
          : Op(Op), LHS(std::move(LHS)), RHS(std::move(RHS)) {}
    std::string codegen(SourceTarget* target) override;
};

class AssignExpr : public Expr {
    std::unique_ptr<Expr> LHS, RHS;
public:
    AssignExpr(std::unique_ptr<Expr> LHS, std::unique_ptr<Expr> RHS)
             : LHS(std::move(LHS)), RHS(std::move(RHS)) {}
    std::string codegen(SourceTarget* target) override;
};

class CallExpr : public Expr {
    std::string Callee;
    std::vector<std::unique_ptr<Expr>> Args;
public:
    CallExpr(const std::string &Callee,
                std::vector<std::unique_ptr<Expr>> Args)
        : Callee(Callee), Args(std::move(Args)) {}
    std::string codegen(SourceTarget* target) override;
};

class IfExpr : public Expr {
    std::unique_ptr<Expr> Cond, Then, Else;
public:
    IfExpr(std::unique_ptr<Expr> Cond, std::unique_ptr<Expr> Then,
              std::unique_ptr<Expr> Else)
        : Cond(std::move(Cond)), Then(std::move(Then)), Else(std::move(Else)) {}
    std::string codegen(SourceTarget* target) override;
};

class LoopExpr : public Expr {
    std::string VarName;
    std::unique_ptr<Expr> Start, End, Step, Body;
public:
    LoopExpr(const std::string &VarName, std::unique_ptr<Expr> Start,
             std::unique_ptr<Expr> End, std::unique_ptr<Expr> Step,
             std::unique_ptr<Expr> Body)
           : VarName(VarName), Start(std::move(Start)), End(std::move(End)),
             Step(std::move(Step)), Body(std::move(Body)) {}
    std::string codegen(SourceTarget* target) override;
};

class VarExpr : public Expr {
    std::vector<std::pair<std::string, std::unique_ptr<Expr>>> VarNames;
    std::unique_ptr<Expr> Body;
public:
    VarExpr(std::vector<std::pair<std::string, std::unique_ptr<Expr>>> VarNames,
            std::unique_ptr<Expr> Body)
          : VarNames(std::move(VarNames)), Body(std::move(Body)) {}
    std::string codegen(SourceTarget* target) override;
};

class ListExpr : public Expr {
    std::vector<std::unique_ptr<Expr>> exprs;
public:
    ListExpr() = default;
    ListExpr(std::vector<std::unique_ptr<Expr>> exprs)
        : exprs(std::move(exprs)) {}
    void push_back(std::unique_ptr<Expr> expr) {
        exprs.emplace_back(std::move(expr));
    }
    std::string codegen(SourceTarget* target) override;
};

class PrototypeAST {
    std::string Name;
    int inputArgNum, outputArgNum;
public:
    PrototypeAST(const std::string &Name, int inputNum, int outputNum)
        : Name(Name), inputArgNum(inputNum), outputArgNum(outputNum) {}

    const std::string &getName() const { return Name; }
    std::string codegen(SourceTarget* target);
};

class FunctionAST {
    std::unique_ptr<PrototypeAST> Proto;
    std::unique_ptr<Expr> Body;
public:
    FunctionAST(std::unique_ptr<PrototypeAST> Proto,
                std::unique_ptr<Expr> Body)
        : Proto(std::move(Proto)), Body(std::move(Body)) {}
    std::string codegen(SourceTarget* target);
};
} // end AST namespace