package jnpf.elnService.template.docx;

import jnpf.exception.DataException;
import org.w3c.dom.Document;
import org.w3c.dom.Element;
import org.w3c.dom.NamedNodeMap;
import org.w3c.dom.Node;
import org.w3c.dom.NodeList;
import org.xml.sax.InputSource;
import org.xml.sax.SAXException;
import org.xml.sax.SAXParseException;
import org.xml.sax.helpers.DefaultHandler;

import javax.xml.XMLConstants;
import javax.xml.parsers.DocumentBuilder;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.transform.OutputKeys;
import javax.xml.transform.Transformer;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Set;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;
import java.util.zip.ZipOutputStream;

/** 受限 DOCX writer：只替换按稳定 Tag 唯一定位的文本内容控件内容。 */
public class DocxContentControlWriter {

    private static final String WORD_NAMESPACE =
            "http://schemas.openxmlformats.org/wordprocessingml/2006/main";
    private static final int MAX_PART_BYTES = 10 * 1024 * 1024;
    private static final int MAX_TOTAL_BYTES = 50 * 1024 * 1024;
    private static final int MAX_ENTRIES = 10000;
    private static final int MAX_COMPRESSION_RATIO = 100;

    public byte[] write(byte[] source, Map<String, String> valuesByTag) {
        if (source == null || source.length == 0 || source.length > MAX_TOTAL_BYTES) {
            throw new DataException("DOCX 文件大小无效");
        }
        Map<String, String> requested = requiredValues(valuesByTag);
        Map<String, Integer> matches = new LinkedHashMap<>();
        for (String tag : requested.keySet()) matches.put(tag, 0);
        Set<String> entryNames = new HashSet<>();
        boolean documentFound = false;
        int entries = 0;
        long extracted = 0;
        try (ZipInputStream input = new ZipInputStream(new ByteArrayInputStream(source));
             ByteArrayOutputStream bytes = new ByteArrayOutputStream(source.length);
             ZipOutputStream output = new ZipOutputStream(bytes)) {
            ZipEntry entry;
            byte[] buffer = new byte[8192];
            while ((entry = input.getNextEntry()) != null) {
                if (++entries > MAX_ENTRIES) {
                    throw new DataException("DOCX ZIP 条目超过 10000 个限制");
                }
                String name = entry.getName();
                requireSafeEntryName(name);
                if (!entryNames.add(name)) throw new DataException("DOCX ZIP 条目名称重复");
                if ("word/document.xml".equals(name)) documentFound = true;
                ByteArrayOutputStream part = new ByteArrayOutputStream();
                long entryBytes = 0;
                int read;
                while ((read = input.read(buffer)) != -1) {
                    entryBytes += read;
                    extracted += read;
                    if (entryBytes > MAX_PART_BYTES && isTargetPart(name)) {
                        throw new DataException("DOCX XML 部件超过 10MB 限制");
                    }
                    if (extracted > MAX_TOTAL_BYTES) {
                        throw new DataException("DOCX 解压后超过 50MB 限制");
                    }
                    part.write(buffer, 0, read);
                }
                rejectExcessiveCompression(entry, entryBytes);
                byte[] content = part.toByteArray();
                if (isTargetPart(name) && content.length > 0) {
                    content = writePart(content, requested, matches);
                }
                ZipEntry replacement = new ZipEntry(name);
                replacement.setTime(entry.getTime());
                if (entry.getComment() != null) replacement.setComment(entry.getComment());
                if (entry.getExtra() != null) replacement.setExtra(entry.getExtra());
                output.putNextEntry(replacement);
                output.write(content);
                output.closeEntry();
                input.closeEntry();
            }
            output.finish();
            if (!documentFound) throw new DataException("DOCX 缺少 word/document.xml");
            for (Map.Entry<String, Integer> match : matches.entrySet()) {
                if (match.getValue() == 0) {
                    throw new DataException("DOCX 缺少目标字段控件: " + match.getKey());
                }
                if (match.getValue() > 1) {
                    throw new DataException("DOCX 目标字段 Tag 重复: " + match.getKey());
                }
            }
            byte[] result = bytes.toByteArray();
            if (result.length == 0 || result.length > MAX_TOTAL_BYTES) {
                throw new DataException("DOCX 写入后大小无效");
            }
            return result;
        } catch (DataException e) {
            throw e;
        } catch (Exception e) {
            throw new DataException("DOCX 内容控件写入失败: " + e.getMessage());
        }
    }

