import inspect
from datetime import datetime
from types import MappingProxyType
from typing import Annotated, Any, Dict, List, Optional, Set, Tuple, Type, get_args, get_origin

import csp
from ccflow import ArrowSchema, ArrowTable, NDArray
from csp import Struct
from csp.impl.struct import StructMeta
from csp.impl.types.container_type_normalizer import ContainerTypeNormalizer
from csp.impl.types.typing_utils import CspTypingUtils
from numpy import ndarray
from pyarrow import Array, Schema, Table
from pydantic import BaseModel, create_model, model_validator
from pydantic.fields import FieldInfo
from pydantic_core import CoreConfig, core_schema

from ..id_generator import get_counter
from .psp import PerspectiveUtilityMixin

IdType = str

__all__ = ("GatewayStruct", "IdType", "PydanticizedCspStruct")


class _DynamicModel(BaseModel):
    @classmethod
    def _validate(cls, v):
        # Without the statement below, pydantic would convert a list containing an invalid request to
        # a single request without any fields set, so we must prevent this accidental conversion
        if isinstance(v, list):
            raise ValueError("converting list inputs to a single value is disabled")
        if _is_pydanticized_struct(v):
            pydantic_type = v.__class__.__pydantic_model__
            if not issubclass(v.__class__.__pydantic_model__, cls):
                raise ValueError(f"Cannot convert struct {type(v)} (with pydantic model {pydantic_type}) to pydantic model {cls}.")
            starting_dict = {}

            for k in v.metadata():
                # Pydantic doesn't support fields that start with underscore
                if k.startswith("_"):
                    continue
                value = getattr(v, k, None)
                if k in pydantic_type.model_fields and value is not None:
                    # Normally pydantic would handle this part automatically. However, we need our own implementation to
                    # handle the case where the value is a subclass of the declared type
                    if _is_pydanticized_struct(value):
                        value = _to_pydantic(value)
                    if isinstance(value, list):
                        value = [_to_pydantic(vi) if _is_pydanticized_struct(vi) else vi for vi in value]
                    if isinstance(value, dict):  # pydantic should handle this automatically, but not done in v1
                        value = {ki: _to_pydantic(vi) if _is_pydanticized_struct(vi) else vi for ki, vi in value.items()}
                    starting_dict[k] = value
            return starting_dict
        else:
            return v

    @model_validator(mode="before")
    def _model_validate(cls, v, info):
        # Automatically convert "GatewayStruct" structs to their pydantic versions on pydantic validation
        return cls._validate(v)

    model_config = {"arbitrary_types_allowed": True, "extra": "forbid", "frozen": True, "coerce_numbers_to_str": True}


def _to_pydantic(self: Struct) -> _DynamicModel:
    """Construct a pydantic model from `self`, a `Struct`.
    Since `Struct.to_dict()` is recursive, we must be as well.
    """
    pydantic_type = self.__class__.__pydantic_model__
    return pydantic_type.model_validate(self)


def _to_csp_struct(
    self: _DynamicModel,
    exclude: Set[str] = None,
) -> Struct:
    """Construct a Struct from self, a pydantic model.
    Args:
        exclude: This only applies to GatewayStructs.
            Determines which fields to exclude from the returned Struct and acts recursively on all GatewayStruct members.
            The usecase in mind is to allow stripping id and timestamp before constructing
            the underlying GatewayStruct, to allow those fields to get autogenerated.
    Returns:
        Instantiated csp.Struct class
    Example:
        class MyClass(GatewayStruct):
            x: int
        MyClass(x=10).to_pydantic().csp(exclude=set(["id"]))
        # The returned object will have "id" autogenerated, replacing the original "id"
    """
    struct_type = self.__class__.__csp_struct__

    starting_dict = {}
    for k, v in dict(self).items():
        if k in struct_type.metadata() and v is not None:
            if isinstance(v, _DynamicModel):
                v = _to_csp_struct(v, exclude=exclude)
            if isinstance(v, list):
                v = [_to_csp_struct(vi, exclude=exclude) if isinstance(vi, _DynamicModel) else vi for vi in v]
            if isinstance(v, dict):
                v = {ki: _to_csp_struct(vi, exclude=exclude) if isinstance(vi, _DynamicModel) else vi for ki, vi in v.items()}
            starting_dict[k] = v
    if exclude is not None and issubclass(struct_type, GatewayStruct):
        for val in exclude:
            starting_dict.pop(val, None)
    return struct_type(**starting_dict)


