#!/usr/bin/env python
#
# Copyright (c) 2022, Alkemy Spa
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import base64
import struct
import zlib
from dataclasses import dataclass
from typing import Dict, List, Union

import sqlparse
from sqlparse.sql import (
    Comment,
    Comparison,
    Function,
    Identifier,
    IdentifierList,
    Parenthesis,
    Statement,
    Where,
)
from sqlparse.tokens import Token

from dbt.adapters.sas.constants import DUMMY_FROM, WORK_SCHEMA

__all__ = ["prepare_query", "Query"]

# Remove FINAL from the keyword list
try:
    del sqlparse.keywords.KEYWORDS['FINAL']
except:
    pass


@dataclass
class Query:
    pre: str
    post: str
    query: str
    ctes: Dict[str, str]  # CTEs  identifier => view name


def get_identifier(token: Token) -> str:
    if isinstance(token, Function):
        token = token[0]
    return str(token)


def is_whitespace(token: Token) -> bool:
    return token.ttype in [Token.Text.Whitespace, Token.Text.Whitespace.Newline]


def get_token_skip_whitespace(tokens: List["Token"]) -> Token:
    "Get the next whitespace token from tokens"
    return pop_token_skip_whitespace(list(tokens))


def pop_token_skip_whitespace(tokens: List["Token"]) -> Token:
    "Pop the next whitespace token from tokens (remote the element from tokens)"
    while tokens:
        token = tokens.pop(0)
        if not is_whitespace(token):
            return token
    return None


def process_cte_definition(token: Identifier, pre: List[str], post: List[str], ctes: Dict[str, str]) -> None:
    # print('>>>', token, type(token), token.ttype)
    tokens = list(token.tokens)
    # Get identifier
    token = pop_token_skip_whitespace(tokens)
    identifier = get_identifier(token)
    # identifier = f"{WORK_SCHEMA}.{identifier}"
    # Check 'AS' keyword
    token = pop_token_skip_whitespace(tokens)
    if token.ttype != Token.Keyword or token.value.upper() != "AS":
        raise Exception(f"Unexpected token {token} (`AS` expected)")
    # Get definition
    token = pop_token_skip_whitespace(tokens)
    if not isinstance(token, Parenthesis):
        raise Exception(f"Unexpected token {token} (parenthesis expected)")
    cte_query = []
    process_parenthesis(token, pre, post, cte_query, ctes)
    query = "".join(cte_query)

    crc = zlib.crc32(query.encode('utf-8'))
    encoded_crc = base64.b32encode(struct.pack('I', crc)).decode('ascii').strip('=')[:7]
    view_identifier = f"{WORK_SCHEMA}.{identifier}_{encoded_crc}"
    ctes[identifier.lower()] = view_identifier

    pre.append(f"create view {view_identifier} as {query};")
    post.append(f"drop view {view_identifier};")


def process_cte(statement: Statement, pre: List[str], post: List[str], ctes: Dict[str, str]) -> None:
    token = pop_token_skip_whitespace(statement)
    if isinstance(token, IdentifierList):  # Multiple CTE
        with_tokens = list(token.tokens)
        while with_tokens:
            with_token = pop_token_skip_whitespace(with_tokens)
            if with_token.ttype != Token.Punctuation:
                process_cte_definition(with_token, pre, post, ctes)
    elif isinstance(token, Identifier):  # Single CTE
        process_cte_definition(token, pre, post, ctes)
    else:
        raise Exception(f"Unexpected token {token}")


def process_parenthesis(
    parenthesis: Parenthesis,
    pre: List[str],
    post: List[str],
    query: List[str],
    ctes: Dict[str, str],
    insert_schema: bool = False,
) -> None:
    p_tokens = parenthesis.tokens
    assert p_tokens[0].value == "("
    assert p_tokens[-1].value == ")"
    query.append("(")
    process_tokens(p_tokens[1:-1], pre, post, query, ctes, insert_schema=insert_schema)
    query.append(")")


def process_identifier(
    identifier: Identifier, pre: List[str], post: List[str], query: List[str], ctes: Dict[str, str], insert_schema: bool = False
) -> None:
    i_tokens = identifier.tokens
    has_schema = False
    while i_tokens:
        token = i_tokens.pop(0)
        if isinstance(token, Parenthesis):
            process_parenthesis(token, pre, post, query, ctes)
            has_schema = False
        elif token.ttype == Token.Punctuation and token.value == ",":
            has_schema = False
            query.append(str(token))
        elif token.ttype == Token.Punctuation and token.value == ".":
            has_schema = True
            query.append(str(token))
        elif isinstance(token, Function) and token.tokens[0].value.upper() == "CAST":
            cast_tokens = list(token.tokens[1:])
            token = pop_token_skip_whitespace(cast_tokens)
            if not isinstance(token, Parenthesis):
                raise Exception(f"Unexpected token {token} (`(` expected)")
            parenthesis_tokens = list(token.tokens[1:])
            token = pop_token_skip_whitespace(parenthesis_tokens)
            query.append(str(token.tokens[0]))
        elif isinstance(token, Comment):
            pass  # skip comments
        else:
            identifier = str(token)
            # print('>>>', token, type(token), token.ttype)
            if insert_schema and identifier.lower() in ctes and not has_schema:
                next_token = get_token_skip_whitespace(i_tokens)
                if next_token and next_token.value.upper() == "AS":
                    identifier = ctes[identifier.lower()]  # don't add the alias (already present)
                elif next_token and isinstance(next_token, Identifier):
                    identifier = ctes[identifier.lower()]  # don't add the alias (already present)
                else:
                    identifier = ctes[identifier.lower()] + " as " + identifier.lower()
            query.append(identifier)
            has_schema = False