    private static byte[] writePart(byte[] source, Map<String, String> requested,
                                    Map<String, Integer> matches) throws Exception {
        DocumentBuilderFactory factory = DocumentBuilderFactory.newInstance();
        factory.setNamespaceAware(true);
        factory.setXIncludeAware(false);
        factory.setExpandEntityReferences(false);
        factory.setFeature("http://apache.org/xml/features/disallow-doctype-decl", true);
        factory.setFeature("http://xml.org/sax/features/external-general-entities", false);
        factory.setFeature("http://xml.org/sax/features/external-parameter-entities", false);
        DocumentBuilder builder = factory.newDocumentBuilder();
        builder.setErrorHandler(new DefaultHandler() {
            @Override
            public void error(SAXParseException exception) throws SAXException {
                throw exception;
            }

            @Override
            public void fatalError(SAXParseException exception) throws SAXException {
                throw exception;
            }
        });
        Document document = builder.parse(
                new InputSource(new ByteArrayInputStream(source)));
        NodeList controls = document.getElementsByTagNameNS(WORD_NAMESPACE, "sdt");
        boolean changed = false;
        for (int index = 0; index < controls.getLength(); index++) {
            Element control = (Element) controls.item(index);
            Element properties = directChild(control, "sdtPr");
            Element content = directChild(control, "sdtContent");
            String tag = tag(properties);
            if (tag == null || !requested.containsKey(tag)) continue;
            matches.put(tag, matches.get(tag) + 1);
            if (content == null) throw new DataException("DOCX 目标字段控件缺少内容区: " + tag);
            replaceText(document, control, content, requested.get(tag));
            changed = true;
        }
        if (!changed) return source;
        TransformerFactory transformerFactory = TransformerFactory.newInstance();
        transformerFactory.setFeature(XMLConstants.FEATURE_SECURE_PROCESSING, true);
        Transformer transformer = transformerFactory.newTransformer();
        transformer.setOutputProperty(OutputKeys.ENCODING, "UTF-8");
        transformer.setOutputProperty(OutputKeys.OMIT_XML_DECLARATION, "no");
        ByteArrayOutputStream output = new ByteArrayOutputStream(source.length);
        transformer.transform(new DOMSource(document), new StreamResult(output));
        return output.toByteArray();
    }

    private static void replaceText(Document document, Element control,
                                    Element content, String value) {
        Element oldParagraph = firstDescendant(content, "p");
        Element oldRun = firstDescendant(content, "r");
        while (content.hasChildNodes()) content.removeChild(content.getFirstChild());
        boolean inline = hasAncestor(control, "p");
        if (!inline) {
            Element paragraph = document.createElementNS(WORD_NAMESPACE, "w:p");
            copyAttributes(oldParagraph, paragraph);
            Element paragraphProperties = oldParagraph == null
                    ? null : directChild(oldParagraph, "pPr");
            if (paragraphProperties != null) {
                paragraph.appendChild(paragraphProperties.cloneNode(true));
            }
            paragraph.appendChild(run(document, oldRun, value));
            content.appendChild(paragraph);
        } else {
            content.appendChild(run(document, oldRun, value));
        }
    }

    private static Element run(Document document, Element oldRun, String value) {
        Element run = document.createElementNS(WORD_NAMESPACE, "w:r");
        copyAttributes(oldRun, run);
        Element runProperties = oldRun == null ? null : directChild(oldRun, "rPr");
        if (runProperties != null) run.appendChild(runProperties.cloneNode(true));
        Element text = document.createElementNS(WORD_NAMESPACE, "w:t");
        text.setAttributeNS(XMLConstants.XML_NS_URI, "xml:space", "preserve");
        text.setTextContent(value);
        run.appendChild(text);
        return run;
    }

