<reponame>adventum/airbyte
#
# Copyright (c) 2021 Airbyte, Inc., all rights reserved.
#


import json
import random
import string
from abc import ABC
from time import sleep
from typing import Any, Iterable, List, Mapping, Optional, Tuple, Union

import requests
from airbyte_cdk.sources import AbstractSource
from airbyte_cdk.sources.streams import Stream
from airbyte_cdk.sources.streams.http import HttpStream
from airbyte_cdk.sources.streams.http.exceptions import (
    DefaultBackoffException,
    UserDefinedBackoffException,
)

from .schema_fields import CUSTOM_SCHEMA_FIELDS
from .utils import random_name

from airbyte_cdk import logger as airbyte_logger

logger = airbyte_logger.AirbyteLogger()


# Basic full refresh stream
class YandexDirectStream(HttpStream, ABC):

    url_base = "https://api.direct.yandex.com/json/v5/reports"
    http_method = "POST"

    def __init__(self, config: Mapping[str, Any]):
        super().__init__(authenticator=None)
        self.config = config
        self.report_name = random_name(10)

    def _send(
        self, request: requests.PreparedRequest, request_kwargs: Mapping[str, Any]
    ) -> requests.Response:
        while True:
            response = self._session.send(request, **request_kwargs)
            response.encoding = "utf-8"
            sleep_time = int(response.headers.get("retryIn", 5))
            if response.status_code == 200:
                return response
            elif response.status_code == 201:
                sleep(sleep_time)
                logger.info(
                    "Report is creating in offline mode. Re-check after "
                    + str(sleep_time)
                    + " seconds"
                )
            elif response.status_code == 202:
                sleep(sleep_time)
                logger.info(
                    "Report is creating in offline mode. Re-check after "
                    + str(sleep_time)
                    + " seconds"
                )
            else:
                if self.should_retry(response):
                    custom_backoff_time = self.backoff_time(response)
                    if custom_backoff_time:
                        raise UserDefinedBackoffException(
                            backoff=custom_backoff_time,
                            request=request,
                            response=response,
                        )
                    else:
                        raise DefaultBackoffException(
                            request=request, response=response
                        )
                elif self.raise_on_http_errors:
                    # Raise any HTTP exceptions that happened in case there were unexpected ones
                    response.raise_for_status()

    def add_constants_to_record(self, record):
        constants = {
            "__productName": self.config["product_name"],
            "__clientName": self.config["client_name"],
        }
        constants.update(json.loads(self.config.get("custom_json", "{}")))
        record.update(constants)
        return record

    def next_page_token(
        self, response: requests.Response
    ) -> Optional[Mapping[str, Any]]:
        return None

    def parse_response(
        self, response: requests.Response, **kwargs
    ) -> Iterable[Mapping]:
        # parse raw TSV data to list of named dicts
        raw_data_lines = response.text.split("\n")
        header = []
        data = []
        for line_n, line in enumerate(raw_data_lines):
            line_values = line.split("\t")

            # skip empty rows
            if not line.strip():
                continue

            # grab header and skip if it is first line
            if line_n == 0:
                header = line_values
                continue

            # zip values list to named dict
            data_item = {}
            for value_n, value in enumerate(line_values):
                data_item[header[value_n]] = value
            data.append(data_item)

        return map(self.add_constants_to_record, data)


