#  Copyright (C) 2025
#  ABM, Moscow
#
#  UNPUBLISHED PROPRIETARY MATERIAL.
#  ALL RIGHTS RESERVED.
#
#  Authors: Mike Orlov <m.orlov@abm-jsc.ru>
from dataclasses import dataclass, field
from typing import Self, Callable, Any, TypeVar

from frozendict import frozendict
from init_helpers import Jsonable, raise_if, AnyType

from openapi_tools.spec import SpecRef
from openapi_tools.spec.info import Info
from openapi_tools.spec.maybe_referencable_resource import MaybeReferencableResource
from openapi_tools.spec.spec_resource import SpecResource

K = TypeVar('K')
V = TypeVar('V')


class MappedDict(dict[K, V]):
    def __init__(self, mapper: Callable[[Any], K], extract_controller: Callable[[K], bool]) -> None:
        super().__init__()
        self.mapper = mapper
        self.extract_controller = extract_controller

    def __getitem__(self, key: Any) -> V:
        mapped_key = self.mapper(key)
        if self.extract_controller(mapped_key):
            return dict.__getitem__(self, mapped_key)
        return mapped_key

    def __setitem__(self, key: Any, value) -> None:
        dict.__setitem__(self, self.mapper(key), value)


@dataclass
class Spec:
    schema_builder: Callable[[AnyType], SpecResource]
    payload: dict[str, Jsonable] = field(init=False)
    resource_to_ref: dict[SpecResource, SpecRef | frozendict] = field(init=False)

    def __init__(self, info: Info, schema_builder: Callable[[AnyType], SpecResource]):
        self.schema_builder = schema_builder
        self.payload = {'openapi': '3.1.0', 'info': info.get_spec_dict({}), 'paths': {}}
        self.resource_to_ref = MappedDict(mapper=schema_builder,
                                          extract_controller=lambda x: isinstance(x, SpecResource))

    def add_resource(self: Self, resource: MaybeReferencableResource, key: str | None = None) -> None:
        key: str = key or resource.get_key()
        raise_if(key is None, ValueError(f'{key=} is not defined'))
        ref = resource.get_spec_ref()
        target: dict[str, Jsonable] = self.payload
        for step in resource.spec_path:
            target = target.setdefault(step, {})

        target[key] = (
            resource.merge_spec_dict(self.resource_to_ref, present_value)
            if (present_value := target.get(key))
            else resource.get_spec_dict(self.resource_to_ref))
        self.resource_to_ref[resource] = ref

    def prepare_resource(self: Self, resource: SpecResource | AnyType) -> None:
        if resource in self.resource_to_ref:
            return
        resource = self.schema_builder(resource)
        self.resource_to_ref[resource] = resource.get_spec_dict(self.resource_to_ref)

    def has_ref(self: Self, ref: SpecRef) -> bool:
        spec_path = ref.ref.lstrip("#/").split('/')
        assert spec_path, f'bad {spec_path=}'
        target: dict[str, Jsonable] = self.payload
        for step in spec_path[:-1]:
            target = target.setdefault(step, {})
        return spec_path[-1] in target
