from dataclasses import dataclass
from typing import Dict, List, Union

from ._BaseModel import BaseModel
from ._RequestBodyModel import RequestBodyModel
from ._SchemaModel import SchemaModel
from ._SecuritySchemaModel import SecuritySchemaModel
from ..utils import set_dict_path_value


@dataclass
class ComponentsModel(BaseModel):
    security_schemes: Dict[str, SecuritySchemaModel] = None
    request_bodies: Dict[str, RequestBodyModel] = None
    schemas: Dict[str, SchemaModel] = None

    def add_security_scheme(self, name: str, model: SecuritySchemaModel):
        if self.security_schemes is None:
            self.security_schemes = {}
        self.security_schemes[name] = model

    def add_request_body(self, name: str, model: RequestBodyModel):
        if self.request_bodies is None:
            self.request_bodies = {}
        self.request_bodies[name] = model

    def add_schema(self, path: Union[str, List[str]], model: SchemaModel):
        if self.schemas is None:
            self.schemas = {}

        set_dict_path_value(
            dict_value=self.schemas,
            path=path,
            value=model
        )
