package org.atlasapi.output.annotation;

import com.google.common.base.Function;
import com.google.common.base.Splitter;
import com.google.common.base.Strings;
import com.google.common.collect.ImmutableMultimap;
import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.Futures;
import org.atlasapi.channel.Channel;
import org.atlasapi.channel.ChannelGroup;
import org.atlasapi.channel.ChannelGroupMembership;
import org.atlasapi.channel.ChannelResolver;
import org.atlasapi.criteria.attribute.Attributes;
import org.atlasapi.entity.Id;
import org.atlasapi.entity.util.Resolved;
import org.atlasapi.media.channel.ChannelQuery;
import org.atlasapi.output.ChannelWithChannelGroupMembership;
import org.atlasapi.output.FieldWriter;
import org.atlasapi.output.OutputContext;
import org.atlasapi.query.v4.channelgroup.ChannelGroupChannelWriter;

import java.io.IOException;
import java.util.Set;
import java.util.concurrent.TimeUnit;

import static com.google.common.base.Preconditions.checkNotNull;

public class ChannelGroupChannelsAnnotation extends OutputAnnotation<org.atlasapi.channel.ChannelGroup> {

    private final ChannelGroupChannelWriter channelWriter;
    private final ChannelResolver channelResolver;

    public ChannelGroupChannelsAnnotation(ChannelGroupChannelWriter channelWriter, ChannelResolver channelResolver) {
        this.channelWriter = checkNotNull(channelWriter);
        this.channelResolver = checkNotNull(channelResolver);
    }

    @Override
    public void write(ChannelGroup entity, FieldWriter writer, OutputContext ctxt) throws IOException {
        final ImmutableMultimap.Builder<Id, ChannelGroupMembership> builder = ImmutableMultimap.builder();
        for (ChannelGroupMembership channelGroupMembership : (Set<ChannelGroupMembership>)entity.getChannels()) {
            builder.put(channelGroupMembership.getChannel().getId(), channelGroupMembership);
        }

        ImmutableMultimap<Id, ChannelGroupMembership> channelGroupMemberships = builder.build();
        String genre = ctxt.getRequest().getParameter(Attributes.CHANNEL_GROUP_CHANNEL_GENRES.externalName());

        ChannelQuery.Builder channelQueryBuilder = ChannelQuery.builder()
                .withChannelGroups(ImmutableSet.of(entity.getId().longValue()));
        if (!Strings.isNullOrEmpty(genre)) {
            channelQueryBuilder.withGenres(ImmutableSet.copyOf(Splitter.on(',').split(genre)));
        }
        ChannelQuery channelQuery = channelQueryBuilder.build();
        Iterable<Channel> channels = Futures.get(
                Futures.transform(
                        this.channelResolver.resolveChannels(channelQuery),
                        new Function<Resolved<Channel>, Iterable<Channel>>() {
                            @Override
                            public Iterable<Channel> apply(Resolved<Channel> channelResolved) {
                                return channelResolved.getResources();
                            }
                        }

                ), 1, TimeUnit.MINUTES, IOException.class
        );
        ImmutableSet.Builder<ChannelWithChannelGroupMembership> resultBuilder = ImmutableSet.builder();

        for (Channel channel : channels) {
            for (ChannelGroupMembership channelGroupMembership : channelGroupMemberships.get(channel.getId())) {
                resultBuilder.add(
                        new ChannelWithChannelGroupMembership(
                                channel,
                                channelGroupMembership
                        )
                );
            }
        }
        ImmutableSet<ChannelWithChannelGroupMembership> result = resultBuilder.build();
        writer.writeList(channelWriter, result, ctxt);
    }
}

@@ -70,24 +70,65 @@
 e.Function;%0A
+import com.google.common.base.Predicate;%0A
 import com.g
@@ -289,24 +289,107 @@
 mutableSet;%0A
+import com.google.common.collect.Iterables;%0Aimport com.google.common.collect.Sets;%0A
 import com.g
@@ -2263,379 +2263,8 @@
 );%0A%0A
-        ChannelQuery.Builder channelQueryBuilder = ChannelQuery.builder()%0A                .withChannelGroups(ImmutableSet.of(entity.getId().longValue()));%0A        if (!Strings.isNullOrEmpty(genre)) %7B%0A            channelQueryBuilder.withGenres(ImmutableSet.copyOf(Splitter.on(',').split(genre)));%0A        %7D%0A        ChannelQuery channelQuery = channelQueryBuilder.build();%0A
     
@@ -2396,23 +2396,18 @@
 .resolve
-Channel
+Id
 s(channe
@@ -2407,21 +2407,41 @@
 (channel
-Query
+GroupMemberships.keySet()
 ),%0A     
@@ -2837,32 +2837,462 @@
 lass%0A        );%0A
+        if (!Strings.isNullOrEmpty(genre)) %7B%0A            final ImmutableSet%3CString%3E genres = ImmutableSet.copyOf(Splitter.on(',').split(genre));%0A            channels = Iterables.filter(channels, new Predicate%3CChannel%3E() %7B%0A                @Override%0A                public boolean apply(Channel input) %7B%0A                    return !Sets.intersection(input.getGenres(), genres).isEmpty();%0A                %7D%0A            %7D);%0A        %7D%0A
         Immutabl

Fix channels not being returned for some channel groups.