/*
 * Copyright 2005-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.springframework.xml.xsd.commons;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.ws.commons.schema.XmlSchema;
import org.apache.ws.commons.schema.XmlSchemaCollection;
import org.apache.ws.commons.schema.XmlSchemaExternal;
import org.apache.ws.commons.schema.XmlSchemaImport;
import org.apache.ws.commons.schema.XmlSchemaInclude;
import org.apache.ws.commons.schema.XmlSchemaObject;
import org.apache.ws.commons.schema.resolver.DefaultURIResolver;
import org.apache.ws.commons.schema.resolver.URIResolver;
import org.xml.sax.InputSource;

import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ResourceLoaderAware;
import org.springframework.core.io.Resource;
import org.springframework.core.io.ResourceLoader;
import org.springframework.core.io.UrlResource;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import org.springframework.xml.sax.SaxUtils;
import org.springframework.xml.validation.XmlValidator;
import org.springframework.xml.validation.XmlValidatorFactory;
import org.springframework.xml.xsd.XsdSchema;
import org.springframework.xml.xsd.XsdSchemaCollection;

/**
 * Implementation of the {@link XsdSchemaCollection} that uses Apache WS-Commons XML
 * Schema.
 * <p>
 * Setting the {@link #setInline(boolean) inline} flag to {@code true} will result in all
 * referenced schemas (included and imported) being merged into the referred schema. When
 * including the schemas into a WSDL, this greatly simplifies the deployment of the
 * schemas.
 *
 * @author Arjen Poutsma
 * @since 1.5.0
 * @see <a href="http://ws.apache.org/commons/XmlSchema/">Commons XML Schema</a>
 */
public class CommonsXsdSchemaCollection implements XsdSchemaCollection, InitializingBean, ResourceLoaderAware {

	private static final Log logger = LogFactory.getLog(CommonsXsdSchemaCollection.class);

	private final XmlSchemaCollection schemaCollection = new XmlSchemaCollection();

	private final List<XmlSchema> xmlSchemas = new ArrayList<>();

	private Resource[] xsdResources;

	private boolean inline = false;

	private URIResolver uriResolver = new ClasspathUriResolver();

	private ResourceLoader resourceLoader;

	/**
	 * Constructs a new, empty instance of the {@code CommonsXsdSchemaCollection}.
	 * <p>
	 * A subsequent call to the {@link #setXsds(Resource[])} is required.
	 */
	public CommonsXsdSchemaCollection() {
	}

	/**
	 * Constructs a new instance of the {@code CommonsXsdSchemaCollection} based on the
	 * given resources.
	 * @param resources the schema resources to load
	 */
	public CommonsXsdSchemaCollection(Resource... resources) {
		this.xsdResources = resources;
	}

	/**
	 * Sets the schema resources to be loaded.
	 * @param xsdResources the schema resources to be loaded
	 */
	public void setXsds(Resource... xsdResources) {
		this.xsdResources = xsdResources;
	}

	/**
	 * Defines whether included schemas should be inlined into the including schema.
	 * <p>
	 * Defaults to {@code false}.
	 */
	public void setInline(boolean inline) {
		this.inline = inline;
	}

	/**
	 * Sets the WS-Commons uri resolver to use when resolving (relative) schemas.
	 * <p>
	 * Default is an internal subclass of {@link DefaultURIResolver} which correctly
	 * handles schemas on the classpath.
	 */
	public void setUriResolver(URIResolver uriResolver) {
		Assert.notNull(uriResolver, "'uriResolver' must not be null");
		this.uriResolver = uriResolver;
	}

	@Override
	public void setResourceLoader(ResourceLoader resourceLoader) {
		this.resourceLoader = resourceLoader;
	}

	@Override
	public void afterPropertiesSet() throws IOException {
		Assert.notEmpty(this.xsdResources, "'xsds' must not be empty");

		this.schemaCollection.setSchemaResolver(this.uriResolver);

		Set<XmlSchema> processedIncludes = new HashSet<>();
		Set<XmlSchema> processedImports = new HashSet<>();

		for (Resource xsdResource : this.xsdResources) {
			Assert.isTrue(xsdResource.exists(), xsdResource + " does not exist");
			try {
				XmlSchema xmlSchema = this.schemaCollection.read(SaxUtils.createInputSource(xsdResource));
				this.xmlSchemas.add(xmlSchema);

				if (this.inline) {
					inlineIncludes(xmlSchema, processedIncludes, processedImports);
					findImports(xmlSchema, processedImports, processedIncludes);
				}
			}
			catch (Exception ex) {
				throw new CommonsXsdSchemaException("Schema [" + xsdResource + "] could not be loaded", ex);
			}
		}
		if (logger.isInfoEnabled()) {
			logger.info("Loaded " + StringUtils.arrayToCommaDelimitedString(this.xsdResources));
		}

	}

