# ModuleUnitOfWork.java

    /*
     * Copyright (c) 2009, Rickard Öberg. All Rights Reserved.
     * Copyright (c) 2013-2015, Niclas Hedhman. All Rights Reserved.
     * Copyright (c) 2013-2015, Paul Merlin. All Rights Reserved.
     *
     * 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 http://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.qi4j.runtime.structure;

    import java.util.Collections;
    import java.util.HashMap;
    import java.util.Iterator;
    import java.util.Map;
    import org.qi4j.api.association.AssociationDescriptor;
    import org.qi4j.api.association.AssociationStateHolder;
    import org.qi4j.api.common.QualifiedName;
    import org.qi4j.api.composite.Composite;
    import org.qi4j.api.entity.EntityBuilder;
    import org.qi4j.api.entity.EntityComposite;
    import org.qi4j.api.entity.EntityReference;
    import org.qi4j.api.entity.Identity;
    import org.qi4j.api.entity.IdentityGenerator;
    import org.qi4j.api.entity.LifecycleException;
    import org.qi4j.api.property.Property;
    import org.qi4j.api.property.PropertyDescriptor;
    import org.qi4j.api.property.StateHolder;
    import org.qi4j.api.query.Query;
    import org.qi4j.api.query.QueryBuilder;
    import org.qi4j.api.query.QueryExecutionException;
    import org.qi4j.api.query.grammar.OrderBy;
    import org.qi4j.api.service.NoSuchServiceException;
    import org.qi4j.api.unitofwork.ConcurrentEntityModificationException;
    import org.qi4j.api.unitofwork.EntityTypeNotFoundException;
    import org.qi4j.api.unitofwork.NoSuchEntityException;
    import org.qi4j.api.unitofwork.UnitOfWork;
    import org.qi4j.api.unitofwork.UnitOfWorkCallback;
    import org.qi4j.api.unitofwork.UnitOfWorkCompletionException;
    import org.qi4j.api.unitofwork.UnitOfWorkFactory;
    import org.qi4j.api.usecase.Usecase;
    import org.qi4j.api.util.NullArgumentException;
    import org.qi4j.api.value.ValueBuilder;
    import org.qi4j.api.value.ValueComposite;
    import org.qi4j.functional.Function;
    import org.qi4j.functional.Iterables;
    import org.qi4j.functional.Specification;
    import org.qi4j.runtime.association.AssociationInstance;
    import org.qi4j.runtime.association.ManyAssociationInstance;
    import org.qi4j.runtime.association.NamedAssociationInstance;
    import org.qi4j.runtime.composite.FunctionStateResolver;
    import org.qi4j.runtime.entity.EntityInstance;
    import org.qi4j.runtime.entity.EntityModel;
    import org.qi4j.runtime.property.PropertyModel;
    import org.qi4j.runtime.unitofwork.EntityBuilderInstance;
    import org.qi4j.runtime.unitofwork.UnitOfWorkInstance;
    import org.qi4j.runtime.value.ValueInstance;
    import org.qi4j.spi.entity.EntityState;
    import org.qi4j.spi.entity.EntityStatus;
    import org.qi4j.spi.entity.NamedAssociationState;
    import org.qi4j.spi.entitystore.EntityStore;
    import org.qi4j.spi.module.ModelModule;
    import org.qi4j.spi.query.EntityFinder;
    import org.qi4j.spi.query.EntityFinderException;
    import org.qi4j.spi.query.QueryBuilderSPI;
    import org.qi4j.spi.query.QuerySource;

    import static org.qi4j.api.entity.EntityReference.parseEntityReference;
    import static org.qi4j.functional.Iterables.first;
    import static org.qi4j.functional.Iterables.map;

    /**
     * JAVADOC
     */
    public class ModuleUnitOfWork
        implements UnitOfWork
    {
        private static final QualifiedName IDENTITY_STATE_NAME;

        static
        {
            try
            {
                IDENTITY_STATE_NAME = QualifiedName.fromAccessor( Identity.class.getMethod( "identity" ) );
            }
            catch( NoSuchMethodException e )
            {
                throw new InternalError( "Zest Core Runtime codebase is corrupted. Contact Zest team: ModuleUnitOfWork" );
            }
        }

        private final UnitOfWorkInstance uow;
        private final ModuleInstance module;

        ModuleUnitOfWork( ModuleInstance module, UnitOfWorkInstance uow )
        {
            this.module = module;
            this.uow = uow;
        }

        public ModuleInstance module()
        {
            return module;
        }

        public UnitOfWorkInstance instance()
        {
            return uow;
        }

        @Override
        public UnitOfWorkFactory unitOfWorkFactory()
        {
            return module;
        }

        @Override
        public long currentTime()
        {
            return uow.currentTime();
        }

        @Override
        public Usecase usecase()
        {
            return uow.usecase();
        }

        @Override
        public <T> T metaInfo( Class<T> infoType )
        {
            return uow.metaInfo().get( infoType );
        }

        @Override
        public void setMetaInfo( Object metaInfo )
        {
            uow.metaInfo().set( metaInfo );
        }

        @Override
        @SuppressWarnings( { "raw", "unchecked" } )
        public <T> Query<T> newQuery( QueryBuilder<T> queryBuilder )
        {
            QueryBuilderSPI queryBuilderSPI = (QueryBuilderSPI) queryBuilder;

            return queryBuilderSPI.newQuery( new UoWQuerySource( this ) );
        }

        @Override
        public <T> T newEntity( Class<T> type )
            throws EntityTypeNotFoundException, LifecycleException
        {
            return newEntity( type, null );
        }

        @Override
        public <T> T newEntity( Class<T> type, String identity )
            throws EntityTypeNotFoundException, LifecycleException
        {
            return newEntityBuilder( type, identity ).newInstance();
        }

        @Override
        public <T> EntityBuilder<T> newEntityBuilder( Class<T> type )
            throws EntityTypeNotFoundException
        {
            return newEntityBuilder( type, null );
        }

        @Override
        public <T> EntityBuilder<T> newEntityBuilder( Class<T> type, String identity )
            throws EntityTypeNotFoundException
        {
            ModelModule<EntityModel> model = module.typeLookup().lookupEntityModel( type );

            if( model == null )
            {
                throw new EntityTypeNotFoundException( type.getName(),
                                                       module.name(),
                                                       map( ModelModule.toStringFunction,
                                                            module.findVisibleEntityTypes()
                                                       ) );
            }

            EntityStore entityStore = model.module().entityStore();

            // Generate id if necessary
            if( identity == null )
            {
                IdentityGenerator idGen = model.module().identityGenerator();
                if( idGen == null )
                {
                    throw new NoSuchServiceException( IdentityGenerator.class.getName(), model.module().name() );
                }
                identity = idGen.generate( first( model.model().types() ) );
            }
            EntityBuilder<T> builder;

            builder = new EntityBuilderInstance<>( model,
                                                   this,
                                                   uow.getEntityStoreUnitOfWork( entityStore, module ),
                                                   identity );
            return builder;
        }

        @Override
        public <T> EntityBuilder<T> newEntityBuilderWithState(
            Class<T> type,
            Function<PropertyDescriptor, Object> propertyFunction,
            Function<AssociationDescriptor, EntityReference> associationFunction,
            Function<AssociationDescriptor, Iterable<EntityReference>> manyAssociationFunction,
            Function<AssociationDescriptor, Map<String, EntityReference>> namedAssociationFunction
        )
            throws EntityTypeNotFoundException
        {
            return newEntityBuilderWithState( type, null,
                                              propertyFunction,
                                              associationFunction,
                                              manyAssociationFunction,
                                              namedAssociationFunction );
        }

        @Override
        public <T> EntityBuilder<T> newEntityBuilderWithState(
            Class<T> type, String identity,
            Function<PropertyDescriptor, Object> propertyFunction,
            Function<AssociationDescriptor, EntityReference> associationFunction,
            Function<AssociationDescriptor, Iterable<EntityReference>> manyAssociationFunction,
            Function<AssociationDescriptor, Map<String, EntityReference>> namedAssociationFunction
        )
            throws EntityTypeNotFoundException
        {
            NullArgumentException.validateNotNull( "propertyFunction", propertyFunction );
            NullArgumentException.validateNotNull( "associationFunction", associationFunction );
            NullArgumentException.validateNotNull( "manyAssociationFunction", manyAssociationFunction );
            NullArgumentException.validateNotNull( "namedAssociationFunction", namedAssociationFunction );

            ModelModule<EntityModel> model = module.typeLookup().lookupEntityModel( type );

            if( model == null )
            {
                throw new EntityTypeNotFoundException( type.getName(),
                                                       module.name(),
                                                       map( ModelModule.toStringFunction,
                                                            module.findVisibleEntityTypes()
                                                       ) );
            }

            EntityStore entityStore = model.module().entityStore();

            FunctionStateResolver stateResolver = new FunctionStateResolver(
                propertyFunction, associationFunction, manyAssociationFunction, namedAssociationFunction
            );

            if( identity == null )
            {
                // Use identity from StateResolver if available
                PropertyModel identityModel = model.model().state().findPropertyModelByQualifiedName( IDENTITY_STATE_NAME );
                identity = (String) stateResolver.getPropertyState( identityModel );
                if( identity == null )
                {
                    // Generate identity
                    IdentityGenerator idGen = model.module().identityGenerator();
                    if( idGen == null )
                    {
                        throw new NoSuchServiceException( IdentityGenerator.class.getName(), model.module().name() );
                    }
                    identity = idGen.generate( first( model.model().types() ) );
                }
            }

            return new EntityBuilderInstance<>( model,
                                                this,
                                                uow.getEntityStoreUnitOfWork( entityStore, module ),
                                                identity,
                                                stateResolver );
        }

        @Override
        public <T> T get( Class<T> type, String identity )
            throws EntityTypeNotFoundException, NoSuchEntityException
        {
            Iterable<ModelModule<EntityModel>> models = module.typeLookup().lookupEntityModels( type );

            if( !models.iterator().hasNext() )
            {
                throw new EntityTypeNotFoundException( type.getName(),
                                                       module.name(),
                                                       map( ModelModule.toStringFunction,
                                                            module.findVisibleEntityTypes()
                                                       ) );
            }

            return uow.get( parseEntityReference( identity ), this, models, type );
        }

        @Override
        @SuppressWarnings( "unchecked" )
        public <T> T get( T entity )
            throws EntityTypeNotFoundException
        {
            EntityComposite entityComposite = (EntityComposite) entity;
            EntityInstance compositeInstance = EntityInstance.entityInstanceOf( entityComposite );
            ModelModule<EntityModel> model = new ModelModule<>( compositeInstance.module(), compositeInstance.entityModel() );
            Class<T> type = (Class<T>) first( compositeInstance.types() );
            return uow.get( compositeInstance.identity(), this, Collections.singletonList( model ), type );
        }

        @Override
        public void remove( Object entity )
            throws LifecycleException
        {
            uow.checkOpen();

            EntityComposite entityComposite = (EntityComposite) entity;

            EntityInstance compositeInstance = EntityInstance.entityInstanceOf( entityComposite );

            if( compositeInstance.status() == EntityStatus.NEW )
            {
                compositeInstance.remove( this );
                uow.remove( compositeInstance.identity() );
            }
            else if( compositeInstance.status() == EntityStatus.LOADED || compositeInstance.status() == EntityStatus.UPDATED )
            {
                compositeInstance.remove( this );
            }
            else
            {
                throw new NoSuchEntityException( compositeInstance.identity(), compositeInstance.types(), usecase() );
            }
        }

        @SuppressWarnings( "DuplicateThrows" )
        @Override
        public void complete()
            throws UnitOfWorkCompletionException, ConcurrentEntityModificationException
        {
            uow.complete();
        }

        @Override
        public void discard()
        {
            uow.discard();
        }

        @Override
        public void close()
        {
            discard();
        }

        @Override
        public boolean isOpen()
        {
            return uow.isOpen();
        }

        @Override
        public boolean isPaused()
        {
            return uow.isPaused();
        }

        @Override
        public void pause()
        {
            uow.pause();
        }

        @Override
        public void resume()
        {
            uow.resume();
        }

        @Override
        public void addUnitOfWorkCallback( UnitOfWorkCallback callback )
        {
            uow.addUnitOfWorkCallback( callback );
        }

        @Override
        public void removeUnitOfWorkCallback( UnitOfWorkCallback callback )
        {
            uow.removeUnitOfWorkCallback( callback );
        }

        @Override
        public boolean equals( Object o )
        {
            if( this == o )
            {
                return true;
            }
            if( o == null || getClass() != o.getClass() )
            {
                return false;
            }

            ModuleUnitOfWork that = (ModuleUnitOfWork) o;

            return uow.equals( that.uow );
        }

        @Override
        public int hashCode()
        {
            return uow.hashCode();
        }

        @Override
        public String toString()
        {
            return uow.toString();
        }

        public void addEntity( EntityInstance instance )
        {
            uow.addEntity( instance );
        }

        @Override
        public <T extends Identity> T toValue( Class<T> primaryType, T entityComposite )
        {
            Function<PropertyDescriptor, Object> propertyFunction = new ToValuePropertyMappingFunction( entityComposite );
            Function<AssociationDescriptor, EntityReference> assocationFunction = new ToValueAssociationMappingFunction<>( entityComposite );
            Function<AssociationDescriptor, Iterable<EntityReference>> manyAssocFunction = new ToValueManyAssociationMappingFunction<>( entityComposite );
            Function<AssociationDescriptor, Map<String, EntityReference>> namedAssocFunction = new ToValueNameAssociationMappingFunction<>( entityComposite );

            @SuppressWarnings( "unchecked" )
            ValueBuilder<T> builder = module().newValueBuilderWithState(
                primaryType, propertyFunction, assocationFunction, manyAssocFunction, namedAssocFunction );
            return builder.newInstance();
        }

        @Override
        public <T extends Identity> T toEntity( Class<T> primaryType, T valueComposite )
        {
            Function<PropertyDescriptor, Object> propertyFunction = new ToEntityPropertyMappingFunction<>( valueComposite );
            Function<AssociationDescriptor, EntityReference> assocationFunction = new ToEntityAssociationMappingFunction<>( valueComposite );
            Function<AssociationDescriptor, Iterable<EntityReference>> manyAssocFunction = new ToEntityManyAssociationMappingFunction<>( valueComposite );
            Function<AssociationDescriptor, Map<String, EntityReference>> namedAssocFunction = new ToEntityNameAssociationMappingFunction<>( valueComposite );

            String identity = valueComposite.identity().get();
            try
            {
                T entity = get( primaryType, identity );
                // If successful, then this entity is to by modified.
                EntityInstance instance = EntityInstance.entityInstanceOf( (EntityComposite) entity );
                EntityState state = instance.entityState();
                FunctionStateResolver stateResolver = new FunctionStateResolver( propertyFunction,
                                                                                 assocationFunction,
                                                                                 manyAssocFunction,
                                                                                 namedAssocFunction );
                EntityModel model = (EntityModel) EntityInstance.entityInstanceOf( (EntityComposite) entity ).descriptor();
                stateResolver.populateState( model, state );
                return entity;
            }
            catch( NoSuchEntityException e )
            {
                EntityBuilder<T> entityBuilder = newEntityBuilderWithState( primaryType,
                                                                            identity,
                                                                            propertyFunction,
                                                                            assocationFunction,
                                                                            manyAssocFunction,
                                                                            namedAssocFunction );
                return entityBuilder.newInstance();
            }
        }

        private static class UoWQuerySource implements QuerySource
        {
            private final ModuleUnitOfWork moduleUnitOfWork;

            private UoWQuerySource( ModuleUnitOfWork moduleUnitOfWork )
            {
                this.moduleUnitOfWork = moduleUnitOfWork;
            }

            @Override
            public <T> T find( Class<T> resultType,
                               Specification<Composite> whereClause,
                               Iterable<OrderBy> orderBySegments,
                               Integer firstResult,
                               Integer maxResults,
                               Map<String, Object> variables
            )
            {
                final EntityFinder entityFinder = moduleUnitOfWork.module().findService( EntityFinder.class ).get();

                try
                {
                    final EntityReference foundEntity = entityFinder.findEntity( resultType, whereClause, variables == null ? Collections
                        .<String, Object>emptyMap() : variables );
                    if( foundEntity != null )
                    {
                        try
                        {
                            return moduleUnitOfWork.get( resultType, foundEntity.identity() );
                        }
                        catch( NoSuchEntityException e )
                        {
                            return null; // Index is out of sync - entity has been removed
                        }
                    }
                    // No entity was found
                    return null;
                }
                catch( EntityFinderException e )
                {
                    throw new QueryExecutionException( "Finder caused exception", e );
                }
            }

            @Override
            public <T> long count( Class<T> resultType,
                                   Specification<Composite> whereClause,
                                   Iterable<OrderBy> orderBySegments,
                                   Integer firstResult,
                                   Integer maxResults,
                                   Map<String, Object> variables
            )
            {
                final EntityFinder entityFinder = moduleUnitOfWork.module().findService( EntityFinder.class ).get();

                try
                {
                    return entityFinder.countEntities( resultType, whereClause, variables == null ? Collections.<String, Object>emptyMap() : variables );
                }
                catch( EntityFinderException e )
                {
                    e.printStackTrace();
                    return 0;
                }
            }

            @Override
            public <T> Iterator<T> iterator( final Class<T> resultType,
                                             Specification<Composite> whereClause,
                                             Iterable<OrderBy> orderBySegments,
                                             Integer firstResult,
                                             Integer maxResults,
                                             Map<String, Object> variables
            )
            {
                final EntityFinder entityFinder = moduleUnitOfWork.module().findService( EntityFinder.class ).get();

                try
                {
                    final Iterator<EntityReference> foundEntities = entityFinder.findEntities( resultType,
                                                                                               whereClause,
                                                                                               Iterables.toArray( OrderBy.class, orderBySegments ),
                                                                                               firstResult,
                                                                                               maxResults,
                                                                                               variables == null ? Collections
                                                                                                   .<String, Object>emptyMap() : variables )
                        .iterator();

                    return new Iterator<T>()
                    {
                        @Override
                        public boolean hasNext()
                        {
                            return foundEntities.hasNext();
                        }

                        @Override
                        public T next()
                        {
                            final EntityReference foundEntity = foundEntities.next();
                            try
                            {
                                return moduleUnitOfWork.get( resultType, foundEntity.identity() );
                            }
                            catch( NoSuchEntityException e )
                            {
                                // Index is out of sync - entity has been removed
                                return null;
                            }
                        }

                        @Override
                        public void remove()
                        {
                            throw new UnsupportedOperationException();
                        }
                    };
                }
                catch( EntityFinderException e )
                {
                    throw new QueryExecutionException( "Query '" + toString() + "' could not be executed", e );
                }
            }

            @Override
            public String toString()
            {
                return "UnitOfWork( " + moduleUnitOfWork.usecase().name() + " )";
            }
        }

        private class ToValuePropertyMappingFunction
            implements Function<PropertyDescriptor, Object>
        {
            private Object entity;

            public ToValuePropertyMappingFunction( Object entity )
            {
                this.entity = entity;
            }

            @Override
            public Object map( PropertyDescriptor propertyDescriptor )
            {
                EntityState entityState = EntityInstance.entityInstanceOf( (EntityComposite) entity ).entityState();
                return entityState.propertyValueOf( propertyDescriptor.qualifiedName() );
            }
        }

        private class ToValueAssociationMappingFunction<T>
            implements Function<AssociationDescriptor, EntityReference>
        {
            private final T entity;

            public ToValueAssociationMappingFunction( T entity )
            {
                this.entity = entity;
            }

            @Override
            public EntityReference map( AssociationDescriptor associationDescriptor )
            {
                EntityState entityState = EntityInstance.entityInstanceOf( (EntityComposite) entity ).entityState();
                return entityState.associationValueOf( associationDescriptor.qualifiedName() );
            }
        }

        private class ToValueManyAssociationMappingFunction<T>
            implements Function<AssociationDescriptor, Iterable<EntityReference>>
        {
            private final T entity;

            public ToValueManyAssociationMappingFunction( T entity )
            {
                this.entity = entity;
            }

            @Override
            public Iterable<EntityReference> map( AssociationDescriptor associationDescriptor )
            {
                EntityState entityState = EntityInstance.entityInstanceOf( (EntityComposite) entity ).entityState();
                return entityState.manyAssociationValueOf( associationDescriptor.qualifiedName() );
            }
        }

        private class ToValueNameAssociationMappingFunction<T>
            implements Function<AssociationDescriptor, Map<String, EntityReference>>
        {
            private final T entity;

            public ToValueNameAssociationMappingFunction( T entity )
            {
                this.entity = entity;
            }

            @Override
            public Map<String, EntityReference> map( AssociationDescriptor associationDescriptor )
            {
                Map<String, EntityReference> result = new HashMap<>();
                EntityState entityState = EntityInstance.entityInstanceOf( (EntityComposite) entity ).entityState();
                final NamedAssociationState state = entityState.namedAssociationValueOf( associationDescriptor.qualifiedName() );
                for( String name : state )
                {
                    result.put( name, state.get( name ) );
                }
                return result;
            }
        }

        private class ToEntityPropertyMappingFunction<T>
            implements Function<PropertyDescriptor, Object>
        {
            private final T value;

            public ToEntityPropertyMappingFunction( T value )
            {
                this.value = value;
            }

            @Override
            public Object map( PropertyDescriptor propertyDescriptor )
            {
                StateHolder state = ValueInstance.valueInstanceOf( (ValueComposite) value ).state();
                Property<Object> property = state.propertyFor( propertyDescriptor.accessor() );
                return property.get();
            }
        }

        private class ToEntityAssociationMappingFunction<T>
            implements Function<AssociationDescriptor, EntityReference>
        {

            private final T value;

            public ToEntityAssociationMappingFunction( T value )
            {
                this.value = value;
            }

            @Override
            public EntityReference map( AssociationDescriptor associationDescriptor )
            {
                AssociationStateHolder state = ValueInstance.valueInstanceOf( (ValueComposite) value ).state();
                AssociationInstance<T> association = (AssociationInstance<T>) state.associationFor( associationDescriptor.accessor() );
                return association.getAssociationState().get();
            }
        }

        private class ToEntityManyAssociationMappingFunction<T>
            implements Function<AssociationDescriptor, Iterable<EntityReference>>
        {

            private final T value;

            public ToEntityManyAssociationMappingFunction( T valueComposite )
            {
                this.value = valueComposite;
            }

            @Override
            public Iterable<EntityReference> map( AssociationDescriptor associationDescriptor )
            {
                AssociationStateHolder state = ValueInstance.valueInstanceOf( (ValueComposite) value ).state();
                ManyAssociationInstance<T> association =
                    (ManyAssociationInstance<T>) state.manyAssociationFor( associationDescriptor.accessor() );
                return association.getManyAssociationState();
            }
        }

        private class ToEntityNameAssociationMappingFunction<T>
            implements Function<AssociationDescriptor, Map<String, EntityReference>>
        {
            private final T value;

            public ToEntityNameAssociationMappingFunction( T valueComposite )
            {
                this.value = valueComposite;
            }

            @Override
            public Map<String, EntityReference> map( AssociationDescriptor associationDescriptor )
            {
                AssociationStateHolder state = ValueInstance.valueInstanceOf( (ValueComposite) value ).state();
                NamedAssociationInstance<T> association =
                    (NamedAssociationInstance<T>) state.namedAssociationFor( associationDescriptor.accessor() );
                HashMap<String, EntityReference> result = new HashMap<>();
                for( Map.Entry<String, EntityReference> entry : association.getEntityReferences() )
                {
                    result.put( entry.getKey(), entry.getValue() );
                }
                return result;
            }
        }
    }