from collections import defaultdict
from enum import Enum
import re
from typing import Mapping, Type, Callable, Any, List, Optional, Dict

from flask import Blueprint, request, jsonify
from gena.deserializer import generate_deserializer
from peewee import Model as PeeweeModel, DoesNotExist, fn
from werkzeug.exceptions import BadRequest, NotFound

from gena.serializer import Serializer, get_peewee_serializer


class APIFuncs(str, Enum):
    get = "get"
    get_by_ids = "get_by_ids"
    get_one = "get_one"
    has = "has"
    create = "create"
    update = "update"
    delete_by_id = "delete_by_id"
    truncate = "truncate"


def generate_api(
    Model: Type[PeeweeModel],
    deserializers: Optional[Dict[str, Callable[[Any], Any]]] = None,
    serialize: Optional[Callable[[Any], dict]] = None,
    batch_serialize: Optional[Callable[[List[Any]], List[dict]]] = None,
    known_type_serializer: Optional[Dict[Any, Serializer]] = None,
    enable_truncate_table: bool = False,
    skip_funcs: Optional[set[APIFuncs]] = None,
):
    """Generate API from the given Model

    Args:
        Model: peewee model of the table
        deserializers: deserialize raw value, throw value error if value is invalid. you can provide deserializer for some field and the rest is going to be generated automatically
        serialize: function to serialize a record of the table
        batch_serialize: function to serialize multiple records of the table
        known_type_serializer: serializer for known type, only affected if serialize is generated automatically
        enable_truncate_table: whether to enable API to truncate the whole table. Note: deprecated, use skip_funcs instead
        skip_funcs: list of APIFuncs to skip, default is [truncate]
    """
    if skip_funcs is None:
        skip_funcs = set()
    if not enable_truncate_table and APIFuncs.truncate not in skip_funcs:
        skip_funcs.add(APIFuncs.truncate)

    table_name = Model._meta.table_name
    default_limit = str(50)
    name2field = {name: field for name, field in Model._meta.fields.items()}
    op_fields = {"fields", "limit", "offset", "unique", "sorted_by", "group_by"}
    field_reg = re.compile(r"(?P<name>[a-zA-Z_0-9]+)(?:\[(?P<op>[a-zA-Z0-9]+)\])?")

    if (
        deserializers is None
        or len(set(name2field.keys()).difference(deserializers.keys())) > 0
    ):
        if deserializers is None:
            deserializers = {}

        # automatically filling missing deserializers
        deserializers.update(
            **generate_deserializer(
                Model, known_field_deserializers=set(deserializers.keys())
            )
        )

    if len(set(name2field.keys()).difference(deserializers.keys())) > 0:
        raise Exception(
            f"Table {table_name} doesn't have deserializer for field: {set(name2field.keys()).difference(deserializers.keys())}"
        )

    if serialize is None:
        if hasattr(Model, "to_dict"):
            serialize = Model.to_dict
        elif batch_serialize is not None:
            serialize = lambda x: batch_serialize([x])[0]
        else:
            serialize = get_peewee_serializer(
                Model, known_type_serializer=known_type_serializer
            )

    if batch_serialize is None:
        assert serialize is not None
        batch_serialize = lambda lst: [serialize(item) for item in lst]

    assert len(op_fields.intersection(name2field.keys())) == 0

    bp = Blueprint(table_name, table_name)

    def get():
        """Retrieving records matched a query.

        There are fixed keyword arguments:
            - limit: limit the number of records returned
            - offset: offset of the records returned
            - sorted_by: comma separated list of field name to sort by, prefix a field with "-" to sort that field in descending order.
            - group_by: comma separated list of field name to group by.
            - fields: comma separated list of field name to return. If not specified, all fields are returned.

        To filter records, you can apply the condition of a column by <field>=<value> (equal condition).
        Condition on a field such as >, >=, <, <=, `max`, `min`, `in` can be specified using brackets such as: <field>[gt]=10.
        We also support complex conditions:
            1. `max`: keep the record in a group that has the largest value
            2. `min`: keep the record in a group that has the smallest value
            3. `in`: select record that its values are in the given list

        The syntax for `max` and `min` are <field>[op]=<fields>, where op is either `max` or `min` and fields is a comma separated
        list of fields that will be used to group the records.

        We can support another aggregation to select keep all values in a group. However, since each value is for each
        record, we also have multiple ids. Therefore, a natural choice is to use `group_by` operator instead. Note that
        when you use group_by, the output is still a table (not a mapping) so that the client can reuse the code that
        read the data, but they have to group the result themselves.

        Note that we enforce the constraint that only one aggregation (`max`, `min`, `group_by`, etc) is allow in a query
        to ensure the behaviour of the query is deterministic (e.g., apply a group_by and max, which one is apply first?).
        """
        if "fields" in request.args:
            field_names = request.args["fields"].split(",")
            fields = [name2field[name] for name in field_names]
        else:
            field_names = []
            fields = []

        limit = int(request.args.get("limit", default_limit))
        offset = int(request.args.get("offset", "0"))

        # construct select clause
        query = Model.select(*fields)
        unique = request.args.get("unique", "false") == "true"
        if unique:
            query = query.distinct()

        # construct order by clause
        order_by = []
        if "sorted_by" in request.args:
            for field in request.args["sorted_by"].split(","):
                try:
                    if field.startswith("-"):
                        field = field[1:]
                        order_by.append(name2field[field].desc())
                    else:
                        order_by.append(name2field[field])
                except KeyError:
                    raise BadRequest(f"Invalid field name: {field}")

        if len(order_by) > 0:
            query = query.order_by(*order_by)

        # construct group by clause
        group_by = []
        if "group_by" in request.args:
            for field in request.args["group_by"].split(","):
                if field not in name2field:
                    raise BadRequest(f"Invalid field name: {field}")
                group_by.append(name2field[field])

        # construct where clause
        filter_fields = defaultdict(list)
        for name, value in request.args.items():
            if name in op_fields:
                continue
            m = field_reg.match(name)
            if m is None:
                raise BadRequest(f"Invalid field name: {name}")

            name = m.group("name")
            if name not in name2field:
                raise BadRequest(f"Invalid field name: {name}")

            op = m.group("op")
            filter_fields[name].append((op, value))

        pending_ops = defaultdict(dict)
        conditions = defaultdict(list)
        for name, ops in filter_fields.items():
            for op, value in ops:
                field = name2field[name]
                if op in {"max"}:
                    assert op not in pending_ops[name]
                    pending_ops[name][op] = value
                    continue
                elif op == "in":
                    deser = deserializers[name]
                    conditions[field].append(
                        (field.in_([deser(x) for x in value.split(",")]))
                    )
                    continue

                # no special operator
                value = deserializers[name](value)
                if op is None:
                    conditions[field].append((field == value))
                elif op == "gt":
                    conditions[field].append((field > value))
                elif op == "gte":
                    conditions[field].append((field >= value))
                elif op == "lt":
                    conditions[field].append((field < value))
                elif op == "lte":
                    conditions[field].append((field <= value))
                else:
                    raise BadRequest(f"Does not support {op} yet")

        if len(conditions) > 0:
            query = query.where(*[item for lst in conditions.values() for item in lst])

        if len(group_by) > 0:
            if len(pending_ops) > 0:
                raise BadRequest(f"Does not support multiple aggregations")
            # update the select to keep the id
            subquery = query.select(
                *[c.alias(f"gb_c{i}") for i, c in enumerate(group_by)]
            ).group_by(*group_by)

            predicate = group_by[0] == getattr(subquery.c, "gb_c0")
            for i, c in enumerate(group_by[1:], start=1):
                predicate = predicate & (c == getattr(subquery.c, f"gb_c{i}"))
            query = query.join(subquery.limit(limit).offset(offset), on=predicate)
            # they want to get only one record so we save computation knowing that it won't use anyway
            total = subquery.count()
        else:
            for name, ops in pending_ops.items():
                field = name2field[name]
                for op, value in ops.items():
                    if op == "max":
                        # select the record with maximum value in the group
                        # need to do a subquery to select the one with maximum value
                        subquery_group_fields = []
                        subquery_group_field_conditions = []

                        for gfield in value.split(","):
                            if gfield not in name2field:
                                raise BadRequest(f"Invalid group by field: {gfield}")
                            subquery_group_fields.append(name2field[gfield])
                            if gfield in conditions:
                                if gfield in pending_ops:
                                    raise BadRequest(
                                        f"Does not support multiple aggregations"
                                    )
                                subquery_group_field_conditions += conditions[gfield]

                        subquery_name = f"{name}_{op}"
                        field_alias = f"{subquery_name}_{name}"
                        subquery = (
                            Model.select(Model.id, fn.MAX(field).alias(field_alias))
                            .group_by(*subquery_group_fields)
                            .alias(subquery_name)
                        )

                        if len(subquery_group_field_conditions) > 0:
                            subquery = subquery.where(*subquery_group_field_conditions)

                        predicate = (Model.id == subquery.c.id) & (
                            field == getattr(subquery.c, field_alias)
                        )
                        query = query.join(subquery, on=predicate)

            # they want to get only one record so we save computation knowing that it won't use anyway
            total = query.count()
            query = query.limit(limit).offset(offset)

        # perform the query
        items = batch_serialize(list(query))
        if len(fields) > 0:
            items = [{k: item[k] for k in field_names} for item in items]

        return jsonify({"items": items, "total": total})

    if APIFuncs.get not in skip_funcs:
        bp.route(f"/{table_name}", methods=["GET"])(get)

    def get_by_ids():
        if "ids" not in request.json:
            raise BadRequest("Bad request. Missing `ids`")

        ids = request.json["ids"]
        if "fields" in request.args:
            field_names = request.args["fields"].split(",")
            fields = [name2field[name] for name in field_names]
        else:
            field_names = []
            fields = []

        records = list(Model.select(Model.id, *fields).where(Model.id.in_(ids)))
        records = batch_serialize(records)

        if len(field_names) > 0:
            records = {
                item["id"]: {k: item[k] for k in field_names if k in item}
                for item in records
            }

        return jsonify({"items": records, "total": len(records)})

    if APIFuncs.get_by_ids not in skip_funcs:
        bp.route(f"/{table_name}/find_by_ids", methods=["POST"])(get_by_ids)

    def get_one(id):
        if "fields" in request.args:
            field_names = request.args["fields"].split(",")
            fields = [name2field[name] for name in field_names]
        else:
            field_names = []
            fields = []

        records = list(Model.select(*fields).where(Model.id == id))
        if len(records) == 0:
            raise NotFound(f"Record {id} does not exist")

        record = serialize(records[0])
        if len(fields) > 0:
            record = {k: record[k] for k in field_names}

        return jsonify(record)

    if APIFuncs.get_one not in skip_funcs:
        bp.route(f"/{table_name}/<id>", methods=["GET"])(get_one)

    def has(id):
        if not Model.select().where(Model.id == id).exists():
            raise NotFound(f"Record {id} does not exist")
        return jsonify()

    if APIFuncs.has not in skip_funcs:
        bp.route(f"/{table_name}/<id>", methods=["HEAD"])(has)

    def create():
        posted_record = request.json
        raw_record = {}

        for name, field in name2field.items():
            if name in posted_record:
                try:
                    raw_record[name] = deserializers[name](posted_record[name])
                except ValueError as e:
                    raise ValueError(f"Field `{name}` {str(e)}")
        if "id" in raw_record:
            # remove id as this API always creates a new record
            raw_record.pop("id")
        record = Model.create(**raw_record)
        # TODO: correct return types according to RESTful specification https://restfulapi.net/http-methods/
        return jsonify(serialize(record))

    if APIFuncs.create not in skip_funcs:
        bp.route(f"/{table_name}", methods=["POST"])(create)

    def update(id):
        try:
            record = Model.get_by_id(id)
        except DoesNotExist as e:
            raise NotFound(f"Record {id} does not exist")

        request_json = request.get_json()
        if request_json is None:
            raise BadRequest("Missing request body")

        for name, field in name2field.items():
            if name in request_json:
                try:
                    value = deserializers[name](request_json[name])
                except ValueError as e:
                    raise ValueError(f"Field `{name}` {str(e)}")

                setattr(record, name, value)
        record.save()

        return jsonify(serialize(record))

    if APIFuncs.update not in skip_funcs:
        bp.route(f"/{table_name}/<id>", methods=["PUT"])(update)

    def delete_by_id(id):
        try:
            Model.get_by_id(id).delete_instance()
        except DoesNotExist as e:
            raise NotFound(f"Record {id} does not exist")

        return jsonify({"status": "success"})

    if APIFuncs.delete_by_id not in skip_funcs:
        bp.route(f"/{table_name}/<id>", methods=["DELETE"])(delete_by_id)

    def truncate():
        Model.truncate_table()
        return jsonify({"status": "success"})

    if APIFuncs.truncate not in skip_funcs:
        bp.route(f"/{table_name}", methods=["DELETE"])(truncate)

    return bp