class CustomReport(YandexDirectStream):
    def path(self, **kwargs) -> str:
        return ""

    @property
    def primary_key(self) -> Optional[Union[str, List[str], List[List[str]]]]:
        return self.config["fields"][0]

    def request_body_json(self, **kwargs) -> Optional[Mapping]:
        date_range = self.config["transformed_data_range"]
        params = {
            "params": {
                "SelectionCriteria": {},
                "FieldNames": self.config["fields"],
                "ReportName": self.report_name,
                "ReportType": "CUSTOM_REPORT",
                "DateRangeType": date_range["data_range_title"],
                "Format": "TSV",
                "IncludeVAT": "NO",
                "IncludeDiscount": "NO",
            }
        }
        if date_range["data_range_title"] == "CUSTOM_DATE":
            params["params"]["SelectionCriteria"].update(
                {
                    "DateFrom": date_range["date_from"],
                    "DateTo": date_range["date_to"],
                }
            )

        if self.config.get("parsed_filters"):
            params["SelectionCriteria"]["Filter"] = self.config.get("parsed_filters")
        return params

    def request_headers(self, *args, **kwargs) -> Mapping[str, Any]:
        headers = {
            "Authorization": "Bearer " + self.config["auth_token"],
            "Accept-Language": "ru",
            "processingMode": "auto",
            "skipReportHeader": "true",
            "skipReportSummary": "true",
        }
        if self.config.get("client_login"):
            headers["Client-Login"] = self.config["client_login"]
        return headers

    def get_json_schema(self) -> Mapping[str, Any]:
        schema = super().get_json_schema()
        for field_name in self.config["fields"]:
            field_schema_type = CUSTOM_SCHEMA_FIELDS.get(field_name, None)
            if field_schema_type:
                schema["properties"][field_name] = {"type": ["null", field_schema_type]}

        extra_properties = ["__productName", "__clientName"]
        custom_keys = (
            json.loads(self.config["custom_json"]).keys()
            if self.config.get("custom_json")
            else []
        )
        extra_properties.extend(custom_keys)
        for key in extra_properties:
            schema["properties"][key] = {"type": ["null", "string"]}

        return schema


