/*
 * Copyright 2021-2024 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.namespace.NamespaceContext;
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.Element.CDATA_WORKAROUND_COMMENT;
import static org.opentest4j.reporting.events.api.NamespaceRegistry.XSI_PREFIX;

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

	private final BufferedWriter bufferedWriter;
	private final XMLStreamWriter xmlWriter;
	private final Context context;
	private final Transformer transformer;
	private final SingleElementStreamWriter elementXmlWriter;

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

		elementXmlWriter = new SingleElementStreamWriter(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 {
			bufferedWriter.close();
		}
	}

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

	private static class SingleElementStreamWriter implements XMLStreamWriter {

		private final XMLStreamWriter delegate;
		private final NamespaceRegistry namespaceRegistry;

		public SingleElementStreamWriter(XMLStreamWriter delegate, NamespaceRegistry namespaceRegistry) {
			this.delegate = delegate;
			this.namespaceRegistry = namespaceRegistry;
		}

		@Override
		public void writeStartElement(String localName) throws XMLStreamException {
			delegate.writeStartElement(localName);
		}

		@Override
		public void writeStartElement(String namespaceURI, String localName) throws XMLStreamException {
			delegate.writeStartElement(namespaceURI, localName);
		}

		@Override
		public void writeStartElement(String prefix, String localName, String namespaceURI) throws XMLStreamException {
			delegate.writeStartElement(prefix, localName, namespaceURI);
		}

		@Override
		public void writeEmptyElement(String namespaceURI, String localName) throws XMLStreamException {
			delegate.writeEmptyElement(namespaceURI, localName);
		}

		@Override
		public void writeEmptyElement(String prefix, String localName, String namespaceURI) throws XMLStreamException {
			delegate.writeEmptyElement(prefix, localName, namespaceURI);
		}

		@Override
		public void writeEmptyElement(String localName) throws XMLStreamException {
			delegate.writeEmptyElement(localName);
		}

		@Override
		public void writeEndElement() throws XMLStreamException {
			delegate.writeEndElement();
		}

		@Override
		public void writeEndDocument() {
		}

		@Override
		public void close() throws XMLStreamException {
			delegate.close();
		}

		@Override
		public void flush() throws XMLStreamException {
			delegate.flush();
		}

		@Override
		public void writeAttribute(String localName, String value) throws XMLStreamException {
			delegate.writeAttribute(localName, value);
		}

		@Override
		public void writeAttribute(String prefix, String namespaceURI, String localName, String value)
				throws XMLStreamException {
			delegate.writeAttribute(prefix, namespaceURI, localName, value);
		}

		@Override
		public void writeAttribute(String namespaceURI, String localName, String value) throws XMLStreamException {
			delegate.writeAttribute(namespaceURI, localName, value);
		}

		@Override
		public void writeNamespace(String prefix, String namespaceURI) throws XMLStreamException {
			if (!namespaceRegistry.containsUri(namespaceURI)) {
				delegate.writeNamespace(prefix, namespaceURI);
			}
		}

		@Override
		public void writeDefaultNamespace(String namespaceURI) {
		}

		@Override
		public void writeComment(String data) throws XMLStreamException {
			// Remove bogus comment that avoids CDATA sections from being merged
			if (!CDATA_WORKAROUND_COMMENT.equals(data)) {
				delegate.writeComment(data);
			}
		}

		@Override
		public void writeProcessingInstruction(String target) throws XMLStreamException {
			delegate.writeProcessingInstruction(target);
		}

		@Override
		public void writeProcessingInstruction(String target, String data) throws XMLStreamException {
			delegate.writeProcessingInstruction(target, data);
		}

		@Override
		public void writeCData(String data) throws XMLStreamException {
			delegate.writeCData(data);
		}

		@Override
		public void writeDTD(String dtd) throws XMLStreamException {
			delegate.writeDTD(dtd);
		}

		@Override
		public void writeEntityRef(String name) throws XMLStreamException {
			delegate.writeEntityRef(name);
		}

		@Override
		public void writeStartDocument() {
		}

		@Override
		public void writeStartDocument(String version) {
		}

		@Override
		public void writeStartDocument(String encoding, String version) {
		}

		@Override
		public void writeCharacters(String text) throws XMLStreamException {
			delegate.writeCharacters(text);
		}

		@Override
		public void writeCharacters(char[] text, int start, int len) throws XMLStreamException {
			delegate.writeCharacters(text, start, len);
		}

		@Override
		public String getPrefix(String uri) throws XMLStreamException {
			return delegate.getPrefix(uri);
		}

		@Override
		public void setPrefix(String prefix, String uri) throws XMLStreamException {
			delegate.setPrefix(prefix, uri);
		}

		@Override
		public void setDefaultNamespace(String uri) throws XMLStreamException {
			delegate.setDefaultNamespace(uri);
		}

		@Override
		public void setNamespaceContext(NamespaceContext context) throws XMLStreamException {
			delegate.setNamespaceContext(context);
		}

		@Override
		public NamespaceContext getNamespaceContext() {
			return delegate.getNamespaceContext();
		}

		@Override
		public Object getProperty(String name) throws IllegalArgumentException {
			return delegate.getProperty(name);
		}
	}
}
