package ru.yandex.webmaster3.worker.download.links;

import java.io.IOException;
import java.util.Base64;

import lombok.AllArgsConstructor;
import org.springframework.beans.factory.annotation.Autowired;

import ru.yandex.webmaster3.core.util.functional.ThrowingConsumer;
import ru.yandex.webmaster3.storage.download.common.LinkMdsExportDescriptor;
import ru.yandex.webmaster3.storage.download.common.MdsExportTaskData;
import ru.yandex.webmaster3.storage.util.yt.AsyncTableReader;
import ru.yandex.webmaster3.storage.util.yt.YtPath;
import ru.yandex.webmaster3.storage.util.yt.YtService;
import ru.yandex.webmaster3.storage.util.yt.YtTableRange;
import ru.yandex.webmaster3.storage.util.yt.YtTableReadDriver;
import ru.yandex.webmaster3.worker.download.AbstractMdsDataProvider;

import static ru.yandex.webmaster3.core.util.functional.ThrowingConsumer.rethrowingUnchecked;

@AllArgsConstructor(onConstructor_ = {@Autowired})
public class LinkDownloadMdsDataProvider extends AbstractMdsDataProvider<LinkData> {
    private final static byte[] EMPTY_ARCHIVE_GZIP = Base64.getDecoder().decode("H4sICAF9g14AA2VtcHR5LnRzdgADAAAAAAAAAAAA");
    private final YtService ytService;
    private final YtPath externalLinksYtPath;
    private final YtPath internalLinksYtPath;

    @Override
    public void provide(MdsExportTaskData exportTaskData, ThrowingConsumer<LinkData, Exception> consumer) throws Exception {
        LinkMdsExportDescriptor descriptor = (LinkMdsExportDescriptor) exportTaskData.getDescriptor();

        ytService.withoutTransaction((cypressService) -> {
            AsyncTableReader<LinkData> reader = new AsyncTableReader<>(
                    cypressService,
                    descriptor.isExternal() ? externalLinksYtPath : internalLinksYtPath,
                    YtTableRange.singleKey(descriptor.getHost()),
                    YtTableReadDriver.createYSONDriver(LinkData.class)
            ).withRetry(3).splitInParts(10000L);

            try (AsyncTableReader.TableIterator<LinkData> read = reader.read()) {
                if (!read.hasNext()) {
                    rethrowingUnchecked(consumer).accept(new LinkData("", "", EMPTY_ARCHIVE_GZIP));
                } else {
                    do {
                        rethrowingUnchecked(consumer).accept(read.next());
                    } while (read.hasNext());
                }
            } catch (IOException e) {
                throw new RuntimeException(e);
            }

            return true;
        });
    }

    @Override
    public Class<LinkData> getRowClass() {
        return LinkData.class;
    }
}