# Source
class SourceYandexDirect(AbstractSource):
    def check_connection(self, logger, config) -> Tuple[bool, Any]:
        date_exclusive_fields = []
        for field_name in config["fields"]:
            if field_name not in CUSTOM_SCHEMA_FIELDS.keys():
                return (
                    False,
                    f"Unknown field name {field_name}. Check https://yandex.ru/dev/direct/doc/reports/fields-list.html for available fields (CUSTOM_REPORT column).",
                )
            if not CUSTOM_SCHEMA_FIELDS[field_name]:
                return (
                    False,
                    f"Field name {field_name} can't be used for CUSTOM_REPORT. Check https://yandex.ru/dev/direct/doc/reports/fields-list.html for available fields (CUSTOM_REPORT column).",
                )
            if CUSTOM_SCHEMA_FIELDS[field_name] == "filter":
                return (
                    False,
                    f"Field name {field_name} is only used for filter data and can't be shown as report field. Check https://yandex.ru/dev/direct/doc/reports/fields-list.html for available fields (CUSTOM_REPORT column).",
                )
            if field_name in ["Date", "Week", "Month", "Quarter", "Year"]:
                date_exclusive_fields.append(field_name)
            if (
                field_name
                in [
                    "Impressions",
                    "Ctr",
                    "AvgImpressionPosition",
                    "WeightedImpressions",
                    "WeightedCtr",
                    "AvgTrafficVolume",
                ]
                and "ClickType" in config["fields"]
            ):
                return (
                    False,
                    f"ClickType field is incompatible with {field_name} field. Check https://yandex.ru/dev/direct/doc/reports/compatibility.html",
                )
            if (
                field_name
                in [
                    "AdFormat",
                    "AdId",
                    "Age",
                    "CarrierType",
                    "Gender",
                    "MobilePlatform",
                    "RlAdjustmentId",
                ]
                and "ImpressionShare" in config["fields"]
            ):
                return (
                    False,
                    f"ImpressionShare field is incompatible with {field_name} field. Check https://yandex.ru/dev/direct/doc/reports/compatibility.html",
                )
        if len(date_exclusive_fields) > 1:
            date_fields = ", ".join(date_exclusive_fields)
            return (
                False,
                f"{date_fields} fields are mutually exclusive: only one of them can be present in the report. Check https://yandex.ru/dev/direct/doc/reports/compatibility.html",
            )

        if set(config["fields"]).intersection(
            ["Criterion", "CriterionId", "CriterionType"]
        ) and set(config["fields"]).intersection(
            ["Criteria", "CriteriaId", "CriteriaType"]
        ):
            return (
                False,
                "Criterion, CriterionId, CriterionType fields are incompatible with Criteria, CriteriaId, CriteriaType fields.",
            )

        if config["date_range"].get("last_days_count"):
            if not config["date_range"].get("last_days_count") in [
                3,
                5,
                7,
                14,
                30,
                90,
                365,
            ]:
                return (
                    False,
                    "Last N days field only can be 3, 5, 7, 14, 30, 90 or 365.",
                )

        try:  #  YYYY-MM-DD
            AUTH_TOKEN = config["auth_token"]
            CLIENT_LOGIN = config.get("client_login", "")
            FIELD_NAMES = [
                "Date",
                "CampaignName",
                "LocationOfPresenceName",
                "Impressions",
                "Clicks",
                "Cost",
            ]
            REPORT_NAME = random_name(9)

            body = {
                "params": {
                    "SelectionCriteria": {},
                    "FieldNames": FIELD_NAMES,
                    "ReportName": REPORT_NAME,
                    "ReportType": "CUSTOM_REPORT",
                    "DateRangeType": "TODAY",
                    "Format": "TSV",
                    "IncludeVAT": "NO",
                    "IncludeDiscount": "NO",
                }
            }

            headers = {
                "Authorization": "Bearer " + AUTH_TOKEN,
                "Accept-Language": "ru",
                "processingMode": "auto",
                "skipReportHeader": "true",
                "skipReportSummary": "true",
            }
            if CLIENT_LOGIN:
                headers["Client-Login"] = CLIENT_LOGIN

            req = requests.post(
                url="https://api.direct.yandex.com/json/v5/reports",
                json=body,
                headers=headers,
            )
            if req.status_code not in [200, 201, 202]:
                return (
                    False,
                    f"Test request status code - {req.status_code}. Body: {req.text}",
                )
        except Exception as e:
            return False, e

        try:
            filter_objects = json.loads(config["filters_json"])
            if not isinstance(filter_objects, list):
                return (
                    False,
                    'Filter JSON must be array of filter objects like this: [{"Fiield": "Year", "Operator": "EQUALS", "Values": ["2021"]}, {...}, {...} ...]. Check https://yandex.ru/dev/direct/doc/reports/filters.html',
                )
            for filter_object in filter_objects:
                if not isinstance(filter_object, dict):
                    return (
                        False,
                        f'Filter JSON must be array of filter objects like this: [{{"Fiield": "Year", "Operator": "EQUALS", "Values": ["2021"]}}, {{...}}, {{...}} ...]. "{filter_object}" filter has validation errors. Check https://yandex.ru/dev/direct/doc/reports/filters.html',
                    )
        except:
            pass

        return True, None

    def transform_config(self, config: Mapping[str, Any]) -> Mapping[str, Any]:
        available_data_range = {
            "today": {"data_range_title": "TODAY"},
            "yesterday": {"data_range_title": "YESTERDAY"},
            "custom_date": {
                "data_range_title": "CUSTOM_DATE",
                "date_from": config["date_range"].get("date_from"),
                "date_to": config["date_range"].get("date_to"),
            },
            "last_n_days": {
                "data_range_title": f"LAST_"
                + str(config["date_range"].get("last_days_count"))
                + "_DAYS"
            },
            "all_time": {"data_range_title": "ALL_TIME"},
        }
        config["transformed_data_range"] = available_data_range[
            config["date_range"]["data_range_type"]
        ]

        config["parsed_filters"] = json.loads(config["filters_json"])
        return config

    def streams(self, config: Mapping[str, Any]) -> List[Any]:
        return [CustomReport(config=self.transform_config(config))]

#
# Copyright (c) 2021 Airbyte, Inc., all rights reserved.
#


