# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/014_AsyncAPI.ipynb.

# %% auto 0
__all__ = ['logger', 'sec_scheme_name_mapping', 'KafkaMessage', 'SecurityType', 'APIKeyLocation', 'SecuritySchema', 'KafkaBroker',
           'ContactInfo', 'KafkaServiceInfo', 'KafkaBrokers', 'yaml_file_cmp', 'export_async_spec']

# %% ../../nbs/014_AsyncAPI.ipynb 1
import json
import platform
import shutil
import subprocess  # nosec: B404: Consider possible security implications associated with the subprocess module.
import tempfile
from datetime import timedelta
from enum import Enum
from pathlib import Path
from typing import *

from pydantic import ConfigDict, BaseModel, Field, HttpUrl, model_serializer
from pydantic.type_adapter import TypeAdapter

from .aiokafka_consumer_loop import ConsumeCallable
from .docs_dependencies import _check_npm_with_local
from .helpers import unwrap_list_type
from .logger import get_logger
from fastkafka._components.producer_decorator import (
    ProduceCallable,
    unwrap_from_kafka_event,
)

# %% ../../nbs/014_AsyncAPI.ipynb 3
logger = get_logger(__name__)

# %% ../../nbs/014_AsyncAPI.ipynb 5
class KafkaMessage(BaseModel):
    # This following config is used to properly format timedelta in ISO 8601 format
    model_config = ConfigDict(ser_json_timedelta="iso8601")

# %% ../../nbs/014_AsyncAPI.ipynb 7
class SecurityType(str, Enum):
    plain = "plain"
    userPassword = "userPassword"
    apiKey = "apiKey"
    X509 = "X509"
    symmetricEncryption = "symmetricEncryption"
    asymmetricEncryption = "asymmetricEncryption"
    httpApiKey = "httpApiKey"
    http = "http"
    oauth2 = "oauth2"
    openIdConnect = "openIdConnect"
    scramSha256 = "scramSha256"
    scramSha512 = "scramSha512"
    gssapi = "gssapi"


class APIKeyLocation(str, Enum):
    user = "user"
    password = "password"  # nosec
    query = "query"
    header = "header"
    cookie = "cookie"


sec_scheme_name_mapping = {"security_type": "type", "api_key_loc": "in"}


