from typing import Dict, Any

from .clause_tokenizer import ClauseTokenizer
from .json_parser import JsonParser
from .utils import CaseInsensitiveDict


class FromParser:
    def __init__(self):
        self.clauses = None

    def parse(self, from_clause) -> Dict[str, str]:
        """
        Parse a FROM-clause in a PARTIQL query
        :param from_clause: a string of format `a AS b, x AS y` where `a` and `x` can contain commas
        :return: a dictionary of format `[b:a, y:x]`
        """
        clauses: Dict[str, Any] = dict()
        section = None  # NAME/AS/ALIAS
        current_phrase = ""
        name = alias = None
        from_clause_parser = ClauseTokenizer(from_clause)
        while True:
            c = from_clause_parser.next()
            if c is None:
                break
            if c == "[":
                if section is None:
                    # Beginning of the FROM-clause - probably a document in its own right, instead of a name
                    current_phrase = "[" + from_clause_parser.next_until(["]"]) + "]"
                    section = "NAME"
                    continue
            if c == " ":
                if section == "AS" and current_phrase.upper() == "AS":
                    current_phrase = ""
                    section = "ALIAS"
                elif section == "NAME":
                    name = current_phrase
                    current_phrase = ""
                    section = "AS"
                elif section == "ALIAS":
                    alias = current_phrase  # noqa
                    current_phrase = ""
                    section = "NAME"
                continue
            if c == ",":
                if section == "NAME":
                    clauses[current_phrase] = current_phrase
                elif section == "ALIAS":
                    clauses[current_phrase] = name
                    section = "NAME"
                current_phrase = ""
                from_clause_parser.skip_white_space()
                # new phrase
                section = None
                continue

            if section is None:
                section = "NAME"

            current_phrase += c
        if section == "NAME":
            clauses[current_phrase] = current_phrase
        else:
            clauses[current_phrase] = name
        #
        # One FROM clause may point to the alias of another
        # {'s': 'sensors', 'r': 's.readings'} --> {'s': 'sensors', 'r': 'sensors.readings'}
        aliases = [(f"{key}.", f"{value}.") for key, value in clauses.items()]
        for key, value in clauses.items():
            for short, long in aliases:
                if value.startswith(short):
                    clauses[key] = value.replace(short, long)

        # {alias: full_name_of_table_or_file}
        self.clauses = clauses
        return clauses


class S3FromParser(FromParser):
    def get_source_data(self, documents: Dict[str, str]):
        from_alias = list(self.clauses.keys())[0].lower()
        from_query = list(self.clauses.values())[0].lower()
        if "." in from_query:
            return self._get_nested_source_data(documents)

        key_has_asterix = from_query.endswith("[*]")
        from_query = from_query[0:-3] if key_has_asterix else from_query
        from_alias = from_alias[0:-3] if from_alias.endswith("[*]") else from_alias
        doc_is_list = documents[from_query].startswith("[") and documents[
            from_query
        ].endswith("]")

        source_data = JsonParser().parse(documents[from_query])
        if doc_is_list:
            return {"_1": source_data}
        elif from_alias:
            if isinstance(source_data, list):
                return [CaseInsensitiveDict({from_alias: doc}) for doc in source_data]
            else:
                return CaseInsensitiveDict({from_alias: source_data})
        else:
            return source_data

    def _get_nested_source_data(self, documents: Dict[str, str]):
        """
        Our FROM-clauses are nested, meaning we need to dig into the provided document to return the key that we need
           --> FROM s3object.name as name
        """
        root_doc = True
        source_data = documents
        iterate_over_docs = False
        entire_key = list(self.clauses.values())[0].lower().split(".")
        alias = list(self.clauses.keys())[0]
        if alias.endswith("[*]"):
            alias = alias[0:-3]
        key_so_far = []
        for key in entire_key:
            key_so_far.append(key)
            key_has_asterix = key.endswith("[*]") and key[0:-3] in source_data
            new_key = key[0:-3] if key_has_asterix else key
            if iterate_over_docs and isinstance(source_data, list):
                # The previous key ended in [*]
                # Iterate over all docs in the result, and only return the requested source key
                if key_so_far == entire_key:
                    # If we have an alias, we have to use that instead of the original name
                    source_data = [{alias: doc.get(new_key, {})} for doc in source_data]
                else:
                    source_data = [doc.get_original(new_key, {}) for doc in source_data]
            else:
                # The previous key was a regular key
                # Assume that the result consists of a singular JSON document
                if new_key in source_data:
                    doc_is_list = source_data[new_key].startswith("[") and source_data[
                        new_key
                    ].endswith("]")
                    source_data = JsonParser().parse(source_data[new_key])
                    if root_doc and doc_is_list:
                        # AWS behaviour when the root-document is a list
                        source_data = {"_1": source_data}
                    elif key_so_far == entire_key:
                        if isinstance(source_data, list):
                            source_data = [{alias: doc} for doc in source_data]
                        else:
                            source_data = {alias: source_data}
                else:
                    source_data = {}

            iterate_over_docs = key_has_asterix
            root_doc = False

        return source_data


class DynamoDBFromParser(FromParser):
    def get_source_data(self, documents: Dict[str, str]):
        source_data = documents[list(self.clauses.values())[0].lower()]
        return JsonParser().parse(source_data)
