// Copyright (C) 2018-2026 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//

#include "append_list_unpack_replacer.hpp"

#include <memory>
#include <utility>

#include "openvino/core/graph_util.hpp"
#include "openvino/core/rt_info.hpp"
#include "openvino/core/validation_util.hpp"
#include "openvino/frontend/sequence_mark.hpp"
#include "openvino/op/constant.hpp"
#include "openvino/op/split.hpp"
#include "openvino/op/squeeze.hpp"
#include "openvino/op/util/framework_node.hpp"
#include "openvino/pass/pattern/matcher.hpp"
#include "openvino/pass/pattern/op/wrap_type.hpp"
#include "utils.hpp"

namespace ov {
namespace frontend {
namespace pytorch {
namespace pass {

using namespace ov::op;

AppendListUnpackReplacer::AppendListUnpackReplacer() {
    auto list_unpack =
        ov::pass::pattern::wrap_type<ov::op::util::FrameworkNode>(fw_node_predicate({"prim::ListUnpack"}));

    ov::matcher_pass_callback callback = [](ov::pass::pattern::Matcher& m) {
        auto list_unpack = m.get_match_root();

        NodeVector rt_copy_from;
        auto input_node = list_unpack->input_value(0).get_node_shared_ptr();

        // Optional aten::__getitem__ node.
        auto getitem_node = cast_fw_node(input_node, "aten::__getitem__");
        if (getitem_node) {
            rt_copy_from.push_back(getitem_node);
            input_node = getitem_node->input(0).get_source_output().get_node_shared_ptr();
        }

        auto seq_mark = ov::as_type_ptr<SequenceMark>(input_node);
        if (!seq_mark) {
            return false;
        }
        auto inputs = seq_mark->get_sequence();
        if (getitem_node) {
            // If aten::__getitem__, expect inputs to be equivalent of pytorch Tensor[][].
            // Tensor selected by aten::__getitem__ index needs to be split along axis 0.
            auto getitem_index_const = ov::util::get_constant_from_source(getitem_node->input_value(1));
            if (!getitem_index_const)
                return false;
            auto index_val = getitem_index_const->cast_vector<int64_t>();
            if (index_val.size() != 1) {
                add_exception_to_fw_node(list_unpack, "prim::ListUnpack: index of aten::__getitem__ is not scalar.");
                return false;
            }
            auto index = index_val[0];
            if (index_val[0] < 0) {
                index = inputs.size() + index;
            }
            auto axis_0 = v0::Constant::create(element::i32, Shape{}, {0});
            auto split = std::make_shared<v1::Split>(inputs[index], axis_0, list_unpack->get_output_size());
            NodeVector to_copy_rt{axis_0, split};
            OutputVector res;
            for (auto& output : split->outputs()) {
                auto squeeze = std::make_shared<v0::Squeeze>(output, axis_0);
                to_copy_rt.push_back(squeeze);
                res.push_back(squeeze);
            }
            copy_runtime_info_and_name(list_unpack, std::move(to_copy_rt), rt_copy_from);
            replace_node(list_unpack, res);
            return true;
        } else {
            // Without aten::__getitem__, expect inputs to be equivalent of pytorch Tensor[].
            // Return all inputs.
            replace_node(list_unpack, inputs);
            return true;
        }
        return false;
    };

    auto m = std::make_shared<ov::pass::pattern::Matcher>(list_unpack,
                                                          "ov::frontend::pytorch::pass::AppendListUnpackReplacer");
    this->register_matcher(m, callback);
};

}  // namespace pass
}  // namespace pytorch
}  // namespace frontend
}  // namespace ov