import json
import random
import string
from abc import ABC
from time import sleep
from typing import Any, Iterable, List, Mapping, Optional, Tuple, Union

import requests
from airbyte_cdk.sources import AbstractSource
from airbyte_cdk.sources.streams import Stream
from airbyte_cdk.sources.streams.http import HttpStream
from airbyte_cdk.sources.streams.http.exceptions import (
    DefaultBackoffException,
    UserDefinedBackoffException,
)

from .schema_fields import CUSTOM_SCHEMA_FIELDS
from .utils import random_name

from airbyte_cdk import logger as airbyte_logger

logger = airbyte_logger.AirbyteLogger()


# Basic full refresh stream
class YandexDirectStream(HttpStream, ABC):

    url_base = "https://api.direct.yandex.com/json/v5/reports"
    http_method = "POST"

    def __init__(self, config: Mapping[str, Any]):
        super().__init__(authenticator=None)
        self.config = config
        self.report_name = random_name(10)

    def _send(
        self, request: requests.PreparedRequest, request_kwargs: Mapping[str, Any]
    ) -> requests.Response:
        while True:
            response = self._session.send(request, **request_kwargs)
            response.encoding = "utf-8"
            sleep_time = int(response.headers.get("retryIn", 5))
            if response.status_code == 200:
                return response
            elif response.status_code == 201:
                sleep(sleep_time)
                logger.info(
                    "Report is creating in offline mode. Re-check after "
                    + str(sleep_time)
                    + " seconds"
                )
            elif response.status_code == 202:
                sleep(sleep_time)
                logger.info(
                    "Report is creating in offline mode. Re-check after "
                    + str(sleep_time)
                    + " seconds"
                )
            else:
                if self.should_retry(response):
                    custom_backoff_time = self.backoff_time(response)
                    if custom_backoff_time:
                        raise UserDefinedBackoffException(
                            backoff=custom_backoff_time,
                            request=request,
                            response=response,
                        )
                    else:
                        raise DefaultBackoffException(
                            request=request, response=response
                        )
                elif self.raise_on_http_errors:
                    # Raise any HTTP exceptions that happened in case there were unexpected ones
                    response.raise_for_status()

    def add_constants_to_record(self, record):
        constants = {
            "__productName": self.config["product_name"],
            "__clientName": self.config["client_name"],
        }
        constants.update(json.loads(self.config.get("custom_json", "{}")))
        record.update(constants)
        return record

    def next_page_token(
        self, response: requests.Response
    ) -> Optional[Mapping[str, Any]]:
        return None

    def parse_response(
        self, response: requests.Response, **kwargs
    ) -> Iterable[Mapping]:
        # parse raw TSV data to list of named dicts
        raw_data_lines = response.text.split("\n")
        header = []
        data = []
        for line_n, line in enumerate(raw_data_lines):
            line_values = line.split("\t")

            # skip empty rows
            if not line.strip():
                continue

            # grab header and skip if it is first line
            if line_n == 0:
                header = line_values
                continue

            # zip values list to named dict
            data_item = {}
            for value_n, value in enumerate(line_values):
                data_item[header[value_n]] = value
            data.append(data_item)

        return map(self.add_constants_to_record, data)