def generate_readonly_api_4dict(
    name: str,
    id2ent: Mapping[str, Any],
    unique_field_funcs: Dict[str, Callable[[str], str]] = None,
    serialize: Optional[Callable[[Any], dict]] = None,
    batch_serialize: Optional[Callable[[List[Any]], List[dict]]] = None,
):
    """Generate API for a dictionary.

    Make it as similar to the API of DB as possible. However, due to its limitation, we
    only support accessing by unique key (via id or other unique field)

    Args:
        name: name of the endpoint
        id2ent: dictionary of entities
        unique_field_funcs: unique field and function to transform it to id
        serialize: function to serialize an entity
        batch_serialize: function to serialize a list of entities
    """
    op_fields = {"fields", "limit", "offset", "unique", "sorted_by"}
    if batch_serialize is None:
        assert serialize is not None
        batch_serialize = gen_batch_serialize(serialize)
    elif serialize is None:
        serialize = lambda x: batch_serialize([x])[0]

    if unique_field_funcs is None:
        unique_field_funcs = {}
    bp = Blueprint(name, name)

    @bp.route(f"/{name}", methods=["GET"])
    def get():
        """Retrieving records matched a query."""
        if "fields" in request.args:
            field_names = request.args["fields"].split(",")
        else:
            field_names = []

        lst = []
        for name in request.args.keys():
            if name in op_fields:
                continue
            if name not in unique_field_funcs:
                raise BadRequest(f"Invalid field name: {name}")
            lst.append(name)

        if len(lst) > 1:
            raise BadRequest(f"Invalid query. Only one field is allowed but get: {lst}")

        if len(lst) == 0:
            raise BadRequest(f"Invalid query. Must provide at least one field")

        id = unique_field_funcs[lst[0]](request.args[lst[0]])
        record = serialize(id2ent[id])
        if len(field_names) > 0:
            record = {k: record[k] for k in field_names if k in record}

        return jsonify({"items": [record], "total": 1})

    @bp.route(f"/{name}/find_by_ids", methods=["POST"])
    def find_by_ids():
        if "ids" not in request.json:  # type: ignore
            raise BadRequest("Bad request. Missing `ids`")

        if "fields" in request.args:
            field_names = request.args["fields"].split(",")
        else:
            field_names = []

        ids = []
        ents = []
        for id in request.json["ids"]:
            if id in id2ent:
                ents.append(id2ent[id])
                ids.append(id)

        records = batch_serialize(ents)
        if len(field_names) > 0:
            records = [
                {k: item[k] for k in field_names if k in item} for item in records
            ]

        return jsonify({"items": dict(zip(ids, records)), "total": len(ents)})

    @bp.route(f"/{name}/<id>", methods=["GET"])
    def find_by_id(id: str):
        if id not in id2ent:
            raise NotFound(f"Record {id} does not exist")

        record = serialize(id2ent[id])

        if "fields" in request.args:
            field_names = request.args["fields"].split(",")
        else:
            field_names = []

        if len(field_names) > 0:
            record = {k: record[k] for k in field_names if k in record}

        return jsonify(record)

    return bp


def gen_batch_serialize(
    serialize: Callable[[Any], dict]
) -> Callable[[List[Any]], List[dict]]:
    def batch_serialize(lst):
        return [serialize(item) for item in lst]

    return batch_serialize