    private static String tag(Element properties) {
        if (properties == null) return null;
        NodeList tags = properties.getElementsByTagNameNS(WORD_NAMESPACE, "tag");
        if (tags.getLength() == 0) return null;
        Element tag = (Element) tags.item(0);
        String value = tag.getAttributeNS(WORD_NAMESPACE, "val");
        if (value.isEmpty()) value = tag.getAttribute("val");
        return value.isEmpty() ? null : value;
    }

    private static Element directChild(Element parent, String localName) {
        if (parent == null) return null;
        for (Node child = parent.getFirstChild(); child != null; child = child.getNextSibling()) {
            if (child.getNodeType() == Node.ELEMENT_NODE
                    && WORD_NAMESPACE.equals(child.getNamespaceURI())
                    && localName.equals(child.getLocalName())) return (Element) child;
        }
        return null;
    }

    private static Element firstDescendant(Element parent, String localName) {
        if (parent == null) return null;
        NodeList result = parent.getElementsByTagNameNS(WORD_NAMESPACE, localName);
        return result.getLength() == 0 ? null : (Element) result.item(0);
    }

    private static boolean hasAncestor(Element element, String localName) {
        for (Node parent = element.getParentNode(); parent != null; parent = parent.getParentNode()) {
            if (parent.getNodeType() == Node.ELEMENT_NODE
                    && WORD_NAMESPACE.equals(parent.getNamespaceURI())
                    && localName.equals(parent.getLocalName())) return true;
        }
        return false;
    }

    private static void copyAttributes(Element source, Element target) {
        if (source == null) return;
        NamedNodeMap attributes = source.getAttributes();
        for (int index = 0; index < attributes.getLength(); index++) {
            Node attribute = attributes.item(index);
            target.setAttributeNS(attribute.getNamespaceURI(), attribute.getNodeName(),
                    attribute.getNodeValue());
        }
    }

    private static Map<String, String> requiredValues(Map<String, String> valuesByTag) {
        if (valuesByTag == null || valuesByTag.isEmpty()) {
            throw new DataException("DOCX 目标字段控件不能为空");
        }
        Map<String, String> result = new LinkedHashMap<>();
        for (Map.Entry<String, String> entry : valuesByTag.entrySet()) {
            String tag = entry.getKey() == null ? null : entry.getKey().trim();
            if (tag == null || (!tag.startsWith("eln.field.") && !tag.startsWith("eln.ref."))) {
                throw new DataException("DOCX 目标字段 Tag 无效");
            }
            if (entry.getValue() == null) throw new DataException("DOCX 目标字段值不能为空");
            result.put(tag, entry.getValue());
        }
        return result;
    }

    private static boolean isTargetPart(String name) {
        return "word/document.xml".equals(name)
                || "word/footnotes.xml".equals(name)
                || "word/endnotes.xml".equals(name)
                || isNumberedPart(name, "word/header")
                || isNumberedPart(name, "word/footer");
    }

    private static boolean isNumberedPart(String name, String prefix) {
        return name.startsWith(prefix) && name.endsWith(".xml")
                && name.indexOf('/', prefix.length()) < 0;
    }

    private static void requireSafeEntryName(String name) {
        if (name == null || name.isEmpty() || name.startsWith("/") || name.startsWith("\\")
                || name.indexOf('\\') >= 0) {
            throw new DataException("DOCX ZIP 条目路径无效");
        }
        String[] segments = name.split("/", -1);
        for (String segment : segments) {
            if ("..".equals(segment) || ".".equals(segment)) {
                throw new DataException("DOCX ZIP 条目不允许路径跳转");
            }
        }
    }

    private static void rejectExcessiveCompression(ZipEntry entry, long uncompressedBytes) {
        long compressedBytes = entry.getCompressedSize();
        if (uncompressedBytes > 0 && compressedBytes == 0) {
            throw new DataException("DOCX ZIP 条目压缩比无效");
        }
        if (compressedBytes > 0 && uncompressedBytes > compressedBytes * MAX_COMPRESSION_RATIO) {
            throw new DataException("DOCX ZIP 条目解压比超过 100:1 限制");
        }
    }
}