# AUTOGENERATED! DO NOT EDIT! File to edit: ../../../nbs/018_Avro_Encode_Decoder.ipynb.

# %% auto 0
__all__ = ['logger', 'AvroBase', 'avro_encoder', 'avro_decoder', 'avsc_to_pydantic']

# %% ../../../nbs/018_Avro_Encode_Decoder.ipynb 1
import io
import json
from typing import *

import fastavro
from pydantic import BaseModel, create_model

from ..logger import get_logger
from ..meta import export

# %% ../../../nbs/018_Avro_Encode_Decoder.ipynb 4
logger = get_logger(__name__)

# %% ../../../nbs/018_Avro_Encode_Decoder.ipynb 7
@export("fastkafka.encoder")
class AvroBase(BaseModel):
    """This is base pydantic class that will add some methods"""

    @classmethod
    def avro_schema_for_pydantic_object(
        cls,
        pydantic_model: BaseModel,
        by_alias: bool = True,
        namespace: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Returns the Avro schema for the given Pydantic object.

        Args:
            pydantic_model (BaseModel): The Pydantic object.
            by_alias (bool, optional): Generate schemas using aliases defined. Defaults to True.
            namespace (Optional[str], optional): Optional namespace string for schema generation.

        Returns:
            Dict[str, Any]: The Avro schema for the model.
        """

        schema = pydantic_model.__class__.model_json_schema(by_alias=by_alias)

        if namespace is None:
            # default namespace will be based on title
            namespace = schema["title"]

        return cls._avro_schema(schema, namespace)

    @classmethod
    def avro_schema_for_pydantic_class(
        cls,
        pydantic_model: Type[BaseModel],
        by_alias: bool = True,
        namespace: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Returns the Avro schema for the given Pydantic class.

        Args:
            pydantic_model (Type[BaseModel]): The Pydantic class.
            by_alias (bool, optional): Generate schemas using aliases defined. Defaults to True.
            namespace (Optional[str], optional): Optional namespace string for schema generation.

        Returns:
            Dict[str, Any]: The Avro schema for the model.
        """

        schema = pydantic_model.model_json_schema(by_alias=by_alias)

        if namespace is None:
            # default namespace will be based on title
            namespace = schema["title"]

        return cls._avro_schema(schema, namespace)

    @classmethod
    def avro_schema(
        cls, by_alias: bool = True, namespace: Optional[str] = None
    ) -> Dict[str, Any]:
        """
        Returns the Avro schema for the Pydantic class.

        Args:
            by_alias (bool, optional): Generate schemas using aliases defined. Defaults to True.
            namespace (Optional[str], optional): Optional namespace string for schema generation.

        Returns:
            Dict[str, Any]: The Avro schema for the model.
        """
        schema = cls.schema(by_alias=by_alias)

        if namespace is None:
            # default namespace will be based on title
            namespace = schema["title"]

        return cls._avro_schema(schema, namespace)

    @staticmethod
    def _avro_schema(schema: Dict[str, Any], namespace: str) -> Dict[str, Any]:
        """Return the avro schema for the given pydantic schema"""

        classes_seen = set()

        def get_definition(ref: str, schema: Dict[str, Any]) -> Dict[str, Any]:
            """Reading definition of base schema for nested structs"""
            id = ref.replace("#/definitions/", "")
            d = schema.get("definitions", {}).get(id)
            if d is None:
                raise RuntimeError(f"Definition {id} does not exist")
            return d  # type: ignore

        def get_type(value: Dict[str, Any]) -> Dict[str, Any]:
            """Returns a type of a single field"""
            t = value.get("type")
            f = value.get("format")
            r = value.get("$ref")
            a = value.get("additionalProperties")
            avro_type_dict: Dict[str, Any] = {}
            if "default" in value:
                avro_type_dict["default"] = value.get("default")
            if "description" in value:
                avro_type_dict["doc"] = value.get("description")
            if "allOf" in value and len(value["allOf"]) == 1:
                r = value["allOf"][0]["$ref"]
            if r is not None:
                class_name = r.replace("#/definitions/", "")
                if class_name in classes_seen:
                    avro_type_dict["type"] = class_name
                else:
                    d = get_definition(r, schema)
                    if "enum" in d:
                        avro_type_dict["type"] = {
                            "type": "enum",
                            "symbols": [str(v) for v in d["enum"]],
                            "name": d["title"],
                        }
                    else:
                        avro_type_dict["type"] = {
                            "type": "record",
                            "fields": get_fields(d),
                            # Name of the struct should be unique true the complete schema
                            # Because of this the path in the schema is tracked and used as name for a nested struct/array
                            "name": class_name,
                        }
                    classes_seen.add(class_name)
            elif t == "array":
                items = value.get("items")
                tn = get_type(items)  # type: ignore
                # If items in array are a object:
                if "$ref" in items:  # type: ignore
                    tn = tn["type"]
                # If items in array are a logicalType
                if (
                    isinstance(tn, dict)
                    and isinstance(tn.get("type", {}), dict)
                    and tn.get("type", {}).get("logicalType") is not None
                ):
                    tn = tn["type"]
                avro_type_dict["type"] = {"type": "array", "items": tn}
            elif t == "string" and f == "date-time":
                avro_type_dict["type"] = {
                    "type": "long",
                    "logicalType": "timestamp-micros",
                }
            elif t == "string" and f == "date":
                avro_type_dict["type"] = {
                    "type": "int",
                    "logicalType": "date",
                }
            elif t == "string" and f == "time":
                avro_type_dict["type"] = {
                    "type": "long",
                    "logicalType": "time-micros",
                }
            elif t == "string" and f == "uuid":
                avro_type_dict["type"] = {
                    "type": "string",
                    "logicalType": "uuid",
                }
            elif t == "string":
                avro_type_dict["type"] = "string"
            elif t == "number":
                avro_type_dict["type"] = "double"
            elif t == "integer":
                # integer in python can be a long
                avro_type_dict["type"] = "long"
            elif t == "boolean":
                avro_type_dict["type"] = "boolean"
            elif t == "object":
                if a is None:
                    value_type = "string"
                else:
                    value_type = get_type(a)  # type: ignore
                if isinstance(value_type, dict) and len(value_type) == 1:
                    value_type = value_type.get("type")  # type: ignore
                avro_type_dict["type"] = {"type": "map", "values": value_type}
            else:
                raise NotImplementedError(
                    f"Type '{t}' not support yet, "
                    f"please report this at https://github.com/godatadriven/pydantic-avro/issues"
                )
            return avro_type_dict

        def get_fields(s: Dict[str, Any]) -> List[Dict[str, Any]]:
            """Return a list of fields of a struct"""
            fields = []

            required = s.get("required", [])
            for key, value in s.get("properties", {}).items():
                if "type" not in value and "anyOf" in value:
                    any_of_types = value.pop("anyOf")
                    types = [x["type"] for x in any_of_types if x["type"] != "null"]
                    value["type"] = types[0]
                avro_type_dict = get_type(value)
                avro_type_dict["name"] = key

                if key not in required:
                    if avro_type_dict.get("default") is None:
                        avro_type_dict["type"] = ["null", avro_type_dict["type"]]
                        avro_type_dict["default"] = None

                fields.append(avro_type_dict)
            return fields

        fields = get_fields(schema)

        return {
            "type": "record",
            "namespace": namespace,
            "name": schema["title"],
            "fields": fields,
        }

# %% ../../../nbs/018_Avro_Encode_Decoder.ipynb 11
@export("fastkafka.encoder")
def avro_encoder(msg: BaseModel) -> bytes:
    """
    Encoder to encode pydantic instances to avro message

    Args:
        msg: An instance of pydantic basemodel

    Returns:
        A bytes message which is encoded from pydantic basemodel
    """
    schema = fastavro.schema.parse_schema(AvroBase.avro_schema_for_pydantic_object(msg))
    bytes_writer = io.BytesIO()

    d = msg.model_dump()
    for k, v in d.items():
        if "pydantic_core" in str(type(v)):
            d[k] = str(v)

    fastavro.schemaless_writer(bytes_writer, schema, d)
    raw_bytes = bytes_writer.getvalue()
    return raw_bytes

# %% ../../../nbs/018_Avro_Encode_Decoder.ipynb 13
@export("fastkafka.encoder")
def avro_decoder(raw_msg: bytes, cls: Type[BaseModel]) -> Any:
    """
    Decoder to decode avro encoded messages to pydantic model instance

    Args:
        raw_msg: Avro encoded bytes message received from Kafka topic
        cls: Pydantic class; This pydantic class will be used to construct instance of same class

    Returns:
        An instance of given pydantic class
    """
    schema = fastavro.schema.parse_schema(AvroBase.avro_schema_for_pydantic_class(cls))

    bytes_reader = io.BytesIO(raw_msg)
    msg_dict = fastavro.schemaless_reader(bytes_reader, schema)

    return cls(**msg_dict)

# %% ../../../nbs/018_Avro_Encode_Decoder.ipynb 16
@export("fastkafka.encoder")
def avsc_to_pydantic(schema: Dict[str, Any]) -> Type[BaseModel]:
    """
    Generate pydantic model from given Avro Schema

    Args:
        schema: Avro schema in dictionary format

    Returns:
        Pydantic model class built from given avro schema
    """
    if "type" not in schema or schema["type"] != "record":
        raise AttributeError("Type not supported")
    if "name" not in schema:
        raise AttributeError("Name is required")
    if "fields" not in schema:
        raise AttributeError("fields are required")

    classes = {}

    def get_python_type(t: Union[str, Dict[str, Any]]) -> str:
        """Returns python type for given avro type"""
        optional = False
        if isinstance(t, str):
            if t == "string":
                py_type = "str"
            elif t == "long" or t == "int":
                py_type = "int"
            elif t == "boolean":
                py_type = "bool"
            elif t == "double" or t == "float":
                py_type = "float"
            elif t in classes:
                py_type = t
            else:
                raise NotImplementedError(f"Type {t} not supported yet")
        elif isinstance(t, list):
            if "null" in t:
                optional = True
            if len(t) > 2 or (not optional and len(t) > 1):
                raise NotImplementedError("Only a single type ia supported yet")
            c = t.copy()
            c.remove("null")
            py_type = get_python_type(c[0])
        elif t.get("logicalType") == "uuid":
            py_type = "UUID"
        elif t.get("logicalType") == "decimal":
            py_type = "Decimal"
        elif (
            t.get("logicalType") == "timestamp-millis"
            or t.get("logicalType") == "timestamp-micros"
        ):
            py_type = "datetime"
        elif (
            t.get("logicalType") == "time-millis"
            or t.get("logicalType") == "time-micros"
        ):
            py_type = "time"
        elif t.get("logicalType") == "date":
            py_type = "date"
        elif t.get("type") == "enum":
            enum_name = t.get("name")
            if enum_name not in classes:
                enum_class = f"class {enum_name}(str, Enum):\n"
                for s in t.get("symbols"):  # type: ignore
                    enum_class += f'    {s} = "{s}"\n'
                classes[enum_name] = enum_class
            py_type = enum_name  # type: ignore
        elif t.get("type") == "string":
            py_type = "str"
        elif t.get("type") == "array":
            sub_type = get_python_type(t.get("items"))  # type: ignore
            py_type = f"List[{sub_type}]"
        elif t.get("type") == "record":
            record_type_to_pydantic(t)
            py_type = t.get("name")  # type: ignore
        elif t.get("type") == "map":
            value_type = get_python_type(t.get("values"))  # type: ignore
            py_type = f"Dict[str, {value_type}]"
        else:
            raise NotImplementedError(
                f"Type {t} not supported yet, "
                f"please report this at https://github.com/godatadriven/pydantic-avro/issues"
            )
        if optional:
            return f"Optional[{py_type}]"
        else:
            return py_type

    def record_type_to_pydantic(schema: Dict[str, Any]) -> Type[BaseModel]:
        """Convert a single avro record type to a pydantic class"""
        name = (
            schema["name"]
            if "." not in schema["name"]
            else schema["name"].split(".")[-1]
        )
        current = f"class {schema['name']}(BaseModel):\n"

        kwargs: Dict[str, Tuple[str, Any]] = {}

        if len(schema["fields"]) == 0:
            raise ValueError("Avro schema has no fields")

        for field in schema["fields"]:
            n = field["name"]
            t = get_python_type(field["type"])
            default = field.get("default")
            if "default" not in field:
                kwargs[n] = (t, ...)
                current += f"    {n}: {t}\n"
            elif isinstance(default, (bool, type(None))):
                kwargs[n] = (t, default)
                current += f"    {n}: {t} = {default}\n"
            else:
                kwargs[n] = (t, default)
                current += f"    {n}: {t} = {json.dumps(default)}\n"

        classes[name] = current
        pydantic_model = create_model(name, __module__=__name__, **kwargs)  # type: ignore
        return pydantic_model  # type: ignore

    return record_type_to_pydantic(schema)
