#include <search/web/core/rule.h>

#include <search/meta/context.h>

#include <kernel/snippets/strhl/goodwrds.h>

#include <kernel/qtree/richrequest/serialization/serializer.h>

#include <library/cpp/html/face/onchunk.h>
#include <library/cpp/html/html5/parse.h>

#include <util/generic/iterator_range.h>

#include <array>


namespace {

    constexpr std::array SKIP_OPEN_TAGS = {
        TStringBuf("<html>"),
        TStringBuf("<head>"),
        TStringBuf("<body>")
    };

    constexpr std::array SKIP_END_TAGS = {
        TStringBuf("</html>"),
        TStringBuf("</head>"),
        TStringBuf("</body>")
    };

    inline TString HighlightChunk(
        const TStringBuf text,
        const TInlineHighlighter& highlighter,
        const TPaintingOptions& options)
    {
        auto wideText = UTF8ToWide(text);
        highlighter.PaintPassages(wideText, options);
        return WideToUTF8(wideText);
    }


    class TTextNodeHighlighter final : public IParserResult
    {
    public:
        TTextNodeHighlighter(
            IOutputStream& stream,
            const TInlineHighlighter& highlighter,
            const TPaintingOptions& options)
            : Stream(stream)
            , Highlighter(highlighter)
            , Options(options)
        {}

        THtmlChunk* OnHtmlChunk(const THtmlChunk& chunk) override {
            const auto text = chunk.GetText();
            switch (chunk.GetLexType()) {
                case HTLEX_TYPE::HTLEX_TEXT:
                    Stream << HighlightChunk(text, Highlighter, Options);
                    break;
                case HTLEX_TYPE::HTLEX_START_TAG:
                    if (FindPtr(SKIP_OPEN_TAGS, text) == nullptr) {
                        Stream << text;
                    }
                    break;
                case HTLEX_TYPE::HTLEX_END_TAG:
                    if (FindPtr(SKIP_END_TAGS, text) == nullptr) {
                        Stream << text;
                    }
                    break;
                default:
                    Stream << text;
                    break;
            }
            return nullptr;
        }

    private:
        IOutputStream& Stream;
        const TInlineHighlighter& Highlighter;
        const TPaintingOptions& Options;
    };


    inline TString HighlightText(
        const TStringBuf text,
        const TInlineHighlighter& highlighter,
        const TPaintingOptions& options)
    {
        constexpr auto reserveFactor = 1.1;
        TString result(Reserve(text.size() * reserveFactor));
        TStringOutput stream(result);
        TTextNodeHighlighter parser(stream, highlighter, options);
        NHtml5::ParseHtml(text, &parser);
        stream.Flush();
        return result;
    }

} // namespace

namespace NSaas {

    class TPatentHighlighterRuleContext final : public IRearrangeRuleContext {
    public:
        void DoRearrangeAfterFetch(TRearrangeParams& rearrangeParams) override {
            if (!rearrangeParams.Current.Grouping) {
                return;
            }

            const auto& cgi = rearrangeParams.GetRequestCgi();
            if (!cgi.Has("hl_qtree") || !cgi.Has("hl_gta")) {
                return;
            }

            NSearchQuery::TRequest qtree;
            TRichTreeDeserializer{}.DeserializeBase64(cgi.Get("hl_qtree"), qtree);

            TInlineHighlighter highlighter;
            highlighter.AddRequest(*qtree.Root);

            auto groups = rearrangeParams.Current.Grouping;
            Y_ENSURE(groups && groups->Size() == 1);

            auto& documents = groups->GetMetaGroup(0).MetaDocs;
            Y_ENSURE(documents.size() == 1);

            auto attributes = documents.front().MutableAttributes();
            Y_ENSURE(attributes);

            TPaintingOptions options;
            options.SkipAttrs = true;

            /// Highlights text in nodes with names corresponding to passed gta attributes
            const auto [begin, end] = cgi.equal_range("hl_gta");
            for (const auto& [_, attr] : TIteratorRange(begin, end)) {
                TMsString text;
                if (!attributes->GetOneValue(text, attr)) {
                    continue;
                }

                attributes->Erase(attr);
                attributes->Add(attr, HighlightText(text, highlighter, options));
            }
        }
    };


    class TPatentHighlighterRule final : public IRearrangeRule {
    public:
        IRearrangeRuleContext* DoConstructContext() const override {
            return new TPatentHighlighterRuleContext();
        }
    };

} // namespace NSaas

IRearrangeRule* CreatePatentHighlighterRule(const TString&, const TSearchConfig&) {
    return new NSaas::TPatentHighlighterRule();
}

REGISTER_REARRANGE_RULE(PatentHighlighter, CreatePatentHighlighterRule);
