# Copyright 2018-2022 contributors to the OpenLineage project
# SPDX-License-Identifier: Apache-2.0

import logging
from typing import List, Tuple, Optional, Set, Any

import sqlparse
from sqlparse.sql import T, TokenList, Parenthesis, Identifier, IdentifierList
from sqlparse.tokens import Punctuation


log = logging.getLogger(__name__)

# Do not import from this package directly. Instead, import from openlineage.common.sql


def provider():
    return "python"


class DbTableMeta:
    def __init__(self, value: str):
        parts = value.strip().split('.')
        if len(parts) > 3:
            raise ValueError(
                f"Expected 'database.schema.table', found '{value}'."
            )
        self.database = self._get_database(parts)
        self.schema = self._get_schema(parts)
        self.name = self._get_table(parts)
        self.qualified_name = self._get_qualified_name()

    def has_database(self) -> bool:
        return self.database is not None

    def has_schema(self) -> bool:
        return self.schema is not None

    def _get_database(self, parts) -> str:
        # {database.schema.table}
        return parts[0] if len(parts) == 3 else None

    def _get_schema(self, parts) -> str:
        # {database.schema.table) or {schema.table}
        return parts[1] if len(parts) == 3 else (
            parts[0] if len(parts) == 2 else None
        )

    def _get_table(self, parts) -> str:
        # {database.schema.table} or {schema.table} or {table}
        return parts[2] if len(parts) == 3 else (
            parts[1] if len(parts) == 2 else parts[0]
        )

    def _get_qualified_name(self) -> str:
        return (
            f"{self.database}.{self.schema}.{self.name}"
            if self.has_database() else (
                f"{self.schema}.{self.name}" if self.has_schema() else self.name
            )
        )

    def __hash__(self):
        return hash(self.qualified_name)

    def __eq__(self, other):
        return self.database == other.database and \
               self.schema == other.schema and \
               self.name == other.name and \
               self.qualified_name == other.qualified_name

    def __repr__(self):
        # Return the string representation of the instance
        return (
            f"DbTableMeta({self.database!r},{self.schema!r},"
            f"{self.name!r},{self.qualified_name!r})"
        )

    def __str__(self):
        # Return the fully qualified table name as the string representation
        # of this object, otherwise the table name only
        return (
            self.qualified_name
            if (self.has_database() or self.has_schema()) else self.name
        )


def _is_in_table(token):
    return token.match(T.Keyword, [
        'FROM',
        'INNER JOIN',
        'JOIN',
        'FULL JOIN',
        'FULL OUTER JOIN',
        'LEFT JOIN',
        'LEFT OUTER JOIN',
        'RIGHT JOIN',
        'RIGHT OUTER JOIN'
    ])


def _is_out_table(token):
    return token.match(T.DML, ['INSERT', 'UPDATE']) or token.match(T.Keyword, ['INTO'])


def _parse_ident(ident: Identifier, default_schema: Optional[str] = None) -> str:
    # Extract table name from possible schema.table naming
    token_list = ident.flatten()
    table_name = next(token_list).value
    try:
        # Determine if the table contains the schema
        # separated by a dot (format: 'schema.table')
        dot = next(token_list)
        if dot.match(Punctuation, '.'):
            table_name += dot.value
            table_name += next(token_list).value
            # And again, to match bigquery's 'database.schema.table'
            try:
                dot = next(token_list)
                if dot.match(Punctuation, '.'):
                    table_name += dot.value
                    table_name += next(token_list).value
            except StopIteration:
                # Do not insert database name if it's not specified
                pass
        elif default_schema:
            table_name = f'{default_schema}.{table_name}'
    except StopIteration:
        if default_schema:
            table_name = f'{default_schema}.{table_name}'

    table_name = table_name.replace('`', '')
    return table_name


def _parse_ident_list(token: IdentifierList, default_schema: Optional[str] = None) -> List[str]:
    # Handle "comma separated joins" as opposed to explicit JOIN keyword
    tables = []
    gidx = 0
    tables.append(_parse_ident(token.token_first(skip_ws=True, skip_cm=True), default_schema))
    gidx, punc = token.token_next(gidx, skip_ws=True, skip_cm=True)
    while punc and punc.value == ',':
        gidx, name = token.token_next(gidx, skip_ws=True, skip_cm=True)
        tables.append(_parse_ident(name, default_schema))
        gidx, punc = token.token_next(gidx)
    return tables


def _get_in_tables(
        tokens,
        idx,
        default_schema: Optional[str] = None
) -> Tuple[int, List[DbTableMeta]]:
    # Extract table identified by preceding SQL keyword at '_is_in_table'
    idx, token = tokens.token_next(idx=idx)
    tables = []
    if isinstance(token, IdentifierList):
        tables = _parse_ident_list(token, default_schema)
    else:
        tables.append(_parse_ident(token, default_schema))

    return idx, [DbTableMeta(table) for table in tables]


