# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_terms.ipynb.

# %% auto 0
__all__ = ['CURRENT_ROW', 'Term', 'Value', 'NullValue', 'FieldBase', 'Field', 'ArithmeticExpression', 'Criteria', 'kwargs_func',
           'DelayedFunc', 'custom_func', 'OverClause', 'Preceding', 'Following', 'Case']

# %% ../nbs/02_terms.ipynb 2
import inspect
from functools import partial
from fastcore.foundation import patch
from .utils import execute, delegates

# %% ../nbs/02_terms.ipynb 4
class Term:
    def __init__(self) -> None:
        self.alias = None
    
    def as_(self, alias):
        self.alias = alias
        return self


class Value(Term):
    """A simple wrapper for a value that will be used in a query."""
    def __init__(self, value) -> None:
        super().__init__()
        self.value = value
        self.get_sql = self.execute

    def execute(self, **kwargs):
        if type(self.value) is str:
            return f"'{self.value}' as {self.alias}" if self.alias else f"'{self.value}'" 
        elif type(self.value) in (list, tuple):
            return '(' + ', '.join([f"{item}" if type(item) in (int, float) else f"'{item}'" for item in self.value]) + ')'
        elif hasattr(self.value, 'get_sql'):
            return self.value.get_sql(**kwargs)
        else:
            return f"{self.value} as {self.alias}" if self.alias else f"{self.value}"


class NullValue(Term):
    """NULL value for use in queries."""
    def __init__(self) -> None:
        super().__init__()
        self.get_sql = self.execute

    def execute(self, **kwargs):
        return f"NULL as {self.alias}" if self.alias else f"NULL"

# %% ../nbs/02_terms.ipynb 7
class FieldBase(Term):
    """Collection of methods to convert to ArithmeticExpression and Criteria"""
    def __add__(self, other):
        return ArithmeticExpression(self, '+', other)

    def __radd__(self, other):
        return ArithmeticExpression(other, '+', self)

    def __sub__(self, other):
        return ArithmeticExpression(self, '-', other)

    def __rsub__(self, other):
        return ArithmeticExpression(other, '-', self)

    def __mul__(self, other):
        return ArithmeticExpression(self, '*', other)
        
    def __rmul__(self, other):
        return ArithmeticExpression(other, '*', self)

    def __truediv__(self, other):
        return ArithmeticExpression(self, '/', other)

    def __rtruediv__(self, other):
        return ArithmeticExpression(other, '/', self)

    def __gt__(self, other):
        return Criteria(self, '>', other)

    def __ge__(self, other):
        return Criteria(self, '>=', other)
        
    def __eq__(self, other):
        return Criteria(self, '=', other)

    def __ne__(self, other):
        return Criteria(self, '<>', other)

    def __le__(self, other):
        return Criteria(self, '<=', other)

    def __lt__(self, other):
        return Criteria(self, '<', other)

    gt = __gt__
    ge = __ge__
    eq = __eq__
    ne = __ne__
    le = __le__
    lt = __lt__

    def isnull(self):
        return Criteria(self, 'IS', NullValue())

    def notnull(self):
        return Criteria(self, 'IS NOT', NullValue())

    def like(self, other):
        return Criteria(self, 'LIKE', Value(other))
    
    def ilike(self, other):
        return Criteria(self, 'ILIKE', Value(other))

    def not_like(self, other):
        return Criteria(self, 'NOT LIKE', Value(other))

    def not_ilike(self, other):
        return Criteria(self, 'NOT ILIKE', Value(other))
    
    def isin(self, other:list):
        return Criteria(self, 'IN', Value(other))
    
    def notin(self, other):
        return Criteria(self, 'NOT IN', Value(other))

    #TODO: implement new Criteria, and execute for list


class Field(FieldBase):
    """A simple wrapper for a field that will be used in a query. Quotes can be added to the field name by setting the `quote_char` parameter in get_sql() method."""
    def __init__(self, name) -> None: 
        self.name = name
        self.alias = None
        self.get_sql = self.execute

    def quoted_name(self, quote_char):
        name_list = self.name.split('.')
        name = '.'.join([f"{quote_char}{n}{quote_char}" for n in name_list])
        return name

    def execute(self, **kwargs):
        q = kwargs.get('quote_char', '') or ''
        quoted_name = self.quoted_name(q)
        if self.alias:
            return f"{quoted_name} AS {q}{self.alias}{q}"
        else:
            return quoted_name
    

