import Plugin from '@ckeditor/ckeditor5-core/src/plugin';

export class KeepCustomClasses extends Plugin {
	init() {
		const editor = this.editor;

		// Define on which elements the CSS classes should be preserved:
		setupCustomClassConversion( 'h2', 'heading1', editor );
		setupCustomClassConversion( 'h3', 'heading2', editor );
		setupCustomClassConversion( 'h4', 'heading3', editor );
		setupCustomClassConversion( 'p', 'paragraph', editor );

		// Define custom attributes that should be preserved.
		setupCustomAttributeConversion( 'h2', 'heading1', 'id', editor );
		setupCustomAttributeConversion( 'h3', 'heading2', 'id', editor );
		setupCustomAttributeConversion( 'h4', 'heading3', 'id', editor );
		setupCustomAttributeConversion( 'p', 'paragraph', 'id', editor );

		function setupCustomClassConversion( viewElementName, modelElementName, editor ) {
			// The 'customClass' attribute will store custom classes from the data in the model so schema definitions allow this attribute.
			editor.model.schema.extend( modelElementName, { allowAttributes: [ 'customClass' ] } );

			// Define upcast converters for the <img> and <table> elements with a
			// "low" priority so they are run after the default converters.
			editor.conversion.for( 'upcast' ).add( upcastCustomClasses( viewElementName ),
				{ priority: 'low' } );

			// Define downcast converters for a model element with a
			// "low" priority so they are run after the default converters.
			// Use `downcastCustomClassesToFigure` if you'd like to keep your classes on
			// <figure> element or `downcastCustomClassesToChild` if you'd like to keep your
			// classes on a <figure> child element, i.e. <img>.
			editor.conversion.for( 'downcast' ).add( downcastCustomClassesToFigure( modelElementName ),
				{ priority: 'low' } );
		}

		function setupCustomAttributeConversion( viewElementName, modelElementName, viewAttribute, editor ) {
			// Extend the schema to store an attribute in the model.
			const modelAttribute = `custom${ viewAttribute }`;

			editor.model.schema.extend( modelElementName, { allowAttributes: [ modelAttribute ] } );

			editor.conversion.for( 'upcast' ).add( upcastAttribute( viewElementName, viewAttribute, modelAttribute ) );
			editor.conversion.for( 'downcast' ).add( downcastAttribute( modelElementName, viewElementName,
				viewAttribute, modelAttribute ) );
		}

		function upcastCustomClasses( elementName ) {
			return dispatcher => dispatcher.on( `element:${ elementName }`, ( evt, data, conversionApi ) => {
				const viewItem = data.viewItem;
				const modelRange = data.modelRange;

				const modelElement = modelRange && modelRange.start.nodeAfter;

				if ( !modelElement ) {
					return;
				}

				// The upcast conversion picks up classes from the base element and from the <figure> element so it should be extensible.
				const currentAttributeValue = modelElement.getAttribute( 'customClass' ) || [];

				currentAttributeValue.push( ...viewItem.getClassNames() );

				conversionApi.writer.setAttribute( 'customClass', currentAttributeValue, modelElement );
			} );
		}

		function downcastCustomClassesToFigure( modelElementName ) {
			return dispatcher => dispatcher.on( `insert:${ modelElementName }`, ( evt, data, conversionApi ) => {
				const modelElement = data.item;

				const viewFigure = conversionApi.mapper.toViewElement( modelElement );

				if ( !viewFigure ) {
					return;
				}

				// The code below assumes that classes are set on the <figure> element.
				conversionApi.writer.addClass( modelElement.getAttribute( 'customClass' ), viewFigure );
			} );
		}

		/**
         * Helper method that searches for a given view element in all children of the model element.
         *
         * @param {module:engine/view/item~Item} viewElement
         * @param {String} viewElementName
         * @param {module:engine/conversion/downcastdispatcher~DowncastConversionApi} conversionApi
         * @return {module:engine/view/item~Item}
         */
		function findViewChild( viewElement, viewElementName, conversionApi ) {
			const viewChildren = Array.from( conversionApi.writer.createRangeIn( viewElement ).getItems() );

			return viewChildren.find( item => item.is( viewElementName ) );
		}

		/**
         * Returns the custom attribute upcast converter.
         */
		function upcastAttribute( viewElementName, viewAttribute, modelAttribute ) {
			return dispatcher => dispatcher.on( `element:${ viewElementName }`, ( evt, data, conversionApi ) => {
				const viewItem = data.viewItem;
				const modelRange = data.modelRange;

				const modelElement = modelRange && modelRange.start.nodeAfter;

				if ( !modelElement ) {
					return;
				}

				conversionApi.writer.setAttribute( modelAttribute, viewItem.getAttribute( viewAttribute ), modelElement );
			} );
		}

		/**
         * Returns the custom attribute downcast converter.
         */
		function downcastAttribute( modelElementName, viewElementName, viewAttribute, modelAttribute ) {
			return dispatcher => dispatcher.on( `insert:${ modelElementName }`, ( evt, data, conversionApi ) => {
				const modelElement = data.item;

				const viewFigure = conversionApi.mapper.toViewElement( modelElement );
				const viewElement = findViewChild( viewFigure, viewElementName, conversionApi );

				if ( !viewElement ) {
					return;
				}

				conversionApi.writer.setAttribute( viewAttribute, modelElement.getAttribute( modelAttribute ), viewElement );
			} );
		}
	}
}
