#include "cps/TypeSystem.h"
#include "cps/TypeAST.h"
#include "llvm/IR/Constants.h"
#include <algorithm>
#include <utility>

using namespace llvm;
using namespace cps;

bool TypeInfo::isC() const { return !CIdentity.empty(); }

bool TypeInfo::sameType(const TypeInfo *Other) const {
    if (!Other) return false;
    if (isC() || Other->isC())
        return isC() && Other->isC() && CIdentity == Other->CIdentity;
    return sameIdentifier(Name, Other->Name);
}

const RecordFieldInfo *RecordPayload::findField(const std::string &Name) const {
    for (const auto &F : Fields) {
        if (sameIdentifier(F.first, Name)) return &F.second;
    }
    return nullptr;
}

bool TypeInfo::isIntegral() const {
    return Kind == TypeKind::Integer || Kind == TypeKind::Char ||
        (Kind == TypeKind::CScalar && LLVMType->isIntegerTy());
}

bool TypeInfo::isReal() const {
    return Kind == TypeKind::Real || (Kind == TypeKind::CScalar && LLVMType->isFloatingPointTy());
}

bool TypeInfo::isBoolean() const {
    return Kind == TypeKind::Boolean || CIdentity == "_Bool";
}

bool TypeInfo::isString() const {
    return IsClass && sameIdentifier(Name, "STRING");
}

bool TypeInfo::isChar() const {
    return Kind == TypeKind::Char || CIdentity == "char";
}

bool TypeInfo::isVoid() const {
    return Kind == TypeKind::Void || (isC() && CIdentity == "void");
}

bool TypeInfo::isEnum() const {
    return Kind == TypeKind::Enum;
}

bool TypeInfo::isRecord() const {
    return Kind == TypeKind::Record;
}

bool TypeInfo::isPointer() const {
    return Kind == TypeKind::Pointer || Kind == TypeKind::CPointer;
}

bool TypeInfo::isUserKind() const {
    return Kind == TypeKind::Enum || Kind == TypeKind::Record || Kind == TypeKind::Pointer;
}

bool TypeInfo::isOutputtable() const {
    switch (Kind) {
    case TypeKind::Integer:
    case TypeKind::Real:
    case TypeKind::Boolean:
    case TypeKind::Char:
    case TypeKind::Enum: // prints its ordinal
        return true;
    case TypeKind::CScalar:
        return !isVoid();
    case TypeKind::Void:
    case TypeKind::Pointer:
    case TypeKind::CPointer:
        return false;
    case TypeKind::Record:
        return isString();
    }
    return false;
}

const RecordPayload *TypeInfo::asRecord() const {
    return Kind == TypeKind::Record ? &Record : nullptr;
}

const EnumPayload *TypeInfo::asEnum() const {
    return Kind == TypeKind::Enum ? &Enum : nullptr;
}

const PointerPayload *TypeInfo::asPointer() const {
    return isPointer() ? &Pointer : nullptr;
}

TypeSystem::TypeSystem(LLVMContext &Ctx) : Context(Ctx) {
    registerBuiltins();
}

bool TypeSystem::sameType(const std::string &Left, const std::string &Right) const {
    const TypeInfo *L = resolve(Left);
    return L && L->sameType(resolve(Right));
}

const TypeInfo *TypeSystem::resolve(const std::string &TypeName) const {
    auto It = Types.find(TypeName);
    if (It == Types.end()) return ResolveMiss ? ResolveMiss(TypeName) : nullptr;
    if (Visible && !Visible(TypeName)) return nullptr;
    return &It->second;
}

TypeInfo *TypeSystem::registerType(const std::string &TypeName,
                                   Type *LLVMType,
                                   TypeKind Kind) {
    TypeInfo Info;
    Info.Name = TypeName;
    Info.LLVMType = LLVMType;
    Info.Kind = Kind;

    auto Result = Types.insert_or_assign(TypeName, std::move(Info));
    return &Result.first->second;
}