class ArithmeticExpression(FieldBase):
    """Constructor for arithmetic expressions from two terms and automatically adds parentheses in appropriate places."""
    add_order = ["+", "-"]

    def __init__(self, this, op, other) -> None:
        super().__init__()
        self.this, self.op, self.other = this, op, other
        self.alias = None
        self.get_sql = self.execute

    def left_needs_parens(self, left_op, curr_op) -> bool:
        """
        Returns true if the expression on the left of the current operator needs to be enclosed in parentheses.
        :param current_op:
            The current operator.
        :param left_op:
            The highest level operator of the left expression.
        """
        if left_op is None or curr_op in self.add_order:
            # If the left expression is a single item.
            # or if the current operator is '+' or '-'.
            return False
        
        # The current operator is '*' or '/'. 
        # If the left operator is '+' or '-', we need to add parentheses:
        # e.g. (A + B) / ..., (A - B) / ...
        # Otherwise, no parentheses are necessary:
        # e.g. A * B / ..., A / B / ...
        return left_op in self.add_order

    def right_needs_parens(self, curr_op, right_op) -> bool:
        """
        Returns true if the expression on the right of the current operator needs to be enclosed in parentheses.
        :param current_op:
            The current operator.
        :param right_op:
            The highest level operator of the right expression.
        """
        if right_op is None:
            # If the right expression is a single item.
            return False
        # If the right operator is '+' or '-', we always add parentheses:
        # e.g. ... - (A + B), ... - (A - B), ... + (A + B)
        # Otherwise, no parentheses are necessary:
        # e.g. ... - A / B, ... - A * B
        return right_op in self.add_order

    def execute(self, **kwargs):
        q = kwargs.get('quote_char', '') or ''
        if self.this.__class__ is ArithmeticExpression:
            this = self.this.execute(**kwargs)
            this = f"({this})" if self.left_needs_parens(self.this.op, self.op) else this
        elif self.this.__class__ is Field:
            this = self.this.alias or execute(self.this, **kwargs)
        else:
            this = execute(self.this, **kwargs)

        if self.other.__class__ is ArithmeticExpression:
            other = self.other.execute(**kwargs)
            other = f"({other})" if self.right_needs_parens(self.op, self.other.op) else other
        elif self.other.__class__ is Field:
            other = self.other.alias or execute(self.other, **kwargs)
        else:
            other = execute(self.other, **kwargs)

        if self.alias:
            return f"{this} {self.op} {other} AS {q}{self.alias}{q}"
        else:
            return f"{this} {self.op} {other}"


class Criteria:
    """Constructor for criteria from two terms and automatically adds parentheses in appropriate places."""
    compose_ops = ('and', 'or')

    def __init__(self, this, op, other) -> None:
        super().__init__()
        self.this, self.op, self.other = this, op, other
        self.add_parentheses = False
        self.not_ = False
        self.get_sql = self.execute

    def compose_criteria(self, op, other):
        """Add parentheses when operator in criteria is different from op"""
        if other.__class__ is Criteria:
            if self.op in self.compose_ops and op in self.compose_ops and self.op != op:
                self.add_parentheses = True
            if other.op in self.compose_ops and op in self.compose_ops and other.op != op:
                other.add_parentheses = True
        return Criteria(self, op, other)

    @staticmethod
    def resolve(obj, **kwargs):
        if obj.__class__ is Criteria:
            obj_c = obj.execute(**kwargs)
            obj_c = f"({obj_c})" if obj.add_parentheses == True else obj_c
        else:
            obj_c = execute(obj, **kwargs)
        return obj_c
    
    def execute(self, **kwargs):
        this = self.resolve(self.this, **kwargs)
        other = self.resolve(self.other, **kwargs)
        self.result =  f"{this} {self.op} {other}"
        if self.not_:
            return f"NOT ({self.result})"
        else:
            return self.result
    
    def negate(self):
        self.not_ = True
        return self

    def __and__(self, __o):
        return self.compose_criteria('and', __o)

    def __or__(self, __o):
        return self.compose_criteria('or', __o)

# %% ../nbs/02_terms.ipynb 10
def kwargs_func(func, *args, **kwargs):
    "Allow arbitrary kwargs. Only pass those kwargs that are specified in func to func."
    sig = inspect.signature(func)
    param = sig.parameters
    func_kwargs = {k:v for k, v in param.items() if v.default!=inspect._empty}
    kwargs = {k:v for k, v in kwargs.items() if k in func_kwargs}
    return func(*args, **kwargs)


class DelayedFunc(FieldBase):
    """Delay the execution of stored function until exec is run."""
    def __init__(self, 
                 func, 
                 args,
                 kwargs, 
                 window_func=True # whether this function is a window function
                 ) -> None: 
        super().__init__()
        self.func = func
        self.args = args
        self.kwargs = kwargs
        self.window_func = window_func
        self.get_sql = self.execute
    

    def execute(self, **kwargs):
        """keyword arguments can be overwritten with any provided new kwargs."""
        self.kwargs.update(kwargs)
        # recursively resolve all delayed functions
        args = (execute(arg, **self.kwargs) for arg in self.args) 
        func = self.func(*args, **self.kwargs)
        if self.alias:
            return f"{func} AS {self.alias}"
        else:
            return func


