# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# See the License at http://www.apache.org/licenses/LICENSE-2.0
# Distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND.

"""
~~~
                      ┌───────────┐
                      │   USER    │
         ┌────────────┤           ◄────────────┐
         │SQL         └───────────┘            │
  ───────┼─────────────────────────────────────┼──────
         │                                     │
   ┌─────▼─────┐                               │
   │ SQL       │                               │
   │ Rewriter  │                               │
   └─────┬─────┘                               │
         │SQL                                  │Results
   ┌─────▼─────┐                         ┌─────┴─────┐
   │           │                         │           │
   │ Parser    │                         │ Executor  │
   └─────┬─────┘                         └─────▲─────┘
         │AST                                  │Plan
   ┌─────▼─────┐      ┌───────────┐      ┌─────┴─────┐
   │ AST       │      │           │      │ Physical  │
   │ Rewriter  │      │ Catalogue │      │ Planner   │
   └─────┬─────┘      └───────────┘      └─────▲─────┘
         │AST               │Schemas           │Plan
   ┌─────▼─────┐      ┌─────▼─────┐      ┌─────┴─────┐
   │ Logical   │ Plan │           │ Plan │           │
   │   Planner ├──────► Binder    ├──────► Optimizer │
   └───────────┘      └───────────┘      └───────────┘

~~~

The SQL Rewriter does the following:
- strips comments
- normalizes whitespace
- temporal extraction (this is non-standard and not part of the parser)

This compensates for missing temporal table support in the SQL parser (sqlparser-rs).
This is relatively complex for what it appears to be doing - it needs to account for
a number of situations whilst being able to reconstruct the SQL query as the parser
would expect it.

For information on temporal tables see:
https://blog.devgenius.io/a-query-in-time-introduction-to-sql-server-temporal-tables-145ddb1355d9

This supports the following syntaxes:

- FOR <timestamp>
- FOR DATES BETWEEN <timestamp> AND <timestamp>
- FOR DATES IN <range>
- FOR DATES SINCE <timestamp>
- FOR LAST <number> DAYS

"""

import datetime
import re
from typing import List
from typing import Tuple

from opteryx.exceptions import InvalidTemporalRangeFilterError
from opteryx.exceptions import UnsupportedSyntaxError
from opteryx.utils import dates

COLLECT_RELATION = {
    r"ANALYZE\sTABLE",
    r"ANTI\sJOIN",
    r"CREATE\sTABLE",
    r"CROSS\sJOIN",
    r"FROM",
    r"FULL\sJOIN",
    r"FULL\sOUTER\sJOIN",
    r"INNER\sJOIN",
    r"JOIN",
    r"LEFT\sANTI\sJOIN",
    r"LEFT\sJOIN",
    r"LEFT\sOUTER\sJOIN",
    r"LEFT\sSEMI\sJOIN",
    r"NATURAL\sJOIN",
    r"RIGHT\sANTI\sJOIN",
    r"RIGHT\sJOIN",
    r"RIGHT\sOUTER\sJOIN",
    r"RIGHT\sSEMI\sJOIN",
    r"SEMI\sJOIN",
}

COLLECT_TEMPORAL = {r"FOR"}

STOP_COLLECTING = {
    r"GROUP\sBY",
    r"HAVING",
    r"LIKE",
    r"LIMIT",
    r"OFFSET",
    r"ON",
    r"ORDER\sBY",
    r"SHOW",
    r"SELECT",
    r"WHERE",
    r"WITH",
    r"USING",
    r";",
    r",",
    r"UNION",
}

COLLECT_ALIAS = {r"AS"}

BOUNDARIES = {r"(", r")"}

FOR_DATE_CLAUSES = {
    r"DATES\sIN\s\w+",
    r"DATES\sBETWEEN\s[^\r\n\t\f\v]AND\s[^\r\n\t\f\v]",
    r"DATES\sSINCE\s\w+",
    r"LAST\s\d+\sDAYS",
}

FUNCTIONS_WITH_FROM_SYNTAX = {"EXTRACT", "SUBSTRING", "TRIM"}