def _get_out_tables(
        tokens,
        idx,
        default_schema: Optional[str] = None
) -> Tuple[int, List[DbTableMeta]]:
    # Extract table identified by preceding SQL keyword at '_is_out_table'
    idx, token = tokens.token_next(idx=idx)
    # Consume all keywords until reaching to an identifier since keywords
    # can be a sequence (e.g., `INSERT INTO`) in the case of out table.
    while token.is_keyword:
        idx, token = tokens.token_next(idx=idx)
    tables = []
    if isinstance(token, IdentifierList):
        tables = _parse_ident_list(token, default_schema)
    else:
        tables.append(_parse_ident(token, default_schema))

    return idx, [DbTableMeta(table) for table in tables]


class SqlMeta:
    def __init__(self, in_tables: List[DbTableMeta], out_tables: List[DbTableMeta]):
        self.in_tables = in_tables
        self.out_tables = out_tables

    def __repr__(self):
        return f"SqlMeta({self.in_tables!r},{self.out_tables!r})"

    def add_in_tables(self, in_tables: List[DbTableMeta]):
        for in_table in in_tables:
            self.in_tables.append(in_table)

    def add_out_tables(self, out_tables: List[DbTableMeta]):
        for out_table in out_tables:
            self.out_tables.append(out_table)


class SqlParser:
    """
    This class parses a SQL statement.
    """

    def __init__(self, default_schema: Optional[str] = None):
        # In some cases like bigquery we can always get schema/dataset ID from client
        self.default_schema = default_schema
        self.ctes: Set[Any] = set()
        self.intables: Set[Any] = set()
        self.outtables: Set[Any] = set()

    def recurse(self, tokens: TokenList) -> SqlMeta:
        in_tables, out_tables = set(), set()
        idx, token = tokens.token_next_by(t=T.Keyword)
        while token:

            # Main parser switch
            if self.is_cte(token):
                cte_name, cte_intables = self.parse_cte(idx, tokens)
                for intable in cte_intables:
                    if intable not in self.ctes:
                        in_tables.add(intable)
            elif _is_in_table(token):
                idx, extracted_tables = _get_in_tables(tokens, idx, self.default_schema)
                for table in extracted_tables:
                    if table.name not in self.ctes:
                        in_tables.add(table)
            elif _is_out_table(token):
                idx, extracted_tables = _get_out_tables(tokens, idx, self.default_schema)
                out_tables.add(extracted_tables[0])  # assuming only one out_table

            idx, token = tokens.token_next_by(t=T.Keyword, idx=idx)

        return SqlMeta(list(in_tables), list(out_tables))

    def is_cte(self, token: T):
        return token.match(T.Keyword.CTE, values=['WITH'])

    def parse_cte(self, idx, tokens: TokenList):
        gidx, group = tokens.token_next(idx, skip_ws=True, skip_cm=True)

        # handle recursive keyword
        if group.match(T.Keyword, values=['RECURSIVE']):
            gidx, group = tokens.token_next(gidx, skip_ws=True, skip_cm=True)

        if not group.is_group:
            return [], None

        # get CTE name
        offset = 1
        cte_name = group.token_first(skip_ws=True, skip_cm=True)
        self.ctes.add(cte_name.value)

        # AS keyword
        offset, as_keyword = group.token_next(offset, skip_ws=True, skip_cm=True)
        if not as_keyword.match(T.Keyword, values=['AS']):
            raise RuntimeError(f"CTE does not have AS keyword at index {gidx}")

        offset, parens = group.token_next(offset, skip_ws=True, skip_cm=True)
        if isinstance(parens, Parenthesis) or parens.is_group:
            # Parse CTE using recursion.
            return cte_name.value, self.recurse(TokenList(parens.tokens)).in_tables
        raise RuntimeError(f"Parens {parens} are not Parenthesis at index {gidx}")


# Do not import this directly. Instead, import openlineage.common.sql.parse
def parse(
    sql: List[str],
    dialect: Optional[str] = None,
    default_schema: Optional[str] = None
) -> SqlMeta:
    if sql is None:
        raise ValueError("A sql statement must be provided.")

    sql_statements = []
    for statement in sql:
        # Tokenize the SQL statement
        sql_statements += sqlparse.parse(statement)

    sql_parser = SqlParser(default_schema)
    sql_meta = SqlMeta([], [])

    for sql_statement in sql_statements:
        tokens = TokenList(sql_statement.tokens)
        log.debug(f"Successfully tokenized sql statement: {tokens}")

        result = sql_parser.recurse(tokens)

        # Add the in / out tables (if any) to the sql meta
        sql_meta.add_in_tables(result.in_tables)
        sql_meta.add_out_tables(result.out_tables)

    return sql_meta