def process_from(statement: Statement, pre: List[str], post: List[str], query: List[str], ctes: Dict[str, str]) -> None:
    while statement:
        token = statement.pop(0)
        if isinstance(token, Identifier):
            process_identifier(token, pre, post, query, ctes, insert_schema=True)
        elif isinstance(token, IdentifierList):
            id_tokens = list(token.tokens)
            while id_tokens:
                id_token = id_tokens.pop(0)
                if id_token.ttype == Token.Punctuation:
                    query.append(str(id_token))
                elif isinstance(token, Identifier):
                    process_identifier(id_token, pre, post, query, ctes, insert_schema=True)
                else:
                    query.append(str(id_token))
        elif isinstance(token, Parenthesis):
            process_parenthesis(token, pre, post, query, ctes, insert_schema=True)
        elif isinstance(token, Where):
            process_where(token, pre, post, query, ctes, insert_schema=True)
        elif isinstance(token, Comment):
            pass  # skip comments
        else:
            query.append(str(token))


def process_where(
    where: Union[Where, Comparison, Parenthesis],
    pre: List[str],
    post: List[str],
    query: List[str],
    ctes: Dict[str, str],
    insert_schema: bool = False,
) -> None:
    w_tokens = list(where.tokens)
    while w_tokens:
        token = w_tokens.pop(0)
        # if token.ttype != Token.Text.Whitespace and token.ttype != Token.Text.Whitespace.Newline:
        #     print('>>>', token, type(token), token.ttype)
        if token.ttype == Token.Operator.Comparison and token.value == "!=":
            query.append("<>")
        elif token.ttype == Token.Keyword and token.value.upper() == "TRUE":
            query.append("1")
        elif token.ttype == Token.Keyword and token.value.upper() == "FALSE":
            query.append("0")
        elif isinstance(token, Comparison):
            process_where(token, pre, post, query, ctes, insert_schema)
        elif isinstance(token, Parenthesis):
            process_where(token, pre, post, query, ctes, insert_schema)
        elif isinstance(token, Function):
            process_where(token, pre, post, query, ctes, insert_schema)
        elif isinstance(token, Comment):
            pass  # skip comments
        else:
            query.append(str(token))


def process_tokens(
    statement: Statement, pre: List[str], post: List[str], query: List[str], ctes: Dict[str, str], insert_schema=False
) -> None:
    is_a_select: bool = False
    has_from: bool = False
    while statement:
        token = statement.pop(0)
        # if token.ttype != Token.Text.Whitespace and token.ttype != Token.Text.Whitespace.Newline:
        #     print('>>>', token, type(token), token.ttype)
        if token.ttype == Token.Keyword.DML and str(token).upper() == "SELECT":
            is_a_select = True
            query.append(str(token))
        elif token.ttype == Token.Keyword.CTE:
            process_cte(statement, pre, post, ctes)
        elif token.ttype == Token.Keyword and token.value.upper() == "TRUE":
            query.append("1")
        elif token.ttype == Token.Keyword and token.value.upper() == "FALSE":
            query.append("0")
        elif token.ttype == Token.Keyword and token.value.upper().split(" ")[0] == "UNION":
            if is_a_select and not has_from:
                query.append(DUMMY_FROM)
                has_from = False
            query.append(str(token))
        elif isinstance(token, Identifier):
            process_identifier(token, pre, post, query, ctes, insert_schema=insert_schema)
        elif token.ttype == Token.Keyword and token.value.upper() == "FROM":
            query.append(str(token))
            query.append(" ")
            has_from = True
            process_from(statement, pre, post, query, ctes)
        elif isinstance(token, IdentifierList):
            process_tokens(list(token.tokens), pre, post, query, ctes, insert_schema=insert_schema)
        elif isinstance(token, Comment):
            pass  # skip comments
        else:
            query.append(str(token))
    if is_a_select and not has_from:
        query.append(DUMMY_FROM)


def prepare_query(sql: str) -> Query:
    pre: List[str] = []
    post: List[str] = []
    query: List[str] = []
    ctes: Dict[str, str] = {}  # identifier => view name
    for statement in sqlparse.parse(sql):
        statement = list(statement.tokens)
        process_tokens(statement, pre, post, query, ctes)
    # If the query is a "CREATE VIEW" using a CTEs, don't delete the CTEs
    keep_ctes = pre and "".join(query).upper().strip().startswith("CREATE VIEW")
    if keep_ctes:
        post = []
    return Query(
        pre="\n".join(pre),
        post="\n".join(post),
        query="".join(query),
        ctes=ctes,
    )
