#include "Binder.hpp"
#include "../CoreData/PrimitiveTypes.hpp"
#include "../Diagnostics/Logger.hpp"
using namespace MalachScript::Parser;
using namespace MalachScript::Diagnostics;

namespace MalachScript::Binder {
    void Binder::Bind(BoundNamespace* ns, const std::vector<const MalachScript::Parser::ParsedStatement*>& statements,
                      const Binder::log_func& log) {
        for (const auto* s : statements) {
            TypeRegistrationFirstPass(ns, s, {}, log);
        }
        for (const auto* s : statements) {
            TypeRegistrationSecondPass(ns, s, {}, log);
        }
        for (const auto* s : statements) {
            TypeRegistrationThirdPass(ns, s, {}, log);
        }
        FinaliseTypes(ns, log);
    }

    inline static std::optional<const BoundType*> ResolveType(BoundNamespace* ns, const Identifier& identifier) {
        switch (identifier) {
            case "int"_id: return MalachScript::PrimitiveTypes::IntType();
            case "int8"_id: return MalachScript::PrimitiveTypes::Int8Type();
            case "int16"_id: return MalachScript::PrimitiveTypes::Int16Type();
            case "int32"_id: return MalachScript::PrimitiveTypes::Int32Type();
            case "int64"_id: return MalachScript::PrimitiveTypes::Int64Type();

            case "uint"_id: return MalachScript::PrimitiveTypes::UintType();
            case "uint8"_id: return MalachScript::PrimitiveTypes::Uint8Type();
            case "uint16"_id: return MalachScript::PrimitiveTypes::Uint16Type();
            case "uint32"_id: return MalachScript::PrimitiveTypes::Uint32Type();
            case "uint64"_id: return MalachScript::PrimitiveTypes::Uint64Type();

            case "float"_id: return MalachScript::PrimitiveTypes::FloatType();
            case "double"_id: return MalachScript::PrimitiveTypes::DoubleType();

            case "bool"_id: return MalachScript::PrimitiveTypes::BoolType();
        }
        return ns->ResolveType(identifier);
    }

