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

import yaml
from init_helpers import AnyType, custom_dumps

from .info import Info
from .maybe_referencable_resource import MaybeReferencableResource
from .operation import Operation
from .schema.base_schema import BaseSchema
from .schema.build import SchemaBuilder
from .spec import Spec
from .spec_resource import SpecResource
from .spec_ref import SpecRef


class SpecGenerator:
    openapi_version: ClassVar[str] = "3.1.0"
    schema_builder: ClassVar[Callable[[AnyType], BaseSchema]] = SchemaBuilder.build_schema

    @staticmethod
    def resource_ref_condition(res: MaybeReferencableResource, parents: frozenset[SpecResource]) -> bool:
        return isinstance(res, BaseSchema) or len(parents) > 1

    @dataclass
    class Config:
        version: str = ''  # empty means autogenerate based on current datetime

    @dataclass
    class Context:
        company: str
        project_group: str
        project_name: str

    def __init__(
            self, config: Config, context: Context,
    ) -> None:
        self.version = config.version or self._get_version()
        self.config = config
        self.context = context
        self.info = Info(version=self.version,
                         title=f"{self.context.company}:{self.context.project_group}:{self.context.project_name}")
        self.resources: list[MaybeReferencableResource] = []

    def register(self, resource: MaybeReferencableResource) -> None:
        if isinstance(resource, Operation):
            if any(r.operation_id == r.operation_id for r in self.resources if isinstance(r, Operation)):
                KeyError(f"operation_id {resource.operation_id!r} already registered")
        self.resources.append(resource)

    def get_spec_dict(self, operation_filter: Callable[[Operation], bool] | None = None) -> dict:
        return self._get_spec(operation_filter=operation_filter).payload

    def get_spec_json(self, operation_filter: Callable[[Operation], bool] | None = None) -> bytes:
        return custom_dumps(self.get_spec_dict(operation_filter=operation_filter), indent=2)

    def get_spec_yaml(self, operation_filter: Callable[[Operation], bool] | None = None) -> bytes:
        return yaml.dump(json.loads(self.get_spec_json(operation_filter=operation_filter)), sort_keys=False).encode()

    def _wrap_non_spec_resource_element(self, element: Any) -> SpecResource:
        if isinstance(element, SpecResource | SpecRef):
            return element
        if isinstance(element, AnyType):
            return self.schema_builder(element)
        raise TypeError(f'Unsupported element type {type(element)=}, {element=}')

    def _get_spec(self, operation_filter: Callable[[Operation], bool] | None = None) -> Spec:
        resource_to_dependencies: dict[SpecResource, frozenset[SpecResource | SpecRef]] = {}
        pending: list[SpecResource] = (
            [r for r in self.resources if not isinstance(r, Operation) or operation_filter(r)]
            if operation_filter else self.resources[:])

        while pending:
            current: SpecResource | SpecRef = self._wrap_non_spec_resource_element(pending.pop())
            if isinstance(current, SpecRef):
                continue
            if current in resource_to_dependencies:
                continue

            current_dependencies: frozenset[SpecResource | SpecRef] = frozenset(
                self._wrap_non_spec_resource_element(dep) for dep in current.get_spec_dependencies()
            )
            resource_to_dependencies[current] = current_dependencies
            pending.extend(current_dependencies)

        resource_to_parents: dict[SpecResource, set[SpecResource | SpecRef]] = {}
        for resource, dependencies in resource_to_dependencies.items():
            for dependency in dependencies:
                resource_to_parents.setdefault(dependency, set()).add(resource)

        result = Spec(self.info, schema_builder=self._wrap_non_spec_resource_element)

        dependency_spec_refs: set[SpecRef] = set()
        for resource in TopologicalSorter(resource_to_dependencies).static_order():
            if isinstance(resource, SpecRef):
                dependency_spec_refs.add(resource)
                continue
            do_add = False
            key = None
            if isinstance(resource, MaybeReferencableResource) and (key := resource.get_key()):
                parents = resource_to_parents.get(resource, frozenset())
                do_add = not parents or not resource.can_be_inlined or self.resource_ref_condition(resource, parents)

            # print(resource)
            # print()
            if key and do_add:
                result.add_resource(resource, key)
            else:
                result.prepare_resource(resource)
        for ref in dependency_spec_refs:
            assert result.has_ref(ref)
        # print()
        return result

    @staticmethod
    def _get_version() -> str:
        now = datetime.datetime.now()
        return f"{now.year - 2021}.{now.month}.{now.day}{now.hour:02d}{now.minute:02d}"