class SecuritySchema(BaseModel):
    security_type: SecurityType = Field(..., example="plain")
    description: Optional[str] = Field(None, example="My security scheme")
    name: Optional[str] = Field(None, example="my_secret_scheme")
    api_key_loc: Optional[APIKeyLocation] = Field(None, example="user")
    scheme: Optional[str] = None
    bearerFormat: Optional[str] = None
    flows: Optional[str] = None
    openIdConnectUrl: Optional[str] = None

    def __init__(self, **kwargs: Any):
        for k, v in sec_scheme_name_mapping.items():
            if v in kwargs:
                kwargs[k] = kwargs.pop(v)
        super().__init__(**kwargs)

    def model_dump(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
        """Renames internal names of members ('security_type' -> 'type', 'api_key_loc' -> 'in')"""
        d = super().model_dump(*args, **kwargs)

        for k, v in sec_scheme_name_mapping.items():
            d[v] = d.pop(k)

        # removes None values
        d = {k: v for k, v in d.items() if v is not None}

        return d

    def model_dump_json(self, *args: Any, **kwargs: Any) -> str:
        """Serialize into JSON using model_dump()"""
        return json.dumps(self.model_dump(), *args, **kwargs)

# %% ../../nbs/014_AsyncAPI.ipynb 9
class KafkaBroker(BaseModel):
    """Kafka broker"""

    url: str = Field(..., example="localhost")
    description: str = Field("Kafka broker")
    port: Union[str, int] = Field("9092")
    protocol: str = Field("kafka")
    security: Optional[SecuritySchema] = None

    def model_dump(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
        """Makes port a variable and remove it from the dictionary"""
        d = super().model_dump(*args, **kwargs)
        if self.security:
            d["security"] = self.security.model_dump(*args, **kwargs)
        d["variables"] = {"port": {"default": str(self.port)}}
        d.pop("port")

        d = {k: v for k, v in d.items() if v is not None}

        return d

    def model_dump_json(self, *args: Any, **kwargs: Any) -> str:
        """Serialize into JSON using dict()"""
        return json.dumps(self.model_dump(), *args, **kwargs)

# %% ../../nbs/014_AsyncAPI.ipynb 12
class ContactInfo(BaseModel):
    name: str = Field(..., example="My company")
    url: HttpUrl = Field(..., example="https://www.github.com/mycompany")
    email: str = Field(..., example="noreply@mycompany.com")


class KafkaServiceInfo(BaseModel):
    title: str = Field("Title")
    version: str = Field("0.0.1")
    description: str = Field("Description of the service")
    contact: ContactInfo = Field(
        ...,
    )

# %% ../../nbs/014_AsyncAPI.ipynb 14
class KafkaBrokers(BaseModel):
    brokers: Dict[str, Union[List[KafkaBroker], KafkaBroker]]

    def model_dump(self, *args: Any, **kwargs: Any) -> Dict[str, Any]:
        """Transcribe brokers into bootstrap server groups"""
        d = super().model_dump(*args, **kwargs)

        brokers = {}
        for k, v in self.brokers.items():
            if isinstance(v, list):
                brokers.update(
                    {
                        f"{k}-bootstrap-server-{i}": u_v.model_dump()
                        for i, u_v in enumerate(v)
                    }
                )
            else:
                brokers.update({f"{k}": v.model_dump()})
        d["brokers"] = brokers
        d = {k: v for k, v in d.items() if v is not None}

        return d

    def model_dump_json(self, *args: Any, **kwargs: Any) -> str:
        """Serialize into JSON using dict()"""
        return json.dumps(self.model_dump(), *args, **kwargs)

# %% ../../nbs/014_AsyncAPI.ipynb 17
# T = TypeVar("T")


def _get_msg_cls_for_producer(f: ProduceCallable) -> Type[Any]:
    types = get_type_hints(f)
    return_type = types.pop("return", type(None))
    # @app.producer must define a return value
    if return_type == type(None):
        raise ValueError(
            f"Producer function must have a defined return value, got {return_type} as return value"
        )

    return_type = unwrap_from_kafka_event(return_type)
    return_type = unwrap_list_type(return_type)

    if not hasattr(return_type, "json"):
        raise ValueError(f"Producer function return value must have json method")
    return return_type  # type: ignore

# %% ../../nbs/014_AsyncAPI.ipynb 22
def _get_msg_cls_for_consumer(f: ConsumeCallable) -> Type[Any]:
    types = get_type_hints(f)
    return_type = types.pop("return", type(None))
    types_list = list(types.values())
    # @app.consumer does not return a value
    if return_type != type(None):
        raise ValueError(
            f"Consumer function cannot return any value, got {return_type}"
        )
    # @app.consumer first consumer argument must be a msg which is a subclass of BaseModel
    try:
        msg_type = types_list[0]

        msg_type = unwrap_list_type(msg_type)

        if not issubclass(msg_type, BaseModel):
            raise ValueError(
                f"Consumer function first param must be a BaseModel subclass msg, got {types_list}"
            )

        return msg_type  # type: ignore

    except IndexError:
        raise ValueError(
            f"Consumer function first param must be a BaseModel subclass msg, got {types_list}"
        )

# %% ../../nbs/014_AsyncAPI.ipynb 27
def _get_topic_dict(
    f: Callable[[Any], Any],
    direction: str = "publish",
) -> Dict[str, Any]:
    if not direction in ["publish", "subscribe"]:
        raise ValueError(
            f"direction must be one of ['publish', 'subscribe'], but it is '{direction}'."
        )

    #     msg_cls = None

    if direction == "publish":
        msg_cls = _get_msg_cls_for_producer(f)
    elif direction == "subscribe":
        msg_cls = _get_msg_cls_for_consumer(f)

    msg_schema = {"message": {"$ref": f"#/components/messages/{msg_cls.__name__}"}}
    if hasattr(f, "description"):
        msg_schema["description"] = getattr(f, "description")
    elif f.__doc__ is not None:
        msg_schema["description"] = f.__doc__  # type: ignore
    return {direction: msg_schema}

# %% ../../nbs/014_AsyncAPI.ipynb 31
def _get_channels_schema(
    consumers: Dict[str, ConsumeCallable],
    producers: Dict[str, ProduceCallable],
) -> Dict[str, Dict[str, Dict[str, Any]]]:
    topics = {}
    for ms, d in zip([consumers, producers], ["subscribe", "publish"]):
        for topic, f in ms.items():  # type: ignore
            topics[topic] = _get_topic_dict(f, d)
    return topics

# %% ../../nbs/014_AsyncAPI.ipynb 33
def _get_kafka_msg_classes(
    consumers: Dict[str, ConsumeCallable],
    producers: Dict[str, ProduceCallable],
) -> Set[Type[BaseModel]]:
    fc = [_get_msg_cls_for_consumer(consumer) for consumer in consumers.values()]
    fp = [_get_msg_cls_for_producer(producer) for producer in producers.values()]
    return set(fc + fp)


def _get_kafka_msg_definitions(
    consumers: Dict[str, ConsumeCallable],
    producers: Dict[str, ProduceCallable],
) -> Dict[str, Dict[str, Any]]:
    msg_classes = _get_kafka_msg_classes(consumers, producers)
    _, msg_definitions = TypeAdapter.json_schemas(
        [(msg_cls, "validation", TypeAdapter(msg_cls)) for msg_cls in msg_classes]
    )
    return msg_definitions

# %% ../../nbs/014_AsyncAPI.ipynb 35
def _get_example(cls: Type[BaseModel]) -> BaseModel:
    kwargs: Dict[str, Any] = {}
    for k, v in cls.model_fields.items():
        #         try:
        if hasattr(v, "json_schema_extra") and "example" in v.json_schema_extra:  # type: ignore
            example = v.json_schema_extra["example"]  # type: ignore
            kwargs[k] = example
    #         except:
    #             pass
    return json.loads(cls(**kwargs).model_dump_json())  # type: ignore

# %% ../../nbs/014_AsyncAPI.ipynb 37
def _add_example_to_msg_definitions(
    msg_cls: Type[BaseModel], msg_schema: Dict[str, Dict[str, Any]]
) -> None:
    try:
        example = _get_example(msg_cls)
    except Exception as e:
        example = None
    if example is not None:
        msg_schema["$defs"][msg_cls.__name__]["example"] = example


def _get_msg_definitions_with_examples(
    consumers: Dict[str, ConsumeCallable],
    producers: Dict[str, ProduceCallable],
) -> Dict[str, Dict[str, Any]]:
    msg_classes = _get_kafka_msg_classes(consumers, producers)
    msg_schema: Dict[str, Dict[str, Any]]
    _, msg_schema = TypeAdapter.json_schemas(
        [(msg_cls, "validation", TypeAdapter(msg_cls)) for msg_cls in msg_classes]
    )
    for msg_cls in msg_classes:
        _add_example_to_msg_definitions(msg_cls, msg_schema)
    msg_schema = (
        {k: {"payload": v} for k, v in msg_schema["$defs"].items()}
        if "$defs" in msg_schema
        else {}
    )

    return msg_schema

# %% ../../nbs/014_AsyncAPI.ipynb 39
def _get_security_schemes(kafka_brokers: KafkaBrokers) -> Dict[str, Any]:
    security_schemes = {}
    for key, broker in kafka_brokers.brokers.items():
        if isinstance(broker, list):
            kafka_broker = broker[0]
        else:
            kafka_broker = broker

        if kafka_broker.security is not None:
            security_schemes[f"{key}_default_security"] = json.loads(
                kafka_broker.security.model_dump_json()
            )
    return security_schemes

# %% ../../nbs/014_AsyncAPI.ipynb 41
def _get_components_schema(
    consumers: Dict[str, ConsumeCallable],
    producers: Dict[str, ProduceCallable],
    kafka_brokers: KafkaBrokers,
) -> Dict[str, Any]:
    definitions = _get_msg_definitions_with_examples(consumers, producers)
    msg_classes = [cls.__name__ for cls in _get_kafka_msg_classes(consumers, producers)]
    components = {
        "messages": {k: v for k, v in definitions.items() if k in msg_classes},
        "schemas": {k: v for k, v in definitions.items() if k not in msg_classes},
        "securitySchemes": _get_security_schemes(kafka_brokers),
    }
    substitutions = {
        f"#/$defs/{k}": f"#/components/messages/{k}"
        if k in msg_classes
        else f"#/components/schemas/{k}"
        for k in definitions.keys()
    }

    def _sub_values(d: Any, substitutions: Dict[str, str] = substitutions) -> Any:
        if isinstance(d, dict):
            d = {k: _sub_values(v) for k, v in d.items()}
        if isinstance(d, list):
            d = [_sub_values(k) for k in d]
        elif isinstance(d, str):
            for k, v in substitutions.items():
                if d == k:
                    d = v
        return d

    return _sub_values(components)  # type: ignore

# %% ../../nbs/014_AsyncAPI.ipynb 43
def _get_servers_schema(kafka_brokers: KafkaBrokers) -> Dict[str, Any]:
    servers = json.loads(kafka_brokers.model_dump_json(sort_keys=False))["brokers"]

    for key, kafka_broker in servers.items():
        if "security" in kafka_broker:
            servers[key]["security"] = [{f"{key}_default_security": []}]
    return servers  # type: ignore

# %% ../../nbs/014_AsyncAPI.ipynb 45
def _get_asyncapi_schema(
    consumers: Dict[str, ConsumeCallable],
    producers: Dict[str, ProduceCallable],
    kafka_brokers: KafkaBrokers,
    kafka_service_info: KafkaServiceInfo,
) -> Dict[str, Any]:
    #     # we don't use dict because we need custom JSON encoders
    info = json.loads(kafka_service_info.model_dump_json())
    servers = _get_servers_schema(kafka_brokers)
    #     # should be in the proper format already
    channels = _get_channels_schema(consumers, producers)
    components = _get_components_schema(consumers, producers, kafka_brokers)
    return {
        "asyncapi": "2.5.0",
        "info": info,
        "servers": servers,
        "channels": channels,
        "components": components,
    }

# %% ../../nbs/014_AsyncAPI.ipynb 47
def yaml_file_cmp(file_1: Union[Path, str], file_2: Union[Path, str]) -> bool:
    """Compares two YAML files and returns True if their contents are equal, False otherwise.

    Args:
        file_1: Path or string representing the first YAML file.
        file_2: Path or string representing the second YAML file.

    Returns:
        A boolean indicating whether the contents of the two YAML files are equal.
    """
    try:
        import yaml
    except Exception as e:
        msg = "Please install docs version of fastkafka using 'pip install fastkafka[docs]' command"
        logger.error(msg)
        raise RuntimeError(msg)

    def _read(f: Union[Path, str]) -> Dict[str, Any]:
        with open(f) as stream:
            return yaml.safe_load(stream)  # type: ignore

    d = [_read(f) for f in [file_1, file_2]]
    return d[0] == d[1]

# %% ../../nbs/014_AsyncAPI.ipynb 48
def _generate_async_spec(
    *,
    consumers: Dict[str, ConsumeCallable],
    producers: Dict[str, ProduceCallable],
    kafka_brokers: KafkaBrokers,
    kafka_service_info: KafkaServiceInfo,
    spec_path: Path,
    force_rebuild: bool,
) -> bool:
    try:
        import yaml
    except Exception as e:
        msg = "Please install docs version of fastkafka using 'pip install fastkafka[docs]' command"
        logger.error(msg)
        raise RuntimeError(msg)

    # generate spec file
    asyncapi_schema = _get_asyncapi_schema(
        consumers, producers, kafka_brokers, kafka_service_info
    )
    if not spec_path.exists():
        logger.info(
            f"Old async specifications at '{spec_path.resolve()}' does not exist."
        )
    spec_path.parent.mkdir(exist_ok=True, parents=True)
    with tempfile.TemporaryDirectory() as d:
        with open(Path(d) / "asyncapi.yml", "w") as f:
            yaml.dump(asyncapi_schema, f, sort_keys=False)
        spec_changed = not (
            spec_path.exists() and yaml_file_cmp(Path(d) / "asyncapi.yml", spec_path)
        )
        if spec_changed or force_rebuild:
            shutil.copyfile(Path(d) / "asyncapi.yml", spec_path)
            logger.info(
                f"New async specifications generated at: '{spec_path.resolve()}'"
            )
            return True
        else:
            logger.info(
                f"Keeping the old async specifications at: '{spec_path.resolve()}'"
            )
            return False

# %% ../../nbs/014_AsyncAPI.ipynb 50
def _generate_async_docs(
    *,
    spec_path: Path,
    docs_path: Path,
) -> None:
    _check_npm_with_local()
    cmd = [
        "npx",
        "-y",
        "-p",
        "@asyncapi/generator",
        "ag",
        f"{spec_path}",
        "@asyncapi/html-template",
        "-o",
        f"{docs_path}",
        "--force-write",
    ]
    # nosemgrep: python.lang.security.audit.subprocess-shell-true.subprocess-shell-true
    p = subprocess.run(  # nosec: B602, B603 subprocess call - check for execution of untrusted input.
        cmd,
        stderr=subprocess.STDOUT,
        stdout=subprocess.PIPE,
        shell=True if platform.system() == "Windows" else False,
    )
    if p.returncode == 0:
        logger.info(f"Async docs generated at '{docs_path}'")
        logger.info(f"Output of '$ {' '.join(cmd)}'{p.stdout.decode()}")
    else:
        logger.error(f"Generation of async docs failed!")
        logger.info(f"Output of '$ {' '.join(cmd)}'{p.stdout.decode()}")
        raise ValueError(
            f"Generation of async docs failed, used '$ {' '.join(cmd)}'{p.stdout.decode()}"
        )

# %% ../../nbs/014_AsyncAPI.ipynb 52
def export_async_spec(
    *,
    consumers: Dict[str, ConsumeCallable],
    producers: Dict[str, ProduceCallable],
    kafka_brokers: KafkaBrokers,
    kafka_service_info: KafkaServiceInfo,
    asyncapi_path: Union[Path, str],
    force_rebuild: bool = True,
) -> None:
    """Exports the AsyncAPI specification and documentation to the given path.

    Args:
        consumers: Dictionary of consumer functions, where the keys are the channel names and the values are the consumer functions.
        producers: Dictionary of producer functions, where the keys are the channel names and the values are the producer functions.
        kafka_brokers: KafkaBrokers object representing the Kafka brokers configuration.
        kafka_service_info: KafkaServiceInfo object representing the Kafka service info configuration.
        asyncapi_path: Path or string representing the base path where the specification and documentation will be exported.
        force_rebuild: Boolean indicating whether to force a rebuild of the specification file even if it already exists.
    """
    # generate spec file
    spec_path = Path(asyncapi_path) / "spec" / "asyncapi.yml"
    is_spec_built = _generate_async_spec(
        consumers=consumers,
        producers=producers,
        kafka_brokers=kafka_brokers,
        kafka_service_info=kafka_service_info,
        spec_path=spec_path,
        force_rebuild=force_rebuild,
    )

    # generate docs folder
    docs_path = Path(asyncapi_path) / "docs"

    if not is_spec_built and docs_path.exists():
        logger.info(
            f"Skipping generating async documentation in '{docs_path.resolve()}'"
        )
        return

    _generate_async_docs(
        spec_path=spec_path,
        docs_path=docs_path,
    )
