#include <wmconsole/version3/wmcutil/yt/yt_utils.h>

#include "forecaster.h"

namespace NWebmaster {
namespace NForecaster {

const char *F_FORECASTER_BUDGET = "budget";
const char *F_FORECASTER_QUERIES = "queries";
const char *F_FORECASTER_REGIONS = "regions";

void TMapQueriesFromForecaster::Do(TReader *input, TWriter *output) {
    for (; input->IsValid(); input->Next()) {
        const NYT::TNode &row = input->GetRow();
        const NYT::TNode &results = row[F_FORECASTER_BUDGET]["results"][0];

        i64 regionId = 0;
        if (row[F_FORECASTER_REGIONS].IsInt64()) {
            regionId = row[F_FORECASTER_REGIONS].AsInt64();
        } else {
            regionId = row[F_FORECASTER_REGIONS][0].AsInt64();
        }

        output->AddRow(NYT::TNode()
            (F_FORECASTED_BID, results["bid"])
            (F_FORECASTED_BUDGET, results["budget"])
            (F_FORECASTED_CLICKS, results["clicks"])
            (F_FORECASTED_SHOWS, results["shows"])
            (F_REGION_ID, static_cast<ui64>(regionId))
            (F_QUERY, row[F_FORECASTER_QUERIES])
        );
    }
}

REGISTER_MAPPER(TMapQueriesFromForecaster)

void TReduceJoinForecatedData::Do(TReader *input, TWriter *output) {
    const int TABLENO_FORECASTER = 0;
    const int TABLENO_QUERIES = 1;

    bool forecasted = false;
    NYT::TNode forecasterRow;
    for (; input->IsValid(); input->Next()) {
        const int tableNo = input->GetTableIndex();
        NYT::TNode row = input->GetRow();

        if (tableNo == TABLENO_FORECASTER) {
            forecasterRow = row;
            forecasted = true;
        } else if (tableNo == TABLENO_QUERIES && forecasted) {
            if (forecasterRow[F_FORECASTED_BID].AsInt64() > 0 && forecasterRow[F_FORECASTED_CLICKS].AsDouble() > 0.0) {
                output->AddRow(row
                    (F_FORECASTED_BUDGET, forecasterRow[F_FORECASTED_BUDGET])
                    (F_FORECASTED_BID, forecasterRow[F_FORECASTED_BID])
                    (F_FORECASTED_CLICKS, forecasterRow[F_FORECASTED_CLICKS])
                    (F_FORECASTED_SHOWS, forecasterRow[F_FORECASTED_SHOWS])
                );
            }
        }
    }
}

REGISTER_REDUCER(TReduceJoinForecatedData)

void TReduceQueriesToForecaster::Do(TReader *input, TWriter *output) {
    const NYT::TNode &row = input->GetRow();
    output->AddRow(NYT::TNode()
        (F_FORECASTER_QUERIES, row[F_QUERY])
        (F_FORECASTER_REGIONS, row[F_REGION_ID])
    );
}

REGISTER_REDUCER(TReduceQueriesToForecaster)

TForecastTable::TForecastTable(NYT::IClientBasePtr client, const TString &toForecasterPrefix, const TString &fromForecasterPrefix, const TString &tablePrefix)
    : Client(client)
    , Timestamp(Now().Seconds())
    , ToForecasterPrefix(toForecasterPrefix)
    , FromForecasterPrefix(fromForecasterPrefix)
    , TablePrefix(tablePrefix)
    , ForecasterAttr(TString(ATTR_AWAITING_FORECAST) + "_" + TablePrefix)
{
    Cout << ToForecasterPrefix << Endl;
    Cout << ForecasterAttr << Endl;
    time_t awaitingTs = NYTUtils::GetAttrOrDefault<i64>(Client, ToForecasterPrefix, ForecasterAttr, 0);
    if (awaitingTs != 0) {
        Timestamp = awaitingTs;
        Awaiting = true;
    }
}

bool TForecastTable::IsForecastPending() const {
    return Awaiting;
}

TString TForecastTable::GetSourceTable() const {
    return NYTUtils::JoinPath(ToForecasterPrefix, TablePrefix + "_" + ToString(Timestamp));
}

TString TForecastTable::GetForecastTable() const {
    return NYTUtils::JoinPath(FromForecasterPrefix, TablePrefix + "_" + ToString(Timestamp));
}

bool TForecastTable::IsThereForecast() const {
    TDeque<NYTUtils::TTableInfo> tables;
    NYTUtils::GetTableList(Client, FromForecasterPrefix, tables);

    const TString output = GetForecastTable();

    for (const NYTUtils::TTableInfo &table : tables) {
        if (table.Name == output) {
            return true;
        }
    }

    return false;
}

TString TForecastTable::PostTableToForecast(NYT::ITransactionPtr tx) {
    Timestamp = Now().Seconds();
    NYTUtils::SetAttr(tx, ToForecasterPrefix, ForecasterAttr, Timestamp);
    Awaiting = true;
    return GetSourceTable();
}

void TForecastTable::Reset(NYT::ITransactionPtr tx) {
    NYTUtils::RemoveAttr(tx, ToForecasterPrefix, ForecasterAttr);
}

NYT::TTableSchema GetToForecasterSchema() {
    NYT::TTableSchema toForecasterSchema;
    toForecasterSchema.Strict(true);
    toForecasterSchema.AddColumn(NYT::TColumnSchema().Name(F_FORECASTER_QUERIES).Type(NYT::VT_STRING));
    toForecasterSchema.AddColumn(NYT::TColumnSchema().Name(F_FORECASTER_REGIONS).Type(NYT::VT_INT64));
    return toForecasterSchema;
}

} //namespace NForecaster
} //namespace NWebmaster
