#include "cps/CodeGen.h"
#include "cps/ArrayHandler.h"
#include "cps/FileAST.h"
#include "cps/FileHandler.h"
#include "cps/MemoryHandler.h"
#include "cps/ExceptionHandler.h"
#include "cps/TypeAST.h"
#include "cps/TemplateAST.h"

using namespace llvm;
using namespace cps;

void CodeGen::emitStmt(StmtAST *Stmt) {
    size_t Start = StackObjects.size();
    emitStmtImpl(Stmt);
    cleanupTemporaries(Start);
}

void CodeGen::emitStmtImpl(StmtAST *Stmt) {
    SaveAndRestore<int> Location(DiagnosticLine, Stmt && Stmt->getLine() ? Stmt->getLine() : DiagnosticLine);
    if (!Stmt) return;
    if (dynamic_cast<TemplateDeclAST*>(Stmt)) return;

    if (dynamic_cast<IncludeStmtAST*>(Stmt)) return;
    if (dynamic_cast<GlobalBlockAST*>(Stmt)) return;
    if (dynamic_cast<IncludeObjStmtAST*>(Stmt) || dynamic_cast<ExternDeclAST*>(Stmt)) return;

    if (auto *Class = dynamic_cast<ClassTypeDeclAST*>(Stmt)) {
        for (const auto &Method : Class->getMethods()) {
            // Invalid or duplicate members have already been diagnosed.
            if (isRegisteredMethod(Class->getName(), Method->getProto()))
                emitStmt(Method.get());
        }
        return;
    }

    if (auto *FuncDef = dynamic_cast<FunctionDefAST*>(Stmt)) {
        SaveAndRestore<std::string> ClassContext(CurrentClass, FuncDef->getProto()->getOwnerClass());
        for (const auto &Param : FuncDef->getProto()->getParams())
            warnMemberShadow(Param.Name, Param.Line ? Param.Line : DiagnosticLine);
        BasicBlock *SavedBlock = Builder->GetInsertBlock();
        // Function bodies get a fresh array scope (metadata SSA values belong
        // to the emitting function), mirroring FunctionGen's Symbols handling.
        auto SavedArrays = Arrays->exchangeTable({});

        auto SavedObjects = std::move(StackObjects);
        StackObjects.clear();
        auto SavedScopes = std::move(ConstantScopes);
        ConstantScopes = {{NextScopeID++, {}}};
        Function *Emitted = FuncGen->emitFunctionDef(FuncDef, [this](StmtAST *S) {
            this->emitStmt(S);
        }, ConstantScopes.back().ID);
        if (Emitted) finishStackObjects(Emitted);
        StackObjects = std::move(SavedObjects);
        ConstantScopes = std::move(SavedScopes);

        Arrays->exchangeTable(std::move(SavedArrays));
        if (SavedBlock) Builder->SetInsertPoint(SavedBlock);
        return;
    }

    // Keep diagnosing statements after RETURN/THROW without appending code
    // to a terminated block or making those statements reachable.
    if (Builder->GetInsertBlock()->hasTerminator()) {
        auto *F = Builder->GetInsertBlock()->getParent();
        Builder->SetInsertPoint(BasicBlock::Create(*TheContext, "unreachable_code", F));
    }

    if (auto *Delete = dynamic_cast<DeleteStmtAST*>(Stmt)) {
        LValueInfo Object;
        if (!emitLValue(Delete->getTarget(), Object)) return;
        const TypeInfo *Info = resolveType(Object.TypeName);
        if (!Info || !Info->IsClass) {
            reportError("DELETE requires a class object");
            return;
        }
        Memory->emitValidate(Object.Addr, Stmt->getLine());
        if (emitDestroy(Object.TypeName, Object.Addr, Stmt->getLine()))
            Memory->emitRelease(Object.Addr, Stmt->getLine());
        return;
    }
    if (auto *Try = dynamic_cast<TryStmtAST*>(Stmt)) {
        emitTryStmt(Try);
        return;
    }
    if (dynamic_cast<ThrowStmtAST*>(Stmt)) {
        Exceptions->emitThrow();
        return;
    }
    if (auto *ArrDecl = dynamic_cast<ArrayDeclareStmtAST*>(Stmt)) {
        Arrays->emitArrayDeclare(ArrDecl, *this);
        return;
    }
    if (auto *ArrAssign = dynamic_cast<ArrayAssignStmtAST*>(Stmt)) {
        Arrays->emitArrayAssign(ArrAssign, *this);
        return;
    }

    if (auto *Call = dynamic_cast<MemberCallStmtAST*>(Stmt)) {
        Value *Result = emitMemberCall(Call->getCall());
        const TypeInfo *Type = getExprTypeInfo(Call->getCall());
        if (Result && Type && (Type->IsClass || containsString(Type))) materializeOperatorReceiver(Result, Type);
        return;
    }

    if (auto *Call = dynamic_cast<CallStmtAST*>(Stmt)) {
        Value *Result = emitCall(Call->getCallee(), Call->getArgs(), Call->getLine());
        const FuncSig *Sig = FuncGen->hasFunction(Call->getCallee())
            ? resolveCall(Call->getCallee(), Call->getArgs()) : nullptr;
        const TypeInfo *Type = Sig ? resolveType(Sig->ReturnTypeName) : nullptr;
        if (Result && Type && (Type->IsClass || containsString(Type))) materializeOperatorReceiver(Result, Type);
        return;
    }

    if (auto *Ret = dynamic_cast<ReturnStmtAST*>(Stmt)) {
        Function *F = Builder->GetInsertBlock()->getParent();
        const FuncSig *Sig = F->getName() == "main" ? nullptr
            : FuncGen->getSignatureByLink(F->getName().str());
        // The generated wrapper has no signature; user MAIN uses __cps_main.
        if (!Sig) {
            reportError("RETURN is only allowed inside a FUNCTION or PROCEDURE");
            return;
        }
        if (F->getReturnType()->isVoidTy() && Ret->getRetVal()) {
            reportError("RETURN with a value is only allowed in a FUNCTION");
            return;
        }
        if (!F->getReturnType()->isVoidTy() && !Ret->getRetVal()) {
            reportError("RETURN in a FUNCTION requires a value");
            return;
        }
        Value *RetVal = nullptr;
        if (Ret->getRetVal()) {
            std::vector<size_t> Transfers;
            RetVal = emitOwnedValue(Ret->getRetVal(), resolveType(Sig->ReturnTypeName), Transfers);
            transferObjects(Transfers);
            if (!RetVal) return;
        }
        cleanupTemporaries();
        FuncGen->emitReturn(Ret, RetVal);
        return;
    }

    if (auto *Const = dynamic_cast<ConstantStmtAST*>(Stmt)) {
        emitConstantStmt(Const);
        return;
    }

    if (auto *Decl = dynamic_cast<DeclareStmtAST*>(Stmt)) {
        emitDeclareStmt(Decl);
        return;
    }

    if (auto *Assign = dynamic_cast<AssignStmtAST*>(Stmt)) {
        const SymbolInfo *Info = getSymbolInfo(Assign->getName());
        if (!Info) {
            if (Types->lookupEnumConstant(Assign->getName())) {
                reportError("Enum value %s is not an assignable location", Assign->getName().c_str());
            } else {
                reportError("Unknown variable name %s", Assign->getName().c_str());
            }
            return;
        }
        if (!requireWritable(*Info, Assign->getName())) return;
        if (Info->IsArray) {
            reportError("Cannot assign array %s without indices", Assign->getName().c_str());
            return;
        }

        emitAssignment({symbolAddress(*Info), Info->TypeName}, Assign->getExpr(), Info);
        return;
    }

    if (auto *In = dynamic_cast<InputStmtAST*>(Stmt)) {
        LValueInfo Target;
        if (!emitLValue(In->getTarget(), Target)) return;
        const TypeInfo *TypeInfo = resolveType(Target.TypeName);
        if (!TypeInfo) {
            reportError("Unknown type for INPUT %s", Target.TypeName.c_str());
            return;
        }

        std::vector<Value*> Args;
        if (TypeInfo->isC()) {
            emitCInput(Target.Addr, TypeInfo);
        } else if (TypeInfo->isReal()) {
            Args.push_back(ScanfFloatFormatStr);
            Args.push_back(Target.Addr);
            Builder->CreateCall(ScanfFunc, Args);
        } else if (TypeInfo->isBoolean()) {
            AllocaInst *TempInt = CreateEntryBlockAlloca(Builder->GetInsertBlock()->getParent(), "tmp_bool_input");
            Args.push_back(ScanfFormatStr);
            Args.push_back(TempInt);
            Builder->CreateCall(ScanfFunc, Args);

            Value *Val = Builder->CreateLoad(Type::getInt64Ty(*TheContext), TempInt);
            Value *BoolVal = Builder->CreateICmpNE(Val, ConstantInt::get(*TheContext, APInt(64, 0)), "bool_cast");
            Builder->CreateStore(BoolVal, Target.Addr);
        } else if (TypeInfo->isString()) {
            emitMethodCall(TypeInfo->Name, Target.Addr, "READ_INPUT", {}, Stmt->getLine());
        } else if (TypeInfo->isChar()) {
            Value *CharFmt = Builder->CreateGlobalString(" %c", "fmt_in_char", 0, TheModule.get());
            Args.push_back(CharFmt);
            Args.push_back(Target.Addr);
            Builder->CreateCall(ScanfFunc, Args);
        } else if (TypeInfo->Kind == TypeKind::Integer) {
            Args.push_back(ScanfFormatStr);
            Args.push_back(Target.Addr);
            Builder->CreateCall(ScanfFunc, Args);
        } else {
            // Enum/Record/Pointer and any future kind: scanf'ing raw bytes
            // into typed storage would corrupt it — reject explicitly.
            reportError("INPUT is not supported for a value of type %s", TypeInfo->Name.c_str());
        }
        return;
    }

    if (auto *Out = dynamic_cast<OutputStmtAST*>(Stmt)) {
        const auto &Exprs = Out->getExprs();
        for (size_t I = 0; I < Exprs.size(); ++I) {
            ExprAST *Expr = Exprs[I].get();
            if (Arrays->tryEmitArrayOutput(Expr, *this)) continue;

            Value *Val = emitExpr(Expr);
            const TypeInfo *TypeInfo = getExprTypeInfo(Expr);
            if (!TypeInfo) {
                reportError("Cannot infer the type of the OUTPUT expression");
                continue;
            }
            if (TypeInfo->isVoid()) {
                reportError("Cannot OUTPUT the result of a PROCEDURE call");
                continue;
            }
            if (TypeInfo->IsClass && !isObjectLValue(Expr))
                materializeOperatorReceiver(Val, TypeInfo);
            emitOutputValue(Val, TypeInfo, I + 1 == Exprs.size());
        }
        return;
    }

    if (auto *Open = dynamic_cast<OpenFileStmtAST*>(Stmt)) {
        if (isWholeArrayVar(Open->getFileName())) {
            reportError("Cannot use a whole array as the OPENFILE filename");
            return;
        }
        Value *Name = emitCStringExpr(Open->getFileName());
        if (!Name) return;
        Files->emitOpen(Name, Open->getMode(), Open->getLine());
        return;
    }

    if (auto *Read = dynamic_cast<ReadFileStmtAST*>(Stmt)) {
        const SymbolInfo *Info = getSymbolInfo(Read->getVarName());
        if (!Info) {
            reportError("Unknown variable name %s", Read->getVarName().c_str());
            return;
        }
        if (!requireWritable(*Info, Read->getVarName())) return;
        if (Info->IsArray) {
            reportError("READFILE into entire array %s is not supported", Read->getVarName().c_str());
            return;
        }
        if (!sameIdentifier(Info->TypeName, "STRING")) {
            reportError("READFILE variable %s must be of type STRING (got %s)",
                        Read->getVarName().c_str(), Info->TypeName.c_str());
            return;
        }
        if (isWholeArrayVar(Read->getFileName())) {
            reportError("Cannot use a whole array as the READFILE filename");
            return;
        }
        Value *Name = emitCStringExpr(Read->getFileName());
        if (!Name) return;
        Value *LineStr = Files->emitRead(Name, Read->getLine());
        if (!LineStr) return;
        const auto *Type = resolveType(Info->TypeName);
        for (const FuncSig *Sig : FuncGen->getOverloads(Type->Name + ".OPERATOR <-")) {
            if (Sig->Params.size() == 1 && !Sig->Params[0].IsByRef &&
                resolveType("C_STRING")->sameType(resolveType(Sig->Params[0].TypeName)) &&
                checkOperatorAccess(*Sig)) {
                FuncGen->emitCall(*Sig, {symbolAddress(*Info), LineStr});
                auto Free = TheModule->getOrInsertFunction("free", FunctionType::get(
                    Type::getVoidTy(*TheContext), {PointerType::getUnqual(*TheContext)}, false));
                Builder->CreateCall(Free, {LineStr});
                return;
            }
        }
        reportError("READFILE requires STRING.OPERATOR<-(Text : C_STRING)");
        return;
    }

    if (auto *Write = dynamic_cast<WriteFileStmtAST*>(Stmt)) {
        if (isWholeArrayVar(Write->getFileName()) || isWholeArrayVar(Write->getData())) {
            reportError("Cannot use a whole array in WRITEFILE");
            return;
        }
        Value *Name = emitCStringExpr(Write->getFileName());
        Value *Data = emitCStringExpr(Write->getData());
        if (!Name || !Data) return;
        Files->emitWrite(Name, Data, Write->getLine());
        return;
    }

    if (auto *Close = dynamic_cast<CloseFileStmtAST*>(Stmt)) {
        if (isWholeArrayVar(Close->getFileName())) {
            reportError("Cannot use a whole array as the CLOSEFILE filename");
            return;
        }
        Value *Name = emitCStringExpr(Close->getFileName());
        if (!Name) return;
        Files->emitClose(Name, Close->getLine());
        return;
    }

    if (auto *Case = dynamic_cast<CaseStmtAST*>(Stmt)) {
        emitCaseStmt(Case);
        return;
    }
    if (auto *IfStmt = dynamic_cast<IfStmtAST*>(Stmt)) {
        emitIfStmt(IfStmt);
        return;
    }
    if (auto *WhileStmt = dynamic_cast<WhileStmtAST*>(Stmt)) {
        emitWhileStmt(WhileStmt);
        return;
    }
    if (auto *RepeatStmt = dynamic_cast<RepeatStmtAST*>(Stmt)) {
        emitRepeatStmt(RepeatStmt);
        return;
    }
    if (auto *ForStmt = dynamic_cast<ForStmtAST*>(Stmt)) {
        emitForStmt(ForStmt);
        return;
    }

    // TYPE declarations are registered by compile()'s pre-pass (the parser
    // guarantees they are top-level); nothing is emitted here.
    if (dynamic_cast<EnumTypeDeclAST*>(Stmt) || dynamic_cast<PointerTypeDeclAST*>(Stmt) ||
        dynamic_cast<RecordTypeDeclAST*>(Stmt) || dynamic_cast<CTypeDeclAST*>(Stmt)) {
        return;
    }
    if (auto *DA = dynamic_cast<DesignatorAssignStmtAST*>(Stmt)) {
        LValueInfo LV;
        if (!emitLValue(DA->getTarget(), LV)) return;
        emitAssignment(LV, DA->getExpr());
        return;
    }
}