#pragma once

#include <library/cpp/deprecated/atomic/atomic.h>

#include <util/thread/pool.h>
#include <util/thread/lfqueue.h>

#include <thread>
#include <chrono>
#include <optional>

namespace maps {
namespace wiki {
namespace autocart {

template <typename T>
struct QueueCounter {
    void IncCount(const T&) {
        AtomicIncrement(Counter);
    }

    void DecCount(const T&) {
        AtomicDecrement(Counter);
    }

    TAtomic Counter = 0;
};

class ObjectInQueueWithData : public IObjectInQueue {
public:
    ObjectInQueueWithData() {
        AtomicSet(waitData_, 1);
        AtomicSet(running_, 1);
    }

    void DataEnded() {
        AtomicSet(waitData_, 0);
    }

    bool isRunning() const {
        return (0 != AtomicGet(running_));
    }

protected:
    bool isWaitData() const {
        return (0 != AtomicGet(waitData_));
    }

    void Stop() {
        AtomicSet(running_, 0);
    }

private:
    TAtomic waitData_;
    TAtomic running_;
};

template <typename Output>
class DataReader : public IObjectInQueue {
public:
    DataReader() {
        AtomicSet(running_, 1);
    }

    void Process(void* /*threadSpecificResource*/) override {
        Read();
        Finish();
    }

    void SetOutputQueue(TLockFreeQueue<Output, QueueCounter<Output>>* outQueue) {
        outQueue_ = outQueue;
    }

    bool isRunning() const {
        return (0 != AtomicGet(running_));
    }

protected:
    virtual void Read() = 0;

    void Finish() {
        AtomicSet(running_, 0);
    }

    TLockFreeQueue<Output, QueueCounter<Output>>* outQueue_;

private:
    TAtomic running_;
};

template <typename Input, typename Output>
class DataProcessor : public ObjectInQueueWithData {
public:
    void Process(void* /*threadSpecificResource*/) override {
        for (;;) {
            Input inpData;
            if (inpQueue_->Dequeue(&inpData)) {
                std::optional<Output> outData = Do(inpData);
                if (outData) {
                    outQueue_->Enqueue(*outData);
                }
            } else if (!isWaitData()) {
                break;
            } else {
                std::this_thread::sleep_for(std::chrono::milliseconds(20));
            }
        }
        Finish();
    }

    void SetInputQueue(TLockFreeQueue<Input, QueueCounter<Input>>* inpQueue) {
        inpQueue_ = inpQueue;
    }

    void SetOutputQueue(TLockFreeQueue<Output, QueueCounter<Output>>* outQueue) {
        outQueue_ = outQueue;
    }

    virtual DataProcessor* Clone() const = 0;

protected:
    virtual std::optional<Output> Do(const Input& data) = 0;

    void Finish() {
       Stop();
    }

    TLockFreeQueue<Input, QueueCounter<Input>>* inpQueue_;
    TLockFreeQueue<Output, QueueCounter<Output>>* outQueue_;
};

template <typename Input>
class DataWriter : public ObjectInQueueWithData {
public:
    void Process(void* /*threadSpecificResource*/) override {
        for (;;) {
            Input data;
            if (inpQueue_->Dequeue(&data)) {
                Write(data);
            } else if (!isWaitData()) {
                break;
            } else {
                std::this_thread::sleep_for(std::chrono::milliseconds(20));
            }
        }
        Finish();
    }

    void SetInputQueue(TLockFreeQueue<Input, QueueCounter<Input>>* inpQueue) {
        inpQueue_ = inpQueue;
    }

protected:
    virtual void Write(const Input& data) = 0;

    void Finish() {
        Stop();
    }

    TLockFreeQueue<Input, QueueCounter<Input>>* inpQueue_;
};

template <typename Input, typename Output>
class MultithreadDataProcessor {
public:
    void SetReader(TAutoPtr<DataReader<Input>> reader) {
        reader_ = reader;
        reader_->SetOutputQueue(&inpQueue_);
    }

    void SetWriter(TAutoPtr<DataWriter<Output>> writer) {
        writer_ = writer;
        writer_->SetInputQueue(&outQueue_);
    }

    void SetProcessors(TAutoPtr<DataProcessor<Input, Output>> processor,
                       size_t count = 1u) {
        processors_.clear();
        for (size_t i = 0; i < count; i++) {
            TAutoPtr<DataProcessor<Input, Output>> instance(processor->Clone());
            instance->SetInputQueue(&inpQueue_);
            instance->SetOutputQueue(&outQueue_);
            processors_.push_back(instance);
        }
    }

    void Run() {
        // threads for processors, reader and writer
        TAutoPtr<IThreadPool> mtpQueue = CreateThreadPool(processors_.size() + 2);
        mtpQueue->SafeAdd(writer_.Get());
        for (const auto& processor : processors_) {
            mtpQueue->SafeAdd(processor.Get());
        }
        mtpQueue->SafeAdd(reader_.Get());

        while (reader_->isRunning()) {
            Wait();
        }

        for (const auto& processor : processors_) {
            processor->DataEnded();
        }

        while (!inpQueue_.IsEmpty()) {
            Wait();
        }

        bool waitProcessor = true;
        while (waitProcessor) {
            Wait();
            waitProcessor = false;
            for (const auto& processor : processors_) {
                waitProcessor |= processor->isRunning();
            }
        }
        writer_->DataEnded();
        mtpQueue->Stop();
    }

private:
    void Wait() const {
        std::this_thread::sleep_for(std::chrono::milliseconds(100));
    }

    TAutoPtr<DataReader<Input>> reader_;
    std::vector<TAutoPtr<DataProcessor<Input, Output>>> processors_;
    TAutoPtr<DataWriter<Output>> writer_;

    TLockFreeQueue<Input, QueueCounter<Input>> inpQueue_;
    TLockFreeQueue<Output, QueueCounter<Output>> outQueue_;
};

} // autocart
} // wiki
} // maps