def custom_func(func=None, window_func=False, dialect=None):
    """return Field"""
    if func is None: 
        return partial(custom_func, window_func=window_func, dialect=dialect)
    else:
        if dialect is None:
            @delegates(func)
            def wrapper(*args, **kwargs):
                dlf = DelayedFunc(func, args, kwargs, window_func=window_func)
                return dlf
        else:
            # get previously defined func
            func_name = func.__name__
            ori_func = globals().get(func_name)
            @delegates(func)
            def wrapper(*args, **kwargs):
                def new_func(*args, **kwargs):
                    # if the dialect is different from the one defined for this new function, fall back to the original func
                    if kwargs['dialect'] != dialect and ori_func:
                        f = ori_func().func
                    else:
                        f = func
                    return kwargs_func(f, *args, **kwargs)

                # make new delayed function
                dlf = DelayedFunc(new_func, args, kwargs, window_func=window_func)
                return dlf
        return wrapper

# %% ../nbs/02_terms.ipynb 16
class OverClause:
    """Constructor for OVER clause."""
    def __init__(self, expression) -> None:
        if not (
            (hasattr(expression, 'window_func')
             and 
             expression.window_func==True)
            or 
            (type(expression) is str)
        ):
            raise ValueError(f"Expression has to be of the ArithmeticExpression type or a Window Function (DelayedFunc with window_func=True)!!")
        self.expr = expression
        self.alias = None
        self.rows_flag = False
        self.range_flag = False
        self.d = {}
        self.get_sql = self.execute

    def over(self, q):
        self.d['PARTITION BY'] = q
        return self

    def orderby(self, q):
        self.d['ORDER BY'] = q
        return self

    def _check_rows_or_range(self):
        if self.rows_flag==True: 
            raise ValueError(f"ROWS already set!")
        if self.range_flag==True:
            raise ValueError(f"RANGE already set!")

    def rows(self, start, end):
        self._check_rows_or_range()
        self.rows = True
        self.d['ROWS'] = (start, end)
        return self

    def range(self, start, end):
        self._check_rows_or_range()
        self.range = True
        self.d['RANGE'] = (start, end)
        return self

    def as_(self, alias):
        self.alias = alias
        return self

    def _resolve_over_statement(self, **kwargs):
        sql = []
        for k in ['PARTITION BY', 'ORDER BY', 'ROWS', 'RANGE']:
            if k in self.d:
                if k in ['PARTITION BY', 'ORDER BY']:
                    rslvd = f"{k} {execute(self.d[k], **kwargs)}"
                else:
                    start, end = self.d[k]
                    rslvd = f"{k} BETWEEN {execute(start, **kwargs)} AND {execute(end, **kwargs)}"
                sql.append(rslvd)
        return f"{execute(self.expr, **kwargs)} OVER ({' '.join(sql)})"

    def execute(self, **kwargs):
        sql = self._resolve_over_statement(**kwargs)
        if self.alias:
            return f"{sql} AS {self.alias}"
        else:
            return sql


class Preceding:
    def __init__(self, N=None) -> None:
        self.N = N
        self.get_sql = self.execute

    def execute(self, **kwargs):
        if self.N:
            return f"{self.N} PRECEDING"
        else:
            return "UNBOUNDED PRECEDING"


class Following:
    def __init__(self, N=None) -> None:
        self.N = N
        self.get_sql = self.execute

    def execute(self, **kwargs):
        if self.N:
            return f"{self.N} FOLLOWING"
        else:
            return "UNBOUNDED FOLLOWING"


CURRENT_ROW = "CURRENT_ROW"

# %% ../nbs/02_terms.ipynb 17
@patch
def over(self:ArithmeticExpression, partition_by):
    return OverClause(self).over(partition_by)

@patch
def over(self:DelayedFunc, partition_by):
    if self.window_func==True:
        return OverClause(self).over(partition_by)
    else:
        raise ValueError(f"This function is not a window function!!")


# %% ../nbs/02_terms.ipynb 20
class Case(Term):
    """Constructor for CASE statement."""
    def __init__(self) -> None:
        self.dp = []
        self.alias = None
        self.get_sql = self.execute

    def check_prev(self, statement):
        if self.dp:
            prev = self.dp[0][0]
            if prev == statement:
                return True
        return False

    def when(self, q, then):
        if self.check_prev('ELSE'):
            raise ValueError(f"'WHEN' can not follow 'ELSE'!")
        self.dp.append(('WHEN', q, then))
        return self

    def else_(self, q):
        self.dp.append(('ELSE', q))
        return self

    def execute(self, **kwargs):
        sql = ["CASE"]
        for item in self.dp:
            if item[0] == 'WHEN':
                q_resolved = f"WHEN {execute(item[1], **kwargs)} THEN {execute(item[2], **kwargs)}"
            else:
                q_resolved = f"ELSE {execute(item[1], **kwargs)}"
            sql.append(q_resolved)
        if self.alias:
            sql.append(f"END AS {self.alias}")
        else:
            sql.append("END")
        return '\n'.join(sql)