class CustomReport(YandexDirectStream):
    def path(self, **kwargs) -> str:
        return ""

    @property
    def primary_key(self) -> Optional[Union[str, List[str], List[List[str]]]]:
        return self.config["fields"][0]

    def request_body_json(self, **kwargs) -> Optional[Mapping]:
        date_range = self.config["transformed_data_range"]
        params = {
            "params": {
                "SelectionCriteria": {},
                "FieldNames": self.config["fields"],
                "ReportName": self.report_name,
                "ReportType": "CUSTOM_REPORT",
                "DateRangeType": date_range["data_range_title"],
                "Format": "TSV",
                "IncludeVAT": "NO",
                "IncludeDiscount": "NO",
            }
        }
        if date_range["data_range_title"] == "CUSTOM_DATE":
            params["params"]["SelectionCriteria"].update(
                {
                    "DateFrom": date_range["date_from"],
                    "DateTo": date_range["date_to"],
                }
            )

        if self.config.get("parsed_filters"):
            params["SelectionCriteria"]["Filter"] = self.config.get("parsed_filters")
        return params

    def request_headers(self, *args, **kwargs) -> Mapping[str, Any]:
        headers = {
            "Authorization": "Bearer " + self.config["auth_token"],
            "Accept-Language": "ru",
            "processingMode": "auto",
            "skipReportHeader": "true",
            "skipReportSummary": "true",
        }
        if self.config.get("client_login"):
            headers["Client-Login"] = self.config["client_login"]
        return headers

    def get_json_schema(self) -> Mapping[str, Any]:
        schema = super().get_json_schema()
        for field_name in self.config["fields"]:
            field_schema_type = CUSTOM_SCHEMA_FIELDS.get(field_name, None)
            if field_schema_type:
                schema["properties"][field_name] = {"type": ["null", field_schema_type]}

        extra_properties = ["__productName", "__clientName"]
        custom_keys = (
            json.loads(self.config["custom_json"]).keys()
            if self.config.get("custom_json")
            else []
        )
        extra_properties.extend(custom_keys)
        for key in extra_properties:
            schema["properties"][key] = {"type": ["null", "string"]}

        return schema