SQL_PARTS = (
    COLLECT_RELATION.union(COLLECT_TEMPORAL)
    .union(STOP_COLLECTING)
    .union(COLLECT_ALIAS)
    .union(FOR_DATE_CLAUSES)
)

COMBINE_WHITESPACE_REGEX = re.compile(r"\r\n\t\f\v+")

# Precompile regex patterns at module level for performance
_KEYWORDS_REGEX = re.compile(
    r"(\,|\(|\)|;|\t|\n|\->>|\->|@>|@>>|\&\&|@\?|"
    + r"|".join([r"\b" + i.replace(r" ", r"\s") + r"\b" for i in SQL_PARTS])
    + r")",
    re.IGNORECASE,
)

# Match ", ', b", b', `
# We match b prefixes separately after the non-prefix versions
_QUOTED_STRINGS_REGEX = re.compile(
    r'("[^"]*"|\'[^\']*\'|\b[bB]"[^"]*"|\b[bB]\'[^\']*\'|\b[rR]"[^"]*"|\b[rR]\'[^\']*\'|`[^`]*`)'
)

# states for the collection algorithm
WAITING: int = 1
RELATION: int = 2
TEMPORAL: int = 4
ALIAS: int = 8
FUNCTION_RELATION: int = 16

WEEKDAYS: List[str] = [
    "MONDAY",
    "TUESDAY",
    "WEDNESDAY",
    "THURSDAY",
    "FRIDAY",
    "SATURDAY",
    "SUNDAY",
]

# Get current time in UTC but without timezone info
NOW = (
    datetime.datetime.now(tz=datetime.timezone.utc)
    .replace(tzinfo=None)
    .replace(hour=0, minute=0, second=0, microsecond=0)
)


def sql_parts(string):
    """
    Split a SQL statement into clauses
    """

    parts = []
    quoted_strings = _QUOTED_STRINGS_REGEX.split(string)
    for i, part in enumerate(quoted_strings):
        if part and part[-1] in ("'", '"', "`"):
            if part[0] in ("b", "B"):
                parts.append(f"CAST({part[1:]} AS VARBINARY)")
                # if there's no alias, we should add one to preserve the input
                if len(quoted_strings) > i + 1:
                    next_token = quoted_strings[i + 1]
                    if next_token.upper().strip().startswith(("FROM ", "JOIN ")):
                        parts.append("AS ")
                        parts.append(f"{part[2:-1]} ")
            elif part[0] in ("r", "R"):
                # We take the raw string and encode it, pass it into the
                # plan as the encoded string and let the engine decode it
                from opteryx.third_party.alantsd import base64

                encoded_part = base64.encode(part[2:-1].encode()).decode()
                # if there's no alias, we should add one to preserve the input
                parts.append(f"BASE64_DECODE('{encoded_part}')")
                if len(quoted_strings) > i + 1:
                    next_token = quoted_strings[i + 1]
                    if next_token.upper().strip().startswith(("FROM ", "JOIN ")):
                        parts.append("AS ")
                        parts.append(f"{part[2:-1]} ")
            else:
                parts.append(part)
        else:
            for subpart in _KEYWORDS_REGEX.split(part):
                subpart = subpart.strip()
                if subpart:
                    parts.append(subpart)

    return parts


def parse_range(fixed_range):  # pragma: no cover
    fixed_range = fixed_range.upper()

    if fixed_range in ("PREVIOUS_MONTH", "LAST_MONTH"):
        # end the day before the first of this month
        end = NOW.replace(day=1) - datetime.timedelta(days=1)
        # start the first day of that month
        start = end.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
    elif fixed_range == "THIS_MONTH":
        # start the first day of this month
        start = NOW.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
        # end today
        end = NOW
    else:
        if parse_date(fixed_range):
            raise InvalidTemporalRangeFilterError(
                f"`THIS_MONTH`, `LAST_MONTH` expected, got `{fixed_range}`"
            )
        if fixed_range == "LAST":
            raise InvalidTemporalRangeFilterError(
                "`LAST` is not a valid range. Did you mean 'FOR LAST count DAYS'?"
            )
        raise InvalidTemporalRangeFilterError(f"Unknown temporal range `{fixed_range}`")

    return start, end