def _pydanticize_struct_class(cls: Type[Struct]) -> None:
    """Update Struct subclass in place with pydantic version"""
    # Used by pydantic internally
    cls.__pydantic_model__ = CspStructToPydanticModel(cls)
    cls.__pydantic_model__.__csp_struct__ = cls

    # Converters
    # csp uses to_<type>()
    cls.to_pydantic = _to_pydantic
    cls.to_pydantic.__name__ = "to_pydantic"
    # pydantic uses <type>()
    cls.__pydantic_model__.csp = _to_csp_struct
    cls.__pydantic_model__.csp.__name__ = "csp"


def _is_pydanticized_struct(v: Struct) -> bool:
    return isinstance(v, Struct) and hasattr(v.__class__, "__pydantic_model__")


def _normalize_csp_struct_type_annotations_with_pydanticized_type(typ: Any) -> Any:
    """Recursively replace csp.Struct with their pydanticized versions."""
    obj_type = ContainerTypeNormalizer.normalize_type(typ)

    # First, handle Annotated types
    if get_origin(obj_type) is Annotated:
        # Extract the original type and metadata
        original_type, *metadata = get_args(obj_type)
        normalized_type = _normalize_csp_struct_type_annotations_with_pydanticized_type(original_type)
        # Handle Field defaults in metadata
        for meta in metadata:
            if isinstance(meta, FieldInfo) and isinstance(meta.default, Struct):
                if hasattr(meta.default, "to_pydantic"):
                    default = _to_pydantic(meta.default)
                else:
                    _normalize_csp_struct_type_annotations_with_pydanticized_type(meta.default.__class__)
                    default = _to_pydantic(meta.default)
                meta.default = default

        return Annotated[(normalized_type, *metadata)]

    # Handle generic containers
    if CspTypingUtils.is_generic_container(obj_type):
        orig_type = CspTypingUtils.get_origin(obj_type)
        if orig_type in (List, Set, Tuple):
            (expected_item_type,) = obj_type.__args__
            return orig_type[_normalize_csp_struct_type_annotations_with_pydanticized_type(expected_item_type)]
        elif orig_type is Dict:
            expected_key_type, expected_value_type = obj_type.__args__
            return Dict[expected_key_type, _normalize_csp_struct_type_annotations_with_pydanticized_type(expected_value_type)]

    # Handle Struct types
    elif issubclass(typ, Struct):
        if (pydantic_model := getattr(typ, "__pydantic_model__", None)) is not None:
            return pydantic_model
        _pydanticize_struct_class(typ)
        return typ.__pydantic_model__

    return typ


# TODO mainline into csp
def CspStructToPydanticModel(csp_struct: Struct) -> Type[BaseModel]:
    """:meta private:"""
    fields: Dict[str, Any] = {}

    # We use typed to get annotated fields
    for k, v in csp_struct.metadata(typed=True).items():
        if k.startswith("_"):
            continue

        typ = _normalize_csp_struct_type_annotations_with_pydanticized_type(v)
        fields[k] = (typ, None)

    # type munging for pydantic
    for k, v in fields.items():
        raw_field_type = v[0]
        try:
            if inspect.isclass(raw_field_type):
                field_type = raw_field_type
            else:
                field_type = CspTypingUtils.get_origin(raw_field_type)
            if issubclass(field_type, ndarray):
                # use pydantic-compatible generic
                # try to grab the generic
                generic = csp_struct.metadata(typed=True)[k]

                # now try to grab the args
                arg = get_args(generic)[0]

                # now swap with pydantic compatible class
                fields[k] = tuple([NDArray[arg], *v[1:]])
            elif issubclass(field_type, Array):
                # TODO
                raise NotImplementedError("Not implemented, no pyarrow.Array pydantic compat")  # pragma: no cover
            elif issubclass(field_type, Schema):
                # TODO
                raise NotImplementedError(f"Not implemented, have {ArrowSchema} but not implemented")  # pragma: no cover
            elif issubclass(field_type, Table):
                # TODO
                raise NotImplementedError(f"Not implemented, have {ArrowTable} but not implemented")  # pragma: no cover
        except (KeyError, IndexError, TypeError):
            # just ignore and hope for the best! errors will bubble up later
            ...

    for k, v in list(fields.items()):
        if isinstance(v, tuple):
            # TODO this might only be a py38 problem
            if isinstance(v[0], list) and len(v[0]) >= 1:
                typ = v[0][0].__pydantic_model__ if hasattr(v[0][0], "__pydantic_model__") else v[0][0]

                fields[k] = (List[typ], None)  # type: ignore[misc, index]
            if len(v) == 2:
                # We pull out the true default, from the Annotations
                typ, default = v
                if get_origin(typ) is Annotated:
                    _original_type, *metadata = get_args(typ)
                    for meta in metadata:
                        if isinstance(meta, FieldInfo) and isinstance(meta.default, _DynamicModel):
                            fields[k] = (typ, meta.default)
                            break

    for name, default in csp_struct.__defaults__.items():
        # Pydantic doesn't support fields that start with underscore
        if name.startswith("_"):
            continue
        if default is not None:
            # We have to convert csp types to their pydantic version. This exists since for the
            if isinstance(default, Struct):
                if hasattr(default, "to_pydantic"):
                    default = _to_pydantic(default)
                else:
                    _normalize_csp_struct_type_annotations_with_pydanticized_type(default.__class__)
                    default = _to_pydantic(default)
            fields[name] = (fields[name][0], default)
    parent_struct = csp_struct.__mro__[1]
    base = getattr(parent_struct, "__pydantic_model__", _DynamicModel)
    for k, v in fields.items():
        fields[k] = (Optional[v[0]], v[1])
    if hasattr(csp_struct, "__get_validator_dict__") and (validators := csp_struct.__get_validator_dict__()) is not None:
        model = create_model(csp_struct.__name__, __base__=base, __validators__=validators, **fields)
    else:
        model = create_model(csp_struct.__name__, __base__=base, **fields)

    return model


