package jnpf.fileStorage.dms;

import jnpf.exception.DataException;
import org.springframework.stereotype.Service;
import org.springframework.web.multipart.MultipartFile;
import org.w3c.dom.*;

import javax.xml.XMLConstants;
import javax.xml.parsers.DocumentBuilderFactory;
import javax.xml.transform.TransformerFactory;
import javax.xml.transform.dom.DOMSource;
import javax.xml.transform.stream.StreamResult;
import java.io.*;
import java.net.URI;
import java.nio.file.Files;
import java.util.*;
import java.util.zip.*;

/** Cleans a new formal copy, rejecting revisions and leaving the source review attachment intact. */
@Service
public class DmsDocxCommentCleaner {
    private static final int MAX_PARTS = 4096;
    private static final int MAX_BYTES = 50 * 1024 * 1024;
    private static final int MAX_XML = 16 * 1024 * 1024;
    private static final String CT = "http://schemas.openxmlformats.org/package/2006/content-types";
    private static final String REL = "http://schemas.openxmlformats.org/package/2006/relationships";
    private static final Set<String> COMMENT_TYPES = Set.of("comments", "commentsExtended", "commentsIds",
            "commentsExtensible", "people");
    private static final Set<String> MARKERS = Set.of("commentRangeStart", "commentRangeEnd", "commentReference");

    public MultipartFile clean(MultipartFile source) {
        String name = source.getOriginalFilename();
        if (name == null || !name.toLowerCase(Locale.ROOT).endsWith(".docx")) return source;
        try {
            Map<String, byte[]> parts = read(source);
            Document types = parse(required(parts, "[Content_Types].xml"));
            Map<String, Document> xml = new LinkedHashMap<>();
            Set<String> removed = new HashSet<>();
            boolean main = false;
            NodeList overrides = types.getElementsByTagNameNS(CT, "Override");
            for (int i = 0; i < overrides.getLength(); i++) {
                Element type = (Element) overrides.item(i);
                String contentType = type.getAttribute("ContentType");
                String part = target("", type.getAttribute("PartName"));
                if (commentContentType(contentType)) removed.add(part);
                if (contentType.equals("application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml")) {
                    required(parts, part);
                    main = true;
                }
            }
            if (!main) throw new IOException("Missing DOCX main part");
            // Parse every XML part securely. Relationship discovery also handles nonstandard comment part names.
            for (Map.Entry<String, byte[]> part : parts.entrySet()) {
                if (part.getKey().endsWith(".xml") || part.getKey().endsWith(".rels")) {
                    Document doc = parse(part.getValue());
                    xml.put(part.getKey(), doc);
                    if (part.getKey().endsWith(".rels")) {
                        NodeList links = doc.getElementsByTagNameNS(REL, "Relationship");
                        for (int i = 0; i < links.getLength(); i++) {
                            Element link = (Element) links.item(i);
                            if (commentRelationship(link.getAttribute("Type")) && !external(link))
                                removed.add(target(owner(part.getKey()), link.getAttribute("Target")));
                        }
                    }
                }
            }
            // Remove comment-owned resources only if no retained part references them (shared images stay).
            boolean expanded;
            do {
                expanded = false;
                Set<String> candidates = new HashSet<>();
                Set<String> retainedTargets = new HashSet<>();
                for (Map.Entry<String, Document> part : xml.entrySet()) {
                    if (!part.getKey().endsWith(".rels")) continue;
                    String owner = owner(part.getKey());
                    NodeList links = part.getValue().getElementsByTagNameNS(REL, "Relationship");
                    for (int i = 0; i < links.getLength(); i++) {
                        Element link = (Element) links.item(i);
                        if (external(link)) continue;
                        String dest = target(owner, link.getAttribute("Target"));
                        if (removed.contains(owner)) candidates.add(dest);
                        else if (!commentRelationship(link.getAttribute("Type"))) retainedTargets.add(dest);
                    }
                }
                candidates.removeAll(retainedTargets);
                if (removed.addAll(candidates)) expanded = true;
            } while (expanded);
            for (String part : new ArrayList<>(removed)) removed.add(rels(part));
            boolean changed = parts.keySet().removeAll(removed);
            for (Map.Entry<String, Document> part : xml.entrySet()) {
                if (removed.contains(part.getKey())) continue;
                Document doc = part.getValue();
                boolean edited = DmsDocxRevisionRejector.reject(doc);
                edited |= removeMarkers(doc);
                if (part.getKey().endsWith(".rels")) {
                    NodeList links = doc.getElementsByTagNameNS(REL, "Relationship");
                    for (int i = links.getLength() - 1; i >= 0; i--) {
                        Element link = (Element) links.item(i);
                        if (commentRelationship(link.getAttribute("Type")) || !external(link)
                                && removed.contains(target(owner(part.getKey()), link.getAttribute("Target")))) {
                            link.getParentNode().removeChild(link); edited = true;
                        }
                    }
                }
                if (part.getKey().equals("[Content_Types].xml")) {
                    NodeList entries = doc.getElementsByTagNameNS(CT, "Override");
                    for (int i = entries.getLength() - 1; i >= 0; i--) {
                        Element entry = (Element) entries.item(i);
                        if (removed.contains(target("", entry.getAttribute("PartName")))) {
                            entry.getParentNode().removeChild(entry); edited = true;
                        }
                    }
                }
                if (edited) { parts.put(part.getKey(), serialize(doc)); changed = true; }
            }
            if (!changed) return source;
            // Package signatures cannot survive byte changes. Business signature-page content is unaffected.
            if (parts.keySet().stream().anyMatch(p -> p.startsWith("_xmlsignatures/")))
                throw new IOException("Digitally signed package cannot be rewritten");
            ByteArrayOutputStream output = new ByteArrayOutputStream();
            try (ZipOutputStream zip = new ZipOutputStream(output)) {
                for (Map.Entry<String, byte[]> part : parts.entrySet()) {
                    zip.putNextEntry(new ZipEntry(part.getKey())); zip.write(part.getValue()); zip.closeEntry();
                }
            }
            return new CleanedFile(source.getName(), name, source.getContentType(), output.toByteArray());
        } catch (Exception e) {
            throw new DataException("DOCX修订及批注清理失败，正式文件未存储：" + e.getMessage());
        }
    }

