from __future__ import annotations

# Copyright 2014 Cloudera Inc.
#
# 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
#
# http://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.
from io import StringIO


def parse_metadata(descr_table):
    parser = MetadataParser(descr_table)
    return parser.parse()


def _noop(tup):
    return None


def _item_converter(i):
    def _get_item(converter=None):
        def _converter(tup):
            result = tup[i]
            if converter is not None:
                result = converter(result)
            return result

        return _converter

    return _get_item


_get_type = _item_converter(1)
_get_comment = _item_converter(2)


def _try_timestamp(x):
    import pandas as pd

    try:
        ts = pd.Timestamp(x, tz="UTC")
        return ts.to_pydatetime().replace(tzinfo=None)
    except (ValueError, TypeError):
        return x


def _try_unix_timestamp(x):
    try:
        value = int(x)
    except (ValueError, TypeError):
        return x
    else:
        import pandas as pd

        return (
            pd.Timestamp.fromtimestamp(value, tz="UTC")
            .tz_localize(None)
            .to_pydatetime()
        )


def _try_boolean(x):
    try:
        x = x.lower()
        if x in ("true", "yes"):
            return True
        elif x in ("false", "no"):
            return False
        return x
    except (ValueError, TypeError):
        return x


def _try_int(x):
    try:
        return int(x)
    except (ValueError, TypeError):
        return x


class MetadataParser:
    """A simple state machine to parse the results of `DESCRIBE FORMATTED`."""

    def __init__(self, table):
        self.table = table
        self.tuples = list(self.table.itertuples(index=False))

    def _reset(self):
        self.pos = 0
        self.schema = None
        self.partitions = None
        self.info = None
        self.storage = None

    def _next_tuple(self):
        if self.pos == len(self.tuples):
            raise StopIteration

        result = self.tuples[self.pos]
        self.pos += 1
        return result

    def parse(self):
        self._reset()
        self._parse()

        return TableMetadata(
            self.schema, self.info, self.storage, partitions=self.partitions
        )

    def _parse(self):
        self.schema = self._parse_schema()

        next_section = self._next_tuple()
        if "partition" in next_section[0].lower():
            self._parse_partitions()
        else:
            self._parse_info()

    def _parse_partitions(self):
        self.partitions = self._parse_schema()

        next_section = self._next_tuple()
        if "table information" not in next_section[0].lower():
            raise ValueError("Table information not present")

        self._parse_info()

    def _parse_schema(self):
        tup = self._next_tuple()
        if "col_name" not in tup[0]:
            raise ValueError(
                f"DESCRIBE FORMATTED did not return the expected results: {tup}"
            )
        self._next_tuple()

        # Use for both main schema and partition schema (if any)
        schema = []
        while True:
            tup = self._next_tuple()
            if not tup[0].strip():
                break
            schema.append((tup[0], tup[1]))

        return schema

    def _parse_info(self):
        self.info = {}
        while True:
            tup = self._next_tuple()
            orig_key = tup[0].strip(":")
            key = _clean_param_name(tup[0])

            if not key or key.startswith("#"):
                # section is done
                break

            if key == "table parameters":
                self._parse_table_parameters()
            elif key in self._info_cleaners:
                result = self._info_cleaners[key](tup)
                self.info[orig_key] = result
            else:
                self.info[orig_key] = tup[1]

        if "storage information" not in key:
            raise ValueError("Storage information not present")

        self._parse_storage_info()

    _info_cleaners = {
        "database": _get_type(),
        "owner": _get_type(),
        "createtime": _get_type(_try_timestamp),
        "lastaccesstime": _get_type(_try_timestamp),
        "protect mode": _get_type(),
        "retention": _get_type(_try_int),
        "location": _get_type(),
        "table type": _get_type(),
    }

    def _parse_table_parameters(self):
        params = self._parse_nested_params(self._table_param_cleaners)
        self.info["Table Parameters"] = params

    _table_param_cleaners = {
        "external": _try_boolean,
        "column_stats_accurate": _try_boolean,
        "numfiles": _try_int,
        "totalsize": _try_int,
        "stats_generated_via_stats_task": _try_boolean,
        "numrows": _try_int,
        "transient_lastddltime": _try_unix_timestamp,
    }

    def _parse_storage_info(self):
        self.storage = {}
        while True:
            # end of the road
            try:
                tup = self._next_tuple()
            except StopIteration:
                break

            orig_key = tup[0].strip(":")
            key = _clean_param_name(tup[0])

            if not key or key.startswith("#"):
                # section is done
                break

            if key == "storage desc params":
                self._parse_storage_desc_params()
            elif key in self._storage_cleaners:
                result = self._storage_cleaners[key](tup)
                self.storage[orig_key] = result
            else:
                self.storage[orig_key] = tup[1]

    _storage_cleaners = {
        "compressed": _get_type(_try_boolean),
        "num buckets": _get_type(_try_int),
    }

    def _parse_storage_desc_params(self):
        params = self._parse_nested_params(self._storage_param_cleaners)
        self.storage["Desc Params"] = params

    _storage_param_cleaners = {}

    def _parse_nested_params(self, cleaners):
        import pandas as pd

        params = {}
        while True:
            try:
                tup = self._next_tuple()
            except StopIteration:
                break
            if pd.isnull(tup[1]):
                break

            key, value = tup[1:]

            if key.lower() in cleaners:
                cleaner = cleaners[key.lower()]
                value = cleaner(value)
            params[key] = value

        return params


def _clean_param_name(x):
    return x.strip().strip(":").lower()


def _get_meta(attr, key):
    @property
    def f(self):
        data = getattr(self, attr)
        if isinstance(key, list):
            result = data
            for k in key:
                if k not in result:
                    raise KeyError(k)
                result = result[k]
            return result
        else:
            return data[key]

    return f


class TableMetadata:
    """Container for the parsed and wrangled results of `DESCRIBE FORMATTED`."""

    def __init__(self, schema, info, storage, partitions=None):
        self.schema = schema
        self.info = info
        self.storage = storage
        self.partitions = partitions

    def __repr__(self):
        import pprint

        # Quick and dirty for now
        buf = StringIO()
        buf.write(str(type(self)))
        buf.write("\n")

        data = {
            "schema": self.schema,
            "info": self.info,
            "storage info": self.storage,
        }
        if self.partitions is not None:
            data["partition schema"] = self.partitions

        pprint.pprint(data, stream=buf)  # noqa: T203

        return buf.getvalue()

    @property
    def is_partitioned(self):
        return self.partitions is not None

    create_time = _get_meta("info", "CreateTime")
    location = _get_meta("info", "Location")
    owner = _get_meta("info", "Owner")
    num_rows = _get_meta("info", ["Table Parameters", "numRows"])
    hive_format = _get_meta("storage", "InputFormat")

    tbl_properties = _get_meta("info", "Table Parameters")
    serde_properties = _get_meta("storage", "Desc Params")


class TableInfo:
    pass


class TableStorageInfo:
    pass