class PydanticizedCspStruct(StructMeta):
    """A subclass of StructMeta from csp, this class adds additional properties onto csp.Struct classes to link them into the Gateway format.

    Specifically, allows lookups and automatic id generation, in addition to
    creating a Pydantic clone of the given class to leverage the validation and serialization/deserialization benefits of Pydantic.
    """

    def __init__(cls: Any, name: str, bases: Any, attr_dict: Any) -> None:
        super().__init__(name, bases, attr_dict)
        # Automatically construct pydantic model from csp.struct

        # TODO: Remove when can remove pydanticized version of GatewayStructs
        _pydanticize_struct_class(cls)

        # Attach an id generator to every class
        cls.id_generator = get_counter(cls)

        # Allow for looking up by ID
        cls._internal_mapping: Dict[str, Any] = {}

        # But expose this lookup as a readonly mapping proxy
        cls.lookup = MappingProxyType(cls._internal_mapping).get
        cls._include_in_lookup = True

    def omit_from_lookup(cls, omit=True):
        cls._include_in_lookup = not omit

    def included_in_lookup(cls):
        return cls._include_in_lookup


class GatewayStruct(PerspectiveUtilityMixin, Struct, metaclass=PydanticizedCspStruct):
    """Sub-class of csp.Struct specifically designed for usage with csp-gateway.
    These classes inherit from csp.Struct, but each one also contains a pydantic model
    as an attribute that mirrors the underlying struct class.

    The pydantic model can be constructed by running `.to_pydantic()` on an instance of the
    given GatewayStruct. To access the underlying class, one can call `__pydantic_model__` on the specific GatewayStruct sub-class object.
    For example:
        class MyClass(GatewayStruct):
            x: int
        pydantic_model = MyClass.__pydantic_model__

    The pydantic model of the the GatewayStruct allows for data validation. Furthermore, the utilities
    in csp-gateway (such as Rest routes, Kafka, Perspective integrations) are specifically designed for GatewayStruct's.
    """

    id: IdType
    timestamp: datetime

    def __init__(self, **kwargs: Any) -> None:
        # auto generate id on every new construction
        if "id" not in kwargs:
            # TODO consider postfixing with a _ for conflicts
            kwargs["id"] = str(self.__class__.id_generator.next())
        if "timestamp" not in kwargs:
            kwargs["timestamp"] = datetime.utcnow()

        # And put into lookup
        # TODO make this nicer when we switch to pydantic as first class instead of csp struct
        if self.__class__.included_in_lookup():
            self.__class__._internal_mapping[kwargs["id"]] = self

        # and defer to normal csp.struct construction
        super().__init__(**kwargs)

    @classmethod
    def _validate_gateway_struct(cls, val):
        """Validate GatewayStruct after pydantic type validation.
        A validator attached to every GatewayStruct to allow for defining custom
        model-level after validators that run after pydantic type validation.
        If not defined on a child class, the parent's validator will be used.  If defined on a child class, the parent's validator will be ignored. Please call the parent's validator directly if you want to run both.
        Args:
            cls: The class this validator is attached to
            val: The value to validate
        Returns:
            The validated value, possibly modified
        """
        return val

    @staticmethod
    def _get_pydantic_core_schema(cls, source_type, handler):
        # Get parent schema - note the cls parameter
        parent_schema = csp.Struct._get_pydantic_core_schema(cls, source_type, handler)
        core_config = CoreConfig(coerce_numbers_to_str=True)
        # soooo hacky...
        parent_schema["schema"]["config"] = core_config
        return core_schema.no_info_after_validator_function(
            function=cls._validate_gateway_struct, schema=parent_schema, serialization=parent_schema.get("serialization")
        )

    @classmethod
    def generate_id(cls) -> str:
        return str(cls.id_generator.next())