# Source
class SourceYandexDirect(AbstractSource):
    def check_connection(self, logger, config) -> Tuple[bool, Any]:
        date_exclusive_fields = []
        for field_name in config["fields"]:
            if field_name not in CUSTOM_SCHEMA_FIELDS.keys():
                return (
                    False,
                    f"Unknown field name {field_name}. Check https://yandex.ru/dev/direct/doc/reports/fields-list.html for available fields (CUSTOM_REPORT column).",
                )
            if not CUSTOM_SCHEMA_FIELDS[field_name]:
                return (
                    False,
                    f"Field name {field_name} can't be used for CUSTOM_REPORT. Check https://yandex.ru/dev/direct/doc/reports/fields-list.html for available fields (CUSTOM_REPORT column).",
                )
            if CUSTOM_SCHEMA_FIELDS[field_name] == "filter":
                return (
                    False,
                    f"Field name {field_name} is only used for filter data and can't be shown as report field. Check https://yandex.ru/dev/direct/doc/reports/fields-list.html for available fields (CUSTOM_REPORT column).",
                )
            if field_name in ["Date", "Week", "Month", "Quarter", "Year"]:
                date_exclusive_fields.append(field_name)
            if (
                field_name
                in [
                    "Impressions",
                    "Ctr",
                    "AvgImpressionPosition",
                    "WeightedImpressions",
                    "WeightedCtr",
                    "AvgTrafficVolume",
                ]
                and "ClickType" in config["fields"]
            ):
                return (
                    False,
                    f"ClickType field is incompatible with {field_name} field. Check https://yandex.ru/dev/direct/doc/reports/compatibility.html",
                )
            if (
                field_name
                in [
                    "AdFormat",
                    "AdId",
                    "Age",
                    "CarrierType",
                    "Gender",
                    "MobilePlatform",
                    "RlAdjustmentId",
                ]
                and "ImpressionShare" in config["fields"]
            ):
                return (
                    False,
                    f"ImpressionShare field is incompatible with {field_name} field. Check https://yandex.ru/dev/direct/doc/reports/compatibility.html",
                )
        if len(date_exclusive_fields) > 1:
            date_fields = ", ".join(date_exclusive_fields)
            return (
                False,
                f"{date_fields} fields are mutually exclusive: only one of them can be present in the report. Check https://yandex.ru/dev/direct/doc/reports/compatibility.html",
            )

        if set(config["fields"]).intersection(
            ["Criterion", "CriterionId", "CriterionType"]
        ) and set(config["fields"]).intersection(
            ["Criteria", "CriteriaId", "CriteriaType"]
        ):
            return (
                False,
                "Criterion, CriterionId, CriterionType fields are incompatible with Criteria, CriteriaId, CriteriaType fields.",
            )

        if config["date_range"].get("last_days_count"):
            if not config["date_range"].get("last_days_count") in [
                3,
                5,
                7,
                14,
                30,
                90,
                365,
            ]:
                return (
                    False,
                    "Last N days field only can be 3, 5, 7, 14, 30, 90 or 365.",
                )

        try:  #  YYYY-MM-DD
            AUTH_TOKEN = config["auth_token"]
            CLIENT_LOGIN = config.get("client_login", "")
            FIELD_NAMES = [
                "Date",
                "CampaignName",
                "LocationOfPresenceName",
                "Impressions",
                "Clicks",
                "Cost",
            ]
            REPORT_NAME = random_name(9)

            body = {
                "params": {
                    "SelectionCriteria": {},
                    "FieldNames": FIELD_NAMES,
                    "ReportName": REPORT_NAME,
                    "ReportType": "CUSTOM_REPORT",
                    "DateRangeType": "TODAY",
                    "Format": "TSV",
                    "IncludeVAT": "NO",
                    "IncludeDiscount": "NO",
                }
            }

            headers = {
                "Authorization": "Bearer " + AUTH_TOKEN,
                "Accept-Language": "ru",
                "processingMode": "auto",
                "skipReportHeader": "true",
                "skipReportSummary": "true",
            }
            if CLIENT_LOGIN:
                headers["Client-Login"] = CLIENT_LOGIN

            req = requests.post(
                url="https://api.direct.yandex.com/json/v5/reports",
                json=body,
                headers=headers,
            )
            if req.status_code not in [200, 201, 202]:
                return (
                    False,
                    f"Test request status code - {req.status_code}. Body: {req.text}",
                )
        except Exception as e:
            return False, e

        try:
            filter_objects = json.loads(config["filters_json"])
            if not isinstance(filter_objects, list):
                return (
                    False,
                    'Filter JSON must be array of filter objects like this: [{"Fiield": "Year", "Operator": "EQUALS", "Values": ["2021"]}, {...}, {...} ...]. Check https://yandex.ru/dev/direct/doc/reports/filters.html',
                )
            for filter_object in filter_objects:
                if not isinstance(filter_object, dict):
                    return (
                        False,
                        f'Filter JSON must be array of filter objects like this: [{{"Fiield": "Year", "Operator": "EQUALS", "Values": ["2021"]}}, {{...}}, {{...}} ...]. "{filter_object}" filter has validation errors. Check https://yandex.ru/dev/direct/doc/reports/filters.html',
                    )
        except:
            pass

        return True, None

    def transform_config(self, config: Mapping[str, Any]) -> Mapping[str, Any]:
        available_data_range = {
            "today": {"data_range_title": "TODAY"},
            "yesterday": {"data_range_title": "YESTERDAY"},
            "custom_date": {
                "data_range_title": "CUSTOM_DATE",
                "date_from": config["date_range"].get("date_from"),
                "date_to": config["date_range"].get("date_to"),
            },
            "last_n_days": {
                "data_range_title": f"LAST_"
                + str(config["date_range"].get("last_days_count"))
                + "_DAYS"
            },
            "all_time": {"data_range_title": "ALL_TIME"},
        }
        config["transformed_data_range"] = available_data_range[
            config["date_range"]["data_range_type"]
        ]

        config["parsed_filters"] = json.loads(config["filters_json"])
        return config

    def streams(self, config: Mapping[str, Any]) -> List[Any]:
        return [CustomReport(config=self.transform_config(config))]
# # Copyright (c) 2021 Airbyte, Inc., all rights reserved. # # Basic full refresh stream # Raise any HTTP exceptions that happened in case there were unexpected ones # parse raw TSV data to list of named dicts # skip empty rows # grab header and skip if it is first line # zip values list to named dict # Source #  YYYY-MM-DD