bool TypeSystem::declareUserType(const std::string &Name, TypeKind Kind, std::string &Err, bool IsClass) {
    if (Types.count(Name)) {
        Err = "Type " + Name + " is already defined";
        return false;
    }
    if (EnumConstants.count(Name)) {
        Err = "'" + Name + "' is an enum value and cannot name a type";
        return false;
    }

    Type *LLVMType = nullptr;
    switch (Kind) {
    case TypeKind::Enum:
        LLVMType = Type::getInt64Ty(Context); // INTEGER representation; identity is the name
        break;
    case TypeKind::Pointer:
        LLVMType = PointerType::getUnqual(Context);
        break;
    case TypeKind::Record:
        // Opaque named struct; the body is set in pass B.
        LLVMType = StructType::create(Context, std::string(kRecordStructPrefix) + Name);
        break;
    default:
        Err = "Internal: declareUserType called with a builtin kind";
        return false;
    }

    registerType(Name, LLVMType, Kind)->IsClass = IsClass;
    return true;
}

bool TypeSystem::defineEnum(const std::string &Name,
                            const std::vector<std::string> &Values,
                            std::string &Err) {
    auto It = Types.find(Name);
    if (It == Types.end() || It->second.Kind != TypeKind::Enum) {
        Err = "Internal: defineEnum on undeclared enum " + Name;
        return false;
    }

    for (size_t I = 0; I < Values.size(); ++I) {
        const std::string &V = Values[I];
        for (size_t J = 0; J < I; ++J) {
            if (sameIdentifier(Values[J], V)) {
                Err = "Duplicate value '" + V + "' in enum " + Name;
                return false;
            }
        }
        if (Types.count(V)) {
            Err = "Enum value '" + V + "' in " + Name + " collides with a type name";
            return false;
        }
        if (auto CIt = EnumConstants.find(V); CIt != EnumConstants.end()) {
            Err = "Enum value '" + V + "' in " + Name + " is already a value of enum " +
                  CIt->second.TypeName;
            return false;
        }
    }

    It->second.Enum.ValueNames = Values;
    for (size_t I = 0; I < Values.size(); ++I) {
        EnumConstants[Values[I]] = {Name, static_cast<int64_t>(I) + kEnumOrdinalBase};
    }
    return true;
}

bool TypeSystem::defineRecord(const std::string &Name,
                              const std::vector<RecordFieldDecl> &Fields,
                              std::string &Err) {
    auto It = Types.find(Name);
    if (It == Types.end() || It->second.Kind != TypeKind::Record) {
        Err = "Internal: defineRecord on undeclared record " + Name;
        return false;
    }

    // A bad field is skipped so the record still gets a body; Err keeps the
    // first diagnostic and Invalid silences follow-up missing-field errors.
    std::vector<Type*> FieldTypes;
    RecordPayload Payload;
    auto Bad = [&](std::string Message) {
        if (!Payload.Invalid) Err = std::move(Message);
        Payload.Invalid = true;
    };
    for (const RecordFieldDecl &Field : Fields) {
        const std::string &FieldName = Field.Name;
        const std::string &FieldTypeName = Field.TypeName;
        if (Payload.findField(FieldName)) {
            Bad("Duplicate field '" + FieldName + "' in record " + Name);
            continue;
        }
        const TypeInfo *FieldInfo = resolve(FieldTypeName);
        if (!FieldInfo) {
            Bad("Unknown type " + FieldTypeName + " for field " + FieldName +
                " of record " + Name);
            continue;
        }
        if (FieldInfo->isVoid()) {
            Bad("Field " + FieldName + " of record " + Name + " cannot be VOID");
            continue;
        }
        // By-value record fields require a *completed* record (source order);
        // this also rules out self/mutual inclusion. Inclusion via a pointer
        // type is fine (opaque ptr, always complete).
        if (FieldInfo->isRecord() && llvm::cast<StructType>(FieldInfo->LLVMType)->isOpaque()) {
            Bad("Record type " + FieldTypeName + " must be fully defined before field " +
                FieldName + " of record " + Name + " (records cannot contain themselves)");
            continue;
        }

        RecordFieldInfo FI;
        FI.TypeName = FieldInfo->Name;
        FI.Index = static_cast<unsigned>(FieldTypes.size());
        FI.IsPublic = Field.IsPublic;
        Payload.Fields.emplace_back(FieldName, FI);
        FieldTypes.push_back(FieldInfo->LLVMType);
    }

    llvm::cast<StructType>(It->second.LLVMType)->setBody(FieldTypes);
    bool Ok = !Payload.Invalid;
    It->second.Record = std::move(Payload);
    return Ok;
}