    void Binder::TypeRegistrationFirstPass(BoundNamespace* ns, const ParsedStatement* statement,
                                           std::optional<BoundType*> activeType, const log_func& log) {
        switch (statement->GetKind()) {
            case Parser::ParsedStatementKind::Script: {
                const auto* s = static_cast<const ParsedScriptStatement*>(statement);
                for (const auto& child : s->GetStatements()) {
                    TypeRegistrationFirstPass(ns, child.get(), activeType, log);
                }
                break;
            }
            case Parser::ParsedStatementKind::Class: {
                const auto* s = static_cast<const ParsedClassStatement*>(statement);
                auto identifier = s->GetIdentifier();
                auto type = ns->ResolveType(s->GetIdentifier());
                if (!type.has_value()) {
                    if (activeType.has_value()) {
                        type = activeType.value()->ResolveType(s->GetIdentifier());
                    }
                }
                if (type.has_value()) {
                    log(DiagnosticLevel::Error, DiagnosticType::TypeAlreadyDefined, s->GetSpan(),
                        {identifier.GetStdString()});
                    break;
                }

                type = new BoundType(identifier, s->GetClassAttr(), s->GetSpan());
                if (activeType.has_value()) {
                    activeType.value()->RegisterType(identifier, type.value());
                } else {
                    ns->RegisterType(identifier, type.value());
                }
                for (const auto& child : s->GetBody()) {
                    TypeRegistrationFirstPass(ns, child.get(), type, log);
                }
                break;
            }
            case Parser::ParsedStatementKind::Namespace: {
                const auto* s = static_cast<const ParsedNamespaceStatement*>(statement);
                auto identifier = s->GetIdentifier();
                auto* innerNamespace = new BoundNamespace();
                ns->RegisterNamespace(identifier, innerNamespace);
                TypeRegistrationFirstPass(innerNamespace, s->GetScript().get(), activeType, log);
                break;
            }
            default: break;
        }
    }
    void Binder::TypeRegistrationSecondPass(BoundNamespace* ns, const ParsedStatement* statement,
                                            std::optional<BoundType*> activeType, const log_func& log) {
        switch (statement->GetKind()) {
            case ParsedStatementKind::Script: {
                const auto* s = static_cast<const ParsedScriptStatement*>(statement);
                for (const auto& child : s->GetStatements()) {
                    TypeRegistrationSecondPass(ns, child.get(), activeType, log);
                }
                break;
            };
            case ParsedStatementKind::Class: {
                const auto* s = static_cast<const ParsedClassStatement*>(statement);
                auto type = ns->ResolveType(s->GetIdentifier());
                if (!type.has_value()) {
                    if (activeType.has_value()) {
                        type = activeType.value()->ResolveType(s->GetIdentifier());
                    }
                    if (!type.has_value()) {
                        throw std::logic_error("Shouldn't be reached");
                    }
                }
                for (const auto& inherits : s->GetInherits()) {
                    auto inheritType = ResolveType(ns, inherits);
                    if (!inheritType.has_value()) {
                        log(DiagnosticLevel::Error, DiagnosticType::UnknownType, s->GetSpan(),
                            {inherits.GetStdString()});
                    } else {
                        type.value()->AddInheritType(inheritType.value());
                    }
                }

                for (const auto& child : s->GetBody()) {
                    TypeRegistrationSecondPass(ns, child.get(), type.value(), log);
                }
                break;
            }
            case ParsedStatementKind::TypeDef: {
                const auto* s = static_cast<const ParsedTypeDefStatement*>(statement);
                auto defineTo = s->GetDefineTo();
                auto defineFrom = s->GetDefineFrom();
                auto defineFromType = ResolveType(ns, defineFrom);
                if (!defineFromType.has_value()) {
                    log(DiagnosticLevel::Error, DiagnosticType::UnknownType, s->GetSpan(), {defineFrom.GetStdString()});
                } else {
                    if (activeType.has_value()) {
                        activeType.value()->RegisterType(defineTo, defineFromType.value());
                    } else {
                        ns->RegisterType(defineTo, defineFromType.value());
                    }
                }
                break;
            }
            case ParsedStatementKind::Namespace: {
                const auto* s = static_cast<const ParsedNamespaceStatement*>(statement);
                auto identifier = s->GetIdentifier();
                auto innerNamespace = ns->ResolveNamespace(identifier);
                if (!innerNamespace.has_value()) {
                    throw std::logic_error("Shouldn't be reached");
                }
                TypeRegistrationSecondPass(innerNamespace.value(), s->GetScript().get(), activeType, log);
                break;
            };
            default: break;
        }
    }
    void Binder::TypeRegistrationThirdPass(BoundNamespace* ns, const MalachScript::Parser::ParsedStatement* statement,
                                           std::optional<BoundType*> activeType, const Binder::log_func& log) {
        switch (statement->GetKind()) {
            case ParsedStatementKind::Script: {
                const auto* s = static_cast<const ParsedScriptStatement*>(statement);
                for (const auto& child : s->GetStatements()) {
                    TypeRegistrationThirdPass(ns, child.get(), activeType, log);
                }
                break;
            };
            case ParsedStatementKind::Class: {
                const auto* s = static_cast<const ParsedClassStatement*>(statement);
                auto type = ns->ResolveType(s->GetIdentifier());
                if (!type.has_value()) {
                    if (activeType.has_value()) {
                        type = activeType.value()->ResolveType(s->GetIdentifier());
                    }
                    if (!type.has_value()) {
                        throw std::logic_error("Shouldn't be reached");
                    }
                }
                for (const auto& child : s->GetBody()) {
                    TypeRegistrationThirdPass(ns, child.get(), type.value(), log);
                }
                break;
            }
            case ParsedStatementKind::Namespace: {
                const auto* s = static_cast<const ParsedNamespaceStatement*>(statement);
                auto identifier = s->GetIdentifier();
                auto innerNamespace = ns->ResolveNamespace(identifier);
                if (!innerNamespace.has_value()) {
                    throw std::logic_error("Shouldn't be reached");
                }
                TypeRegistrationThirdPass(innerNamespace.value(), s->GetScript().get(), activeType, log);
                break;
            };
            case ParsedStatementKind::Var: {
                const auto* var = static_cast<const ParsedVarStatement*>(statement);
                const auto& typeStatement = var->GetTypeStatement();
                // FIXME: Resolve namespace of scoped identifier
                auto type = ResolveType(ns, typeStatement->GetScopedIdentifier().GetIdentifier());
                if (!type.has_value()) {
                    throw std::logic_error("Shouldn't be reached");
                }
                if (activeType.has_value()) {
                    activeType.value()->AddField(var->GetIdentifier(),
                                                 new BoundVariable(type.value(), var->GetAccess()));
                } else {
                    ns->AddVariable(var->GetIdentifier(), new BoundVariable(type.value(), var->GetAccess()));
                }

                break;
            }
            default: break;
        }
    }

    static void FinaliseType(BoundType* type, std::vector<const BoundType*>& stack, const Binder::log_func& log) {
        if (type->IsInitialised()) {
            return;
        }
        if (std::find(stack.begin(), stack.end(), type) != stack.end()) {
            log(Diagnostics::DiagnosticLevel::Error, Diagnostics::DiagnosticType::CircularTypeReference,
                type->GetSpan().value(),
                {stack[stack.size() - 1]->GetIdentifier().GetStdString(), type->GetIdentifier().GetStdString()});
            return;
        }

        stack.push_back(type);
        size_t size = 0;
        for (auto& inherits : type->GetInherits()) {
            if (!inherits->IsInitialised()) {
                FinaliseType(const_cast<BoundType*>(inherits), stack, log);
            }
            size += inherits->GetSize();
        }
        for (auto& field : type->GetFieldsLookup()) {
            if (!field.second->GetType()->IsInitialised()) {
                FinaliseType(const_cast<BoundType*>(field.second->GetType()), stack, log);
            }
            size += field.second->GetType()->GetSize();
        }
        stack.pop_back();
        type->Finalise(size);
    }

    static void FinaliseNamespace(BoundNamespace* ns, const Binder::log_func& log) {
        std::vector<const BoundType*> s = {};
        for (const auto& t : ns->GetTypes()) {
            FinaliseType(t.second, s, log);
        }
        for (const auto& n : ns->GetNamespaces()) {
            FinaliseNamespace(n.second, log);
        }
    }

    void Binder::FinaliseTypes(BoundNamespace* ns, const Binder::log_func& log) { FinaliseNamespace(ns, log); }
}