package me.qwq.doghouse.flexmark.ext.plantuml;

import java.io.ByteArrayOutputStream;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.zip.Deflater;
import java.util.zip.DeflaterOutputStream;

import org.apache.commons.lang3.StringUtils;
import org.jetbrains.annotations.NotNull;

import com.vladsch.flexmark.ast.FencedCodeBlock;
import com.vladsch.flexmark.html.HtmlRendererOptions;
import com.vladsch.flexmark.html.HtmlWriter;
import com.vladsch.flexmark.html.renderer.NodeRenderer;
import com.vladsch.flexmark.html.renderer.NodeRendererContext;
import com.vladsch.flexmark.html.renderer.NodeRendererFactory;
import com.vladsch.flexmark.html.renderer.NodeRenderingHandler;
import com.vladsch.flexmark.util.data.DataHolder;
import com.vladsch.flexmark.util.html.Attribute;
import com.vladsch.flexmark.util.misc.CharPredicate;
import com.vladsch.flexmark.util.sequence.BasedSequence;

public class PlantUMLNodeRenderer implements NodeRenderer {


    final private static String SAME_ORIGIN = "same-origin";
    final private static String REFERRER_POLICY = "referrerpolicy";
    final private static String PLANTUML_URL_FORMAT = "https://www.plantuml.com/plantuml/%s/%s";
    final private static String PLANTUML_URL_FORMAT_DARK = "https://www.plantuml.com/plantuml/d%s/%s";
    final PlantUMLOptions options;
    final private static Pattern FMT_PATTERN = Pattern.compile("(:(png|svg))");
    final private static Pattern ATTR_PATTERN = Pattern.compile("(\\[([\\w/]+?)(=(.*?))?\\])");
    
    public PlantUMLNodeRenderer(DataHolder options) {
        this.options = new PlantUMLOptions(options);
    }
    
    @Override
    public Set<NodeRenderingHandler<?>> getNodeRenderingHandlers() {
        Set<NodeRenderingHandler<?>> set = new HashSet<>();
        if (options.renderBlockPlantUML) {
            set.add(new NodeRenderingHandler<>(FencedCodeBlock.class, PlantUMLNodeRenderer.this::render));
        }
        return set;
    }
    
    private void render(FencedCodeBlock node, HtmlWriter html, HtmlRendererOptions htmlOptions, String format, Map<String, String> attrs, String encode, boolean darkMode) {

        String clazz = this.options.blockPlantUMLClass;
        if (darkMode && this.options.darkModeClass != null) {
            clazz += ' ' + this.options.darkModeClass;
        }
        html.line()
        .srcPosWithTrailingEOL(node.getChars())
        .attr(Attribute.CLASS_ATTR, clazz)
        .attr("format", format);
        
        for (Entry<String, String> entry : attrs.entrySet()) {
            String[] keys = entry.getKey().split("/");
            String value = entry.getValue();
            for (String key : keys) {
                html.attr(key, value);
            }
        }
        
        if (this.options.imgReferrerPolicySameOrigin) {
            html.attr(REFERRER_POLICY, SAME_ORIGIN);
        }
        html
        .attr("src", String.format(!darkMode ? PLANTUML_URL_FORMAT : PLANTUML_URL_FORMAT_DARK, format, encode));
        html.withAttr().tag("img").line().openPre().closePre().tag("/img").lineIf(htmlOptions.htmlBlockCloseTagEol);
    }
    
    private void render(FencedCodeBlock node, NodeRendererContext context, HtmlWriter html) {
        HtmlRendererOptions htmlOptions = context.getHtmlOptions();
        //BasedSequence language = node.getInfoDelimitedByAny(htmlOptions.languageDelimiterSet);
        BasedSequence language = node.getInfoDelimitedByAny(CharPredicate.ANY_EOL);

        if (this.options.renderBlockPlantUML && language.startsWithIgnoreCase(options.plantUMLLanguage)) {
            
            String format = "png";
            Map<String, String> attrs = new HashMap<>();
            
            Matcher formatMatcher = FMT_PATTERN.matcher(language);
            if (formatMatcher.find()) {
                format = formatMatcher.group(2);
            }
            Matcher attrMatcher = ATTR_PATTERN.matcher(language);
            int find = 0;
            while (attrMatcher.find(find)) {
                attrs.put(attrMatcher.group(2), attrMatcher.group(4));
                find = attrMatcher.end();
            }
            String encode = deflaterCompressString(node.getContentChars().normalizeEOL());
            render(node, html, htmlOptions, format, attrs, encode, options.darkMode);
            
            if (options.bothModes) {
                render(node, html, htmlOptions, format, attrs, encode, !options.darkMode);
            }
            
        } else {
            context.delegateRender();
        }
    }

    public static String deflaterCompressString(String text) {
        if (StringUtils.isEmpty(text)) {
            return null;
        }

        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
        try (DeflaterOutputStream deflaterOutputStream = new DeflaterOutputStream(byteArrayOutputStream, new Deflater(Deflater.DEFAULT_COMPRESSION, true))) {
            deflaterOutputStream.write(text.getBytes(StandardCharsets.UTF_8));
            deflaterOutputStream.flush();
            deflaterOutputStream.finish();
        } catch (Exception e) {
            System.out.println(e);
            return null;
        }
        return PlantUMLEncode.getEncoder().encode(byteArrayOutputStream.toByteArray());
    }
    
    public static class Factory implements NodeRendererFactory {
        @NotNull
        @Override
        public NodeRenderer apply(@NotNull DataHolder options) {
            return new PlantUMLNodeRenderer(options);
        }
    }
}
