/*
 * Copyright 2021-2025 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.opentest4j.reporting.events.api;

import org.opentest4j.reporting.schema.Namespace;
import org.opentest4j.reporting.schema.QualifiedName;

import javax.xml.stream.XMLOutputFactory;
import javax.xml.stream.XMLStreamException;
import javax.xml.stream.XMLStreamWriter;
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.stax.StAXResult;
import java.io.BufferedWriter;
import java.io.IOException;
import java.io.Writer;
import java.util.Optional;
import java.util.function.Consumer;

import static java.util.stream.Collectors.joining;
import static javax.xml.XMLConstants.FEATURE_SECURE_PROCESSING;
import static javax.xml.XMLConstants.W3C_XML_SCHEMA_INSTANCE_NS_URI;
import static org.opentest4j.reporting.events.api.NamespaceRegistry.XSI_PREFIX;

class DefaultDocumentWriter<T extends Element<T>> implements DocumentWriter<T> {

	private final AttributeValueEscapingWriter out;
	private final XMLStreamWriter xmlWriter;
	private final Context context;
	private final Transformer transformer;
	private final SingleElementXMLStreamWriter elementXmlWriter;

	DefaultDocumentWriter(QualifiedName rootElementName, NamespaceRegistry namespaceRegistry, Writer writer)
			throws Exception {
		XMLOutputFactory factory = XMLOutputFactory.newInstance();
		out = new AttributeValueEscapingWriter(buffered(writer));
		xmlWriter = new EscapingXMLStreamWriter(out, factory.createXMLStreamWriter(out));
		xmlWriter.writeStartDocument();
		writeNewLine();
		writeStartElement(rootElementName, namespaceRegistry);
		writeNamespaces(namespaceRegistry);
		writeSchemaLocations(namespaceRegistry);
		xmlWriter.writeCharacters("");

		elementXmlWriter = new SingleElementXMLStreamWriter(xmlWriter, namespaceRegistry);

		context = Context.create(rootElementName, namespaceRegistry);
		TransformerFactory transformerFactory = TransformerFactory.newInstance();
		transformerFactory.setFeature(FEATURE_SECURE_PROCESSING, true);
		transformer = transformerFactory.newTransformer();
		transformer.setOutputProperty(OutputKeys.METHOD, "xml");
	}

	private void writeStartElement(QualifiedName rootElement, NamespaceRegistry namespaceRegistry)
			throws XMLStreamException {
		Optional<String> prefix = namespaceRegistry.getPrefix(rootElement.getNamespace());
		if (prefix.isPresent()) {
			xmlWriter.writeStartElement(prefix.get(), rootElement.getSimpleName(), rootElement.getNamespace().getUri());
		}
		else {
			xmlWriter.writeStartElement(rootElement.getSimpleName());
		}
	}

	private void writeNamespaces(NamespaceRegistry namespaceRegistry) throws XMLStreamException {
		xmlWriter.writeDefaultNamespace(namespaceRegistry.getDefaultNamespace().getUri());
		for (Namespace namespace : namespaceRegistry.getAdditionalNamespaces()) {
			Optional<String> prefix = namespaceRegistry.getPrefix(namespace);
			if (prefix.isPresent()) {
				xmlWriter.writeNamespace(prefix.get(), namespace.getUri());
			}
		}
	}

	private void writeSchemaLocations(NamespaceRegistry namespaceRegistry) throws XMLStreamException {
		if (!namespaceRegistry.getSchemaLocations().isEmpty()) {
			String value = namespaceRegistry.getSchemaLocations().entrySet().stream() //
					.map(e -> e.getKey().getUri() + " " + e.getValue()) //
					.collect(joining(" "));
			xmlWriter.writeAttribute(XSI_PREFIX, W3C_XML_SCHEMA_INSTANCE_NS_URI, "schemaLocation", value);
		}
	}

	private static BufferedWriter buffered(Writer writer) {
		return writer instanceof BufferedWriter //
				? (BufferedWriter) writer //
				: new BufferedWriter(writer);
	}

	@Override
	public synchronized <C extends ChildElement<T, ? super C>> DocumentWriter<T> append(Factory<C> creator,
			Consumer<? super C> configurer) {
		C event = creator.createAndConfigure(context, configurer);
		try {
			writeNewLine();
			org.w3c.dom.Element domElement = event.getDomElement();
			domElement.normalize();
			transformer.transform(new DOMSource(domElement), new StAXResult(elementXmlWriter));
		}
		catch (Exception e) {
			throw new RuntimeException("Failed to write event: " + event, e);
		}
		return this;
	}

	@Override
	public void close() throws IOException {
		try {
			writeNewLine();
			xmlWriter.writeEndDocument();
			xmlWriter.close();
		}
		catch (XMLStreamException e) {
			throw new IOException("Failed to write XML", e);
		}
		finally {
			out.close();
		}
	}

	private void writeNewLine() throws IOException, XMLStreamException {
		xmlWriter.flush();
		out.newLine();
		out.flush();
	}

}
