#include <fesa/io/abaqus/set_resolver.hpp>

#include <algorithm>
#include <charconv>
#include <compare>
#include <cstddef>
#include <cstdint>
#include <map>
#include <set>
#include <string>
#include <string_view>
#include <utility>
#include <vector>

namespace fesa {
namespace {

struct SetNamespace final {
    ResolvedSetScope scope;
    std::string scope_name;
    ResolvedSetKind kind;

    auto operator<=>(const SetNamespace&) const = default;
};

struct SetKey final {
    SetNamespace name_space;
    std::string set_name;

    auto operator<=>(const SetKey&) const = default;
};

struct RawMember final {
    std::string text;
    SourceLocation source;
};

struct RawSet final {
    std::vector<RawMember> members;
};

enum class VisitState { unvisited, visiting, resolved, failed };

const std::string* parameter(
    const DeckRecord& record,
    const std::string_view name) {
    const auto found = record.parameters.find(name);
    return found == record.parameters.end() ? nullptr : &found->second;
}

bool has_parameter(
    const DeckRecord& record,
    const std::string_view name) {
    return record.parameters.contains(name);
}

SourceLocation data_source(
    const DeckRecord& record,
    const std::size_t row) {
    return row < record.data_sources.size()
               ? record.data_sources[row]
               : record.source;
}

bool parse_positive_label(
    const std::string_view text,
    std::int64_t& value) {
    const auto parsed =
        std::from_chars(text.data(), text.data() + text.size(), value);
    return parsed.ec == std::errc{} &&
           parsed.ptr == text.data() + text.size() && value > 0;
}

class SetResolver final {
public:
    explicit SetResolver(const ParsedDeck& deck) : deck_{deck} {}

    [[nodiscard]] SetResolutionResult resolve() {
        collect_scope(
            deck_.global_records,
            {ResolvedSetScope::global, "global", ResolvedSetKind::node});

        for (const ParsedPart& part : deck_.parts) {
            collect_scope(
                part.records,
                {ResolvedSetScope::part,
                 part.name,
                 ResolvedSetKind::node});
        }

        collect_assembly();

        std::vector<ResolvedSet> sets;
        sets.reserve(raw_sets_.size());
        for (const auto& [key, raw_set] : raw_sets_) {
            static_cast<void>(raw_set);
            if (!resolve_set(key)) {
                continue;
            }
            sets.push_back({
                key.name_space.scope_name,
                key.set_name,
                resolved_sets_.at(key),
                key.name_space.scope,
                key.name_space.kind,
            });
        }

        if (!diagnostics_.empty()) {
            sets.clear();
        }
        return {std::move(sets), std::move(diagnostics_)};
    }

private:
    void add_error(
        std::string code,
        std::string message,
        const SourceLocation& source) {
        diagnostics_.push_back({
            DiagnosticStage::semantic,
            Severity::error,
            std::move(code),
            std::move(message),
            source,
        });
    }

    void collect_scope(
        const std::vector<DeckRecord>& records,
        SetNamespace name_space) {
        collect_entities(records, name_space);
        collect_set_records(records, std::move(name_space), nullptr);
    }

    void collect_entities(
        const std::vector<DeckRecord>& records,
        const SetNamespace& base_namespace) {
        for (const DeckRecord& record : records) {
            ResolvedSetKind kind;
            if (record.keyword == "NODE") {
                kind = ResolvedSetKind::node;
            } else if (record.keyword == "ELEMENT") {
                kind = ResolvedSetKind::element;
            } else {
                continue;
            }

            SetNamespace entity_namespace = base_namespace;
            entity_namespace.kind = kind;
            std::set<std::int64_t>& labels = entities_[entity_namespace];
            for (std::size_t row = 0; row < record.data.size(); ++row) {
                if (record.data[row].empty()) {
                    continue;
                }
                std::int64_t label = 0;
                if (!parse_positive_label(record.data[row][0], label)) {
                    continue;
                }
                labels.insert(label);

                if (kind != ResolvedSetKind::element) {
                    continue;
                }
                const std::string* set_name = parameter(record, "ELSET");
                if (set_name == nullptr || set_name->empty()) {
                    continue;
                }
                raw_sets_[{entity_namespace, *set_name}].members.push_back({
                    std::to_string(label),
                    data_source(record, row),
                });
            }
        }
    }

    void collect_set_records(
        const std::vector<DeckRecord>& records,
        const SetNamespace& base_namespace,
        const std::string* active_instance) {
        for (const DeckRecord& record : records) {
            SetNamespace set_namespace = base_namespace;
            std::string_view name_parameter;
            if (record.keyword == "NSET") {
                set_namespace.kind = ResolvedSetKind::node;
                name_parameter = "NSET";
            } else if (record.keyword == "ELSET") {
                set_namespace.kind = ResolvedSetKind::element;
                name_parameter = "ELSET";
            } else {
                continue;
            }

            const std::string* set_name = parameter(record, name_parameter);
            if (set_name == nullptr || set_name->empty()) {
                add_error(
                    "abaqus.semantic.missing_parameter",
                    "*" + record.keyword + " requires parameter " +
                        std::string{name_parameter} + ".",
                    record.source);
                continue;
            }

            if (set_namespace.scope == ResolvedSetScope::assembly) {
                const std::string* instance = parameter(record, "INSTANCE");
                if (active_instance == nullptr || instance == nullptr ||
                    *instance != *active_instance) {
                    add_error(
                        "abaqus.semantic.wrong_instance",
                        "Assembly set '" + *set_name +
                            "' must reference the active Instance.",
                        record.source);
                    continue;
                }
            }

            RawSet& raw_set = raw_sets_[{set_namespace, *set_name}];
            if (has_parameter(record, "GENERATE")) {
                collect_generate(record, raw_set);
            } else {
                collect_explicit(record, raw_set);
            }
        }
    }