bool TypeSystem::definePointer(const std::string &Name,
                               const std::string &PointeeName,
                               std::string &Err) {
    auto It = Types.find(Name);
    if (It == Types.end() || It->second.Kind != TypeKind::Pointer) {
        Err = "Internal: definePointer on undeclared pointer type " + Name;
        return false;
    }

    const TypeInfo *Pointee = resolve(PointeeName);
    if (!Pointee) {
        Err = "Unknown pointee type " + PointeeName + " in pointer type " + Name;
        return false;
    }
    if (Pointee->isVoid()) {
        Err = "Pointer type " + Name + " cannot point to VOID";
        return false;
    }

    It->second.Pointer.PointeeTypeName = Pointee->Name;
    return true;
}

namespace {

// Seen holds the records already being checked, so self-referential
// pointer types (linked lists) end the recursion.
bool checkCStruct(const TypeSystem &Types, const TypeInfo *Info, std::string *Why,
                  std::vector<const TypeInfo *> &Seen) {
    if (!Info || !Info->isRecord() || Info->IsClass) return false;
    if (std::find(Seen.begin(), Seen.end(), Info) != Seen.end()) return true;
    Seen.push_back(Info);
    for (const auto &[FieldName, Field] : Info->Record.Fields) {
        const TypeInfo *Inner = Types.pointerBase(Types.resolve(Field.TypeName));
        if (Inner && Inner->isC() && !Inner->isVoid()) continue;
        if (Inner && Inner->isRecord() && !Inner->IsClass) {
            if (checkCStruct(Types, Inner, Why, Seen)) continue;
            return false;
        }
        if (Why) *Why = "field " + FieldName + " of " + Info->Name + " has type " + Field.TypeName;
        return false;
    }
    return true;
}

} // namespace

const TypeInfo *TypeSystem::pointerBase(const TypeInfo *Info) const {
    std::vector<const TypeInfo *> Chain;
    while (Info && Info->Kind == TypeKind::Pointer) {
        if (std::find(Chain.begin(), Chain.end(), Info) != Chain.end()) return nullptr;
        Chain.push_back(Info);
        Info = resolve(Info->Pointer.PointeeTypeName);
    }
    return Info;
}

bool TypeSystem::isCStruct(const TypeInfo *Info, std::string *Why) const {
    std::vector<const TypeInfo *> Seen;
    return checkCStruct(*this, Info, Why, Seen);
}

bool TypeSystem::isExternType(const TypeInfo *Info) const {
    const TypeInfo *Base = pointerBase(Info);
    if (!Base) return false;
    if (Base->isC()) return !Base->isVoid();
    return isCStruct(Base);
}

const EnumConstant *TypeSystem::lookupEnumConstant(const std::string &Name) const {
    auto It = EnumConstants.find(Name);
    return It == EnumConstants.end() || (Visible && !Visible(It->second.TypeName))
               ? nullptr : &It->second;
}

Type *TypeSystem::getLLVMType(const std::string &TypeName) const {
    const TypeInfo *Info = resolve(TypeName);
    return Info ? Info->LLVMType : nullptr;
}

Constant *TypeSystem::getSizeOfConstant(Type *Ty) const {
    return ConstantExpr::getSizeOf(Ty);
}

Constant *TypeSystem::getZeroValue(const std::string &TypeName) const {
    const TypeInfo *Info = resolve(TypeName);
    if (!Info || !Info->LLVMType || Info->LLVMType->isVoidTy()) {
        return nullptr;
    }

    switch (Info->Kind) {
        case TypeKind::Integer:
            return ConstantInt::get(Context, APInt(64, 0));
        case TypeKind::Real:
            return ConstantFP::get(Context, APFloat(0.0));
        case TypeKind::Boolean:
            return ConstantInt::get(Context, APInt(1, 0));
        case TypeKind::Char:
            return ConstantInt::get(Context, APInt(8, 0));
        case TypeKind::Void:
            return nullptr;
        case TypeKind::Enum:
            return ConstantInt::get(Context, APInt(64, kEnumOrdinalBase));
        case TypeKind::Record:
            return Constant::getNullValue(Info->LLVMType); // zeroinitializer
        case TypeKind::Pointer:
            return ConstantPointerNull::get(cast<PointerType>(Info->LLVMType));
        case TypeKind::CScalar:
        case TypeKind::CPointer:
            return Constant::getNullValue(Info->LLVMType);
    }

    return Constant::getNullValue(Info->LLVMType);
}