    private static boolean commentContentType(String type) {
        for (String name : COMMENT_TYPES) {
            if (type.equals("application/vnd.openxmlformats-officedocument.wordprocessingml." + name + "+xml")
                    || type.equals("application/vnd.ms-word." + name + "+xml")) return true;
        }
        return false;
    }
    private static boolean commentRelationship(String type) {
        return (type.startsWith("http://schemas.openxmlformats.org/officeDocument/2006/relationships/")
                || type.startsWith("http://purl.oclc.org/ooxml/officeDocument/relationships/")
                || type.startsWith("http://schemas.microsoft.com/office/"))
                && COMMENT_TYPES.contains(type.substring(type.lastIndexOf('/') + 1));
    }
    private static boolean external(Element link) { return "External".equals(link.getAttribute("TargetMode")); }
    private static String owner(String rels) throws IOException {
        if (rels.equals("_rels/.rels")) return "";
        int index = rels.lastIndexOf("/_rels/");
        if (index < 0 || !rels.endsWith(".rels")) throw new IOException("Invalid relationship part");
        return rels.substring(0, index + 1) + rels.substring(index + 7, rels.length() - 5);
    }
    private static String rels(String part) {
        int index = part.lastIndexOf('/');
        return part.substring(0, index + 1) + "_rels/" + part.substring(index + 1) + ".rels";
    }
    private static String target(String owner, String value) throws IOException {
        try {
            URI resolved = URI.create("/" + owner).resolve(value).normalize();
            String path = resolved.getPath();
            if (resolved.isAbsolute() || resolved.getAuthority() != null || path == null || !path.startsWith("/")
                    || path.contains("/../")) throw new IOException("Invalid package target");
            return path.substring(1);
        } catch (IllegalArgumentException e) { throw new IOException("Invalid package target", e); }
    }
    private static byte[] required(Map<String, byte[]> parts, String name) throws IOException {
        byte[] data = parts.get(name);
        if (data == null) throw new IOException("Missing package part: " + name);
        return data;
    }
    private static Map<String, byte[]> read(MultipartFile source) throws IOException {
        Map<String, byte[]> parts = new LinkedHashMap<>();
        int total = 0, count = 0;
        try (ZipInputStream zip = new ZipInputStream(source.getInputStream())) {
            ZipEntry entry;
            while ((entry = zip.getNextEntry()) != null) {
                String name = entry.getName();
                if (++count > MAX_PARTS || name.startsWith("/") || name.contains("\\")
                        || Arrays.asList(name.split("/")).contains("..") || parts.containsKey(name))
                    throw new IOException("Unsafe or oversized DOCX package");
                if (entry.isDirectory()) continue;
                ByteArrayOutputStream bytes = new ByteArrayOutputStream();
                byte[] buffer = new byte[8192]; int read;
                while ((read = zip.read(buffer)) != -1) {
                    total += read;
                    if (total > MAX_BYTES) throw new IOException("DOCX exceeds 50MB expanded limit");
                    bytes.write(buffer, 0, read);
                }
                parts.put(name, bytes.toByteArray());
            }
        }
        return parts;
    }
    private static Document parse(byte[] bytes) throws Exception {
        if (bytes.length > MAX_XML) throw new IOException("DOCX XML part exceeds 16MB limit");
        DocumentBuilderFactory factory = DocumentBuilderFactory.newDefaultInstance(); factory.setNamespaceAware(true);
        factory.setFeature(XMLConstants.FEATURE_SECURE_PROCESSING, true);
        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);
        factory.setAttribute(XMLConstants.ACCESS_EXTERNAL_DTD, ""); factory.setAttribute(XMLConstants.ACCESS_EXTERNAL_SCHEMA, "");
        return factory.newDocumentBuilder().parse(new ByteArrayInputStream(bytes));
    }
    private static boolean removeMarkers(Document doc) {
        NodeList nodes = doc.getElementsByTagNameNS("*", "*"); boolean changed = false;
        for (int i = nodes.getLength() - 1; i >= 0; i--) {
            Node node = nodes.item(i);
            if (MARKERS.contains(node.getLocalName()) && ("http://schemas.openxmlformats.org/wordprocessingml/2006/main".equals(node.getNamespaceURI())
                    || "http://purl.oclc.org/ooxml/wordprocessingml/main".equals(node.getNamespaceURI()))) {
                node.getParentNode().removeChild(node); changed = true;
            }
        }
        return changed;
    }
    private static byte[] serialize(Document doc) throws Exception {
        TransformerFactory factory = TransformerFactory.newDefaultInstance();
        factory.setFeature(XMLConstants.FEATURE_SECURE_PROCESSING, true);
        factory.setAttribute(XMLConstants.ACCESS_EXTERNAL_DTD, ""); factory.setAttribute(XMLConstants.ACCESS_EXTERNAL_STYLESHEET, "");
        ByteArrayOutputStream output = new ByteArrayOutputStream();
        factory.newTransformer().transform(new DOMSource(doc), new StreamResult(output)); return output.toByteArray();
    }
    private record CleanedFile(String name, String originalFilename, String contentType, byte[] bytes) implements MultipartFile {
        public String getName() { return name; }
        public String getOriginalFilename() { return originalFilename; }
        public String getContentType() { return contentType; }
        public boolean isEmpty() { return bytes.length == 0; }
        public long getSize() { return bytes.length; }
        public byte[] getBytes() { return bytes.clone(); }
        public InputStream getInputStream() { return new ByteArrayInputStream(bytes); }
        public void transferTo(File dest) throws IOException { Files.write(dest.toPath(), bytes); }
    }
}