    void collect_explicit(
        const DeckRecord& record,
        RawSet& raw_set) {
        for (std::size_t row = 0; row < record.data.size(); ++row) {
            for (const std::string& field : record.data[row]) {
                if (field.empty()) {
                    continue;
                }
                raw_set.members.push_back({field, data_source(record, row)});
            }
        }
    }

    void collect_generate(
        const DeckRecord& record,
        RawSet& raw_set) {
        const SourceLocation source =
            record.data.empty() ? record.source : data_source(record, 0U);
        if (record.data.size() != 1U || record.data[0].size() != 3U ||
            std::ranges::any_of(
                record.data[0],
                [](const std::string& field) { return field.empty(); })) {
            add_error(
                "abaqus.semantic.invalid_generate",
                "*" + record.keyword +
                    ", GENERATE requires exactly start, end, increment.",
                source);
            return;
        }

        std::int64_t start = 0;
        std::int64_t end = 0;
        std::int64_t increment = 0;
        if (!parse_positive_label(record.data[0][0], start) ||
            !parse_positive_label(record.data[0][1], end) ||
            !parse_positive_label(record.data[0][2], increment) ||
            start > end || (end - start) % increment != 0) {
            add_error(
                "abaqus.semantic.invalid_generate",
                "Invalid *" + record.keyword + " generate range.",
                source);
            return;
        }

        for (std::int64_t label = start;; label += increment) {
            raw_set.members.push_back({std::to_string(label), source});
            if (label == end) {
                break;
            }
        }
    }

    void collect_assembly() {
        if (!deck_.assembly.has_value()) {
            return;
        }

        const ParsedAssembly& assembly = *deck_.assembly;
        const ParsedInstance* active_instance =
            assembly.instances.size() == 1U
                ? &assembly.instances.front()
                : nullptr;
        SetNamespace assembly_namespace{
            ResolvedSetScope::assembly,
            assembly.name,
            ResolvedSetKind::node,
        };

        if (active_instance != nullptr) {
            const auto part = std::ranges::find(
                deck_.parts, active_instance->part_name, &ParsedPart::name);
            if (part != deck_.parts.end()) {
                for (const ResolvedSetKind kind : {
                         ResolvedSetKind::node,
                         ResolvedSetKind::element}) {
                    const SetNamespace part_namespace{
                        ResolvedSetScope::part,
                        part->name,
                        kind,
                    };
                    SetNamespace lifted_namespace = assembly_namespace;
                    lifted_namespace.kind = kind;
                    const auto labels = entities_.find(part_namespace);
                    if (labels != entities_.end()) {
                        entities_[lifted_namespace] = labels->second;
                    }
                }
            }
        }

        const std::string* instance_name =
            active_instance == nullptr ? nullptr : &active_instance->name;
        collect_set_records(
            assembly.records, assembly_namespace, instance_name);
    }

    bool resolve_set(const SetKey& key) {
        VisitState& state = states_[key];
        if (state == VisitState::resolved) {
            return true;
        }
        if (state == VisitState::failed) {
            return false;
        }

        state = VisitState::visiting;
        bool succeeded = true;
        std::vector<std::int64_t> labels;
        for (const RawMember& member : raw_sets_.at(key).members) {
            std::int64_t label = 0;
            if (parse_positive_label(member.text, label)) {
                const auto entity_namespace = entities_.find(key.name_space);
                if (entity_namespace == entities_.end() ||
                    !entity_namespace->second.contains(label)) {
                    add_error(
                        "abaqus.semantic.missing_set_member",
                        "Set '" + key.set_name +
                            "' references missing entity label " +
                            std::to_string(label) + ".",
                        member.source);
                    succeeded = false;
                    continue;
                }
                labels.push_back(label);
                continue;
            }

            const SetKey nested_key{key.name_space, member.text};
            const auto nested = raw_sets_.find(nested_key);
            if (nested == raw_sets_.end()) {
                add_error(
                    "abaqus.semantic.missing_set_member",
                    "Set '" + key.set_name + "' references missing set '" +
                        member.text + "'.",
                    member.source);
                succeeded = false;
                continue;
            }

            if (states_[nested_key] == VisitState::visiting) {
                add_error(
                    "abaqus.semantic.set_cycle",
                    "Set '" + key.set_name +
                        "' closes a nested set reference cycle through '" +
                        member.text + "'.",
                    member.source);
                succeeded = false;
                continue;
            }
            if (!resolve_set(nested_key)) {
                succeeded = false;
                continue;
            }
            const std::vector<std::int64_t>& nested_labels =
                resolved_sets_.at(nested_key);
            labels.insert(
                labels.end(), nested_labels.begin(), nested_labels.end());
        }

        if (!succeeded) {
            state = VisitState::failed;
            return false;
        }

        std::ranges::sort(labels);
        labels.erase(std::ranges::unique(labels).begin(), labels.end());
        resolved_sets_[key] = std::move(labels);
        state = VisitState::resolved;
        return true;
    }

    const ParsedDeck& deck_;
    std::map<SetNamespace, std::set<std::int64_t>> entities_;
    std::map<SetKey, RawSet> raw_sets_;
    std::map<SetKey, VisitState> states_;
    std::map<SetKey, std::vector<std::int64_t>> resolved_sets_;
    std::vector<Diagnostic> diagnostics_;
};

}  // namespace

SetResolutionResult resolve_sets(const ParsedDeck& deck) {
    return SetResolver{deck}.resolve();
}

}  // namespace fesa