#  Copyright (C) 2023
#  ABM, Moscow
#
#  UNPUBLISHED PROPRIETARY MATERIAL.
#  ALL RIGHTS RESERVED.
#
#  Authors: Vasiliev Ivan <i.vasiliev@technokert.ru>
import enum
from typing import Type

from .clickhouse_types import BigInteger, LongFloat, UnsignedBigInteger, Float, Integer, ClickHouseType
from .column import Column
from .filter_by import Condition
from .abstract.selectable import Selectable


class AggregationFunctionType(enum.StrEnum):
    avg = "avg"
    sum = "sum"
    min = "min"
    max = "max"
    count = "count"


class AggregationFunction(Selectable):

    def __init__(self,
                 function_type: AggregationFunctionType,
                 attribute: Column,
                 condition: Condition | None = None
                 ):
        self.function_type = function_type
        self.column = attribute
        self.condition = condition
        super(AggregationFunction, self).__init__(self.get_clickhouse_type(), nullable=False, name=self.to_sql())

    def get_clickhouse_type(self) -> Type[ClickHouseType]:
        if self.function_type in {AggregationFunctionType.min, AggregationFunctionType.max}:
            return self.column.clickhouse_type
        elif self.function_type == AggregationFunctionType.avg:
            return LongFloat
        elif self.function_type == AggregationFunctionType.count:
            return UnsignedBigInteger
        elif self.function_type == AggregationFunctionType.sum:
            if issubclass(self.column.clickhouse_type, Float):
                return LongFloat
            elif issubclass(self.column.clickhouse_type, Integer):
                return BigInteger

        raise NotImplementedError()

    def to_selector(self) -> str:
        if self.condition:
            return f"{self.function_type}If({self.column.name}, {self.condition.to_sql()}) as {self.name}"

        return f"{self.function_type}({self.column.name}) as {self.name}"

    def to_sql(self):
        return f"{self.column.name}_{self.function_type}"
