"""
Wrapper for subscript type annotations, like `List[str]`.

Copyright 2024 Vlad Emelianov
"""

from collections.abc import Generator, Iterable, Iterator
from typing import Self

from mypy_boto3_builder.exceptions import TypeAnnotationError
from mypy_boto3_builder.import_helpers.import_record import ImportRecord
from mypy_boto3_builder.type_annotations.fake_annotation import FakeAnnotation
from mypy_boto3_builder.type_annotations.type_parent import TypeParent


class TypeSubscript(TypeParent):
    """
    Wrapper for subscript type annotations, like `List[str]`.

    Arguments:
        parent -- Parent type annotation.
        children -- Children type annotations.
    """

    def __init__(
        self,
        parent: FakeAnnotation,
        children: Iterable[FakeAnnotation] = (),
    ) -> None:
        self.parent: FakeAnnotation = parent
        self.children: list[FakeAnnotation] = list(children)

    def render(self) -> str:
        """
        Render type annotation to a valid Python code for local usage.

        Returns:
            A string with a valid type annotation.
        """
        result = self.parent.render()
        if self.children:
            children = ", ".join([i.render() for i in self.children])
            result = f"{result}[{children}]"

        return result

    def get_import_records(self) -> set[ImportRecord]:
        """
        Get all import records required for using type annotation.
        """
        result: set[ImportRecord] = set()
        result.update(self.parent.get_import_records())
        for child in self.children:
            result.update(child.get_import_records())
        return result

    def iterate_types(self) -> Generator[FakeAnnotation]:
        """
        Extract type annotations from children.
        """
        yield from self.parent.iterate_types()
        for child in self.children:
            yield from child.iterate_types()

    def add_child(self, child: FakeAnnotation) -> None:
        """
        Add new child to Substcript.
        """
        self.children.append(child)

    def is_dict(self) -> bool:
        """
        Whether subscript parent is Dict.
        """
        return self.parent.is_dict()

    def is_list(self) -> bool:
        """
        Whether subscript parent is List.
        """
        return self.parent.is_list()

    def __copy__(self) -> Self:
        """
        Create a copy of type annotation wrapper.
        """
        return self.__class__(
            parent=self.parent,
            children=list(self.children),
        )

    def get_local_types(self) -> list[FakeAnnotation]:
        """
        Get internal types generated by builder.
        """
        result: list[FakeAnnotation] = []
        for child in self.children:
            result.extend(child.get_local_types())
        return result

    def iterate_direct_type_annotations(self) -> Iterator[FakeAnnotation]:
        """
        Iterate over a parent and then children type annotations.
        """
        yield self.parent
        yield from self.children

    def get_children_types(self) -> set[FakeAnnotation]:
        """
        Extract required type annotations from attributes.
        """
        result: set[FakeAnnotation] = set()
        for child in self.children:
            result.update(child.iterate_types())
        return result

    def replace_child(self, child: FakeAnnotation, new_child: FakeAnnotation) -> Self:
        """
        Replace child type annotation with a new one.
        """
        if self.parent is child:
            self.parent = new_child
            return self

        if child not in self.children:
            raise TypeAnnotationError(f"Child not found: {child}")

        index = self.children.index(child)
        self.children[index] = new_child
        return self

    def iterate_children(self) -> Iterator[FakeAnnotation]:
        """
        Iterate over children.
        """
        yield from self.children