def parse_date(date, end: bool = False):  # pragma: no cover
    if not date:
        return None
    if len(date) > 30:
        return None

    if date == "TODAY":
        return NOW
    if date == "YESTERDAY":
        return NOW - datetime.timedelta(days=1)

    if date in WEEKDAYS:
        # Find the weekday number (0=Monday, 1=Tuesday, ..., 6=Sunday)
        target_weekday = WEEKDAYS.index(date)
        # Find the current weekday number
        current_weekday = NOW.weekday()

        # Calculate how many days to subtract to get the last occurrence of the target weekday
        days_to_subtract = (current_weekday - target_weekday) % 7
        if days_to_subtract == 0:
            # If today is the target weekday, adjust to get the last week's same day
            days_to_subtract = 7

        # Calculate the most recent date for the target weekday
        most_recent_day = (NOW - datetime.timedelta(days=days_to_subtract)).replace(
            hour=0, minute=0, second=0, microsecond=0
        )
        return most_recent_day

    if date[0] == date[-1] and date[0] in ("'", '"', "`"):
        date = date[1:-1]

    # Parse ISO dates and convert to naive UTC
    parsed = dates.parse_iso(date)
    if parsed:
        # Convert to UTC and make naive
        parsed = parsed.replace(tzinfo=None)
        if end and len(date) <= 12:
            return parsed.replace(hour=23, minute=59)
        return parsed

    return None


def _temporal_extration_state_machine(
    parts: List[str],
) -> Tuple[List[Tuple[str, str]], str]:
    """
    Utilizes a state machine to extract the temporal information from the query
    and maintain the relation to filter information.

    Parameters:
        parts: List[str]
            SQL statement parts

    Returns:
        Tuple containing two lists, first with the temporal filters, second with the remaining SQL parts.
    """
    # We use a four state machine to extract the temporal information from the query
    # and maintain the relation to filter information.
    #
    # We separate out the two key parts of the algorithm, first we determine the state,
    # then we work out if the state transition means we should do something.
    #
    # We're essentially using a bit mask to record state and transitions.

    in_special_function = False
    special_function_brackets = 0

    state = WAITING
    relation = ""
    temporal = ""
    query_collector = []
    temporal_range_collector = []
    open_count = 0
    for part in parts:
        # record the current state
        transition = [state]
        comparable_part = part.upper().replace(" ", r"\s")

        if comparable_part in FUNCTIONS_WITH_FROM_SYNTAX:
            in_special_function = True
            special_function_brackets = open_count

        # work out what our current state is
        elif comparable_part in BOUNDARIES:
            if comparable_part == "(":
                open_count += 1
            if comparable_part == ")":
                open_count -= 1
                if in_special_function and open_count == special_function_brackets:
                    in_special_function = False
            # function relations, like FAKE(234,234) need the items between the
            # brackets be be consumed
            state = WAITING if relation == "" else FUNCTION_RELATION

        if not in_special_function:
            if comparable_part in STOP_COLLECTING:
                if state == FUNCTION_RELATION and open_count > 0:
                    pass
                else:
                    state = WAITING
            if comparable_part in COLLECT_RELATION:
                state = RELATION
            if comparable_part in COLLECT_TEMPORAL:
                state = TEMPORAL
            if comparable_part in COLLECT_ALIAS:
                state = ALIAS
        transition.append(state)

        # based on what the state was and what it is now, do something
        if in_special_function:
            pass
        elif transition == [TEMPORAL, TEMPORAL]:
            temporal = (temporal + " " + part).strip()
        elif (
            transition
            in (
                [WAITING, WAITING],
                [TEMPORAL, RELATION],
                [RELATION, RELATION],
                [RELATION, WAITING],
                [ALIAS, RELATION],
                [ALIAS, WAITING],  # probably
                [FUNCTION_RELATION, RELATION],
            )
            and relation
        ):
            temporal_range_collector.append((relation, temporal))
            relation = ""
            temporal = ""
            if comparable_part == ",":
                state = RELATION
        elif transition == [RELATION, RELATION]:
            relation = part
        elif transition == [WAITING, TEMPORAL]:
            raise InvalidTemporalRangeFilterError(
                "Temporal `FOR` statements must directly follow the dataset they apply to."
            )
        else:
            pass

        if state != TEMPORAL:
            query_collector.append(part)

    # if we're at the end of we have a relation, emit it
    if relation:
        temporal_range_collector.append((relation, temporal))

    return temporal_range_collector, " ".join(query_collector)