	@Override
	public XsdSchema[] getXsdSchemas() {
		XsdSchema[] result = new XsdSchema[this.xmlSchemas.size()];
		for (int i = 0; i < this.xmlSchemas.size(); i++) {
			XmlSchema xmlSchema = this.xmlSchemas.get(i);
			result[i] = new CommonsXsdSchema(xmlSchema, this.schemaCollection);
		}
		return result;
	}

	@Override
	public XmlValidator createValidator() {
		try {
			Resource[] resources = new Resource[this.xmlSchemas.size()];
			for (int i = this.xmlSchemas.size() - 1; i >= 0; i--) {
				XmlSchema xmlSchema = this.xmlSchemas.get(i);
				String sourceUri = xmlSchema.getSourceURI();
				if (StringUtils.hasLength(sourceUri)) {
					resources[i] = new UrlResource(sourceUri);
				}
			}
			return XmlValidatorFactory.createValidator(resources, XmlValidatorFactory.SCHEMA_W3C_XML);
		}
		catch (IOException ex) {
			throw new CommonsXsdSchemaException(ex.getMessage(), ex);
		}
	}

	private void inlineIncludes(XmlSchema schema, Set<XmlSchema> processedIncludes, Set<XmlSchema> processedImports) {
		processedIncludes.add(schema);

		List<XmlSchemaObject> schemaItems = schema.getItems();
		for (int i = 0; i < schemaItems.size(); i++) {
			XmlSchemaObject schemaObject = schemaItems.get(i);
			if (schemaObject instanceof XmlSchemaInclude) {
				XmlSchema includedSchema = ((XmlSchemaInclude) schemaObject).getSchema();
				if (!processedIncludes.contains(includedSchema)) {
					inlineIncludes(includedSchema, processedIncludes, processedImports);
					findImports(includedSchema, processedImports, processedIncludes);
					schemaItems.addAll(includedSchema.getItems());
				}
				// remove the <include/>
				schemaItems.remove(i);
				i--;
			}
		}
	}

	private void findImports(XmlSchema schema, Set<XmlSchema> processedImports, Set<XmlSchema> processedIncludes) {
		processedImports.add(schema);
		List<XmlSchemaExternal> externals = schema.getExternals();
		for (XmlSchemaExternal external : externals) {
			if (external instanceof XmlSchemaImport schemaImport) {
				XmlSchema importedSchema = schemaImport.getSchema();
				if (!"http://www.w3.org/XML/1998/namespace".equals(schemaImport.getNamespace())
						&& importedSchema != null && !processedImports.contains(importedSchema)) {
					inlineIncludes(importedSchema, processedIncludes, processedImports);
					findImports(importedSchema, processedImports, processedIncludes);
					this.xmlSchemas.add(importedSchema);
				}
				// remove the schemaLocation
				external.setSchemaLocation(null);
			}
		}
	}

	public String toString() {
		StringBuilder builder = new StringBuilder("CommonsXsdSchemaCollection");
		builder.append('{');
		for (int i = 0; i < this.xmlSchemas.size(); i++) {
			XmlSchema schema = this.xmlSchemas.get(i);
			builder.append(schema.getTargetNamespace());
			if (i < this.xmlSchemas.size() - 1) {
				builder.append(',');
			}
		}
		builder.append('}');
		return builder.toString();
	}

	private final class ClasspathUriResolver extends DefaultURIResolver {

		@Override
		public InputSource resolveEntity(String namespace, String schemaLocation, String baseUri) {
			if (CommonsXsdSchemaCollection.this.resourceLoader != null) {
				Resource resource = CommonsXsdSchemaCollection.this.resourceLoader.getResource(schemaLocation);
				if (resource.exists()) {
					return createInputSource(resource);
				}
				else if (StringUtils.hasLength(baseUri)) {
					// let's try and find it relative to the baseUri, see SWS-413
					try {
						Resource baseUriResource = new UrlResource(baseUri);
						resource = baseUriResource.createRelative(schemaLocation);
						if (resource.exists()) {
							return createInputSource(resource);
						}
					}
					catch (IOException ex) {
						// fall through
					}
				}
				// let's try and find it on the classpath, see SWS-362
				String classpathLocation = ResourceLoader.CLASSPATH_URL_PREFIX + "/" + schemaLocation;
				resource = CommonsXsdSchemaCollection.this.resourceLoader.getResource(classpathLocation);
				if (resource.exists()) {
					return createInputSource(resource);
				}
			}
			return super.resolveEntity(namespace, schemaLocation, baseUri);
		}

		private InputSource createInputSource(Resource resource) {
			try {
				return SaxUtils.createInputSource(resource);
			}
			catch (IOException ex) {
				throw new CommonsXsdSchemaException("Could not resolve location", ex);
			}
		}

	}

}