def extract_temporal_filters(parts: str):  # pragma: no cover
    import shlex

    # extract the raw temporal information
    initial_collector, sql = _temporal_extration_state_machine(parts)

    final_collector = []

    for relation, for_date_string in initial_collector:
        start_date = None
        end_date = None

        for_date_string = for_date_string.upper()
        for_date = parse_date(for_date_string)

        if for_date:
            start_date = for_date
            end_date = for_date
            if for_date_string in ("TODAY", "YESTERDAY") or len(for_date_string) <= 12:
                start_date = start_date.replace(hour=0)
                end_date = end_date.replace(hour=23, minute=59)

        elif for_date_string.startswith("DATES BETWEEN "):
            parts = shlex.split(for_date_string)

            if len(parts) != 5:
                raise InvalidTemporalRangeFilterError(
                    "Invalid temporal range, expected format `FOR DATES BETWEEN <start> AND <end>`."
                )
            if parts[3] != "AND":
                raise InvalidTemporalRangeFilterError(
                    f"Invalid temporal range, expected `AND`, found `{parts[3]}`."
                )

            start_date = parse_date(parts[2])
            end_date = parse_date(parts[4], end=True)

            if start_date is None:
                raise InvalidTemporalRangeFilterError(
                    f"Invalid temporal range, expected a date for start of range, found `{parts[2]}`."
                )
            if end_date is None:
                raise InvalidTemporalRangeFilterError(
                    f"Invalid temporal range, expected a date for end of range, found `{parts[4]}`."
                )
            if start_date > end_date:
                raise InvalidTemporalRangeFilterError(
                    "Invalid temporal range, start of range is after end of range."
                )

        elif for_date_string.startswith("LAST "):
            parts = shlex.split(for_date_string)
            if len(parts) != 3 or parts[2] != "DAYS":
                raise InvalidTemporalRangeFilterError(
                    "Invalid temporal range, expected format `FOR LAST <number> DAYS`."
                )
            interval = int(parts[1])
            # Get current time in UTC but without timezone info
            end_date = NOW.replace(hour=23, minute=59, second=0, microsecond=0)
            start_date = (end_date - datetime.timedelta(interval)).replace(
                hour=0, minute=0, second=0, microsecond=0
            )

        elif for_date_string.startswith("DATES IN "):
            parts = shlex.split(for_date_string)
            start_date, end_date = parse_range(parts[2])

        elif for_date_string.startswith("DATES SINCE "):
            parts = shlex.split(for_date_string)
            start_date = parse_date(parts[2])
            # Get current time in UTC but without timezone info
            end_date = NOW

        elif for_date_string:
            raise InvalidTemporalRangeFilterError(
                f"Unable to interpret temporal filter `{for_date_string}`"
            )

        final_collector.append(
            (
                relation,
                start_date,
                end_date,
            )
        )

    # we've rewritten the sql so make it sqlparser-rs compatible
    return sql, final_collector


def rewrite_explain(parts: list) -> list:
    """
    The parser does not support MERMAID format.

    We rewrite it to GRAPHVIZ, we don't support.
    """
    if parts[0] == "EXPLAIN ANALYZE FORMAT GRAPHVIZ":
        raise UnsupportedSyntaxError("GRAPHVIZ format is not supported")
    if parts[0] == "EXPLAIN ANALYZE FORMAT JSON":
        raise UnsupportedSyntaxError("JSON format is not supported")
    if parts[0].upper() == "EXPLAIN ANALYZE FORMAT MERMAID":
        parts[0] = "EXPLAIN ANALYZE FORMAT GRAPHVIZ"
    return parts


def do_sql_rewrite(statement):
    parts = sql_parts(statement)
    parts = rewrite_explain(parts)
    return extract_temporal_filters(parts)
