from operator import attrgetter
from typing import Any, cast

from mthds.client.models.working_memory import WorkingMemoryAbstract
from pydantic import Field, model_validator
from typing_extensions import override

from pipelex import log, pretty_print
from pipelex.core.memory.exceptions import (
    WorkingMemoryConsistencyError,
    WorkingMemoryStuffAttributeNotFoundError,
    WorkingMemoryStuffNotFoundError,
    WorkingMemoryTypeError,
)
from pipelex.core.stuffs.document_content import DocumentContent
from pipelex.core.stuffs.html_content import HtmlContent
from pipelex.core.stuffs.image_content import ImageContent
from pipelex.core.stuffs.list_content import ListContent
from pipelex.core.stuffs.mermaid_content import MermaidContent
from pipelex.core.stuffs.number_content import NumberContent
from pipelex.core.stuffs.stuff import Stuff
from pipelex.core.stuffs.stuff_artefact import StuffArtefact
from pipelex.core.stuffs.stuff_content import StuffContentType
from pipelex.core.stuffs.text_and_images_content import TextAndImagesContent
from pipelex.core.stuffs.text_content import TextContent
from pipelex.tools.misc.context_provider_abstract import ContextProviderAbstract
from pipelex.types import Self

MAIN_STUFF_NAME = "main_stuff"
BATCH_ITEM_STUFF_NAME = "BATCH_ITEM"
PRETTY_PRINT_MAX_LENGTH = 1000
TEST_DUMMY_NAME = "dummy_result"

StuffDict = dict[str, Stuff]
StuffArtefactDict = dict[str, StuffArtefact]


class WorkingMemory(WorkingMemoryAbstract[Stuff], ContextProviderAbstract):
    root: StuffDict = Field(default_factory=dict)
    aliases: dict[str, str] = Field(default_factory=dict)

    @model_validator(mode="after")
    def validate_stuff_names(self) -> Self:
        for key, stuff in self.root.items():
            if key.startswith("_") and key != BATCH_ITEM_STUFF_NAME:
                log.warning(f"Stuff key '{key}' starts with '_', which is reserved for params")

            if not stuff.stuff_name:
                self.root[key].stuff_name = key
            elif key not in {MAIN_STUFF_NAME, stuff.stuff_name}:
                log.warning(f"Stuff name '{stuff.stuff_name}' does not match the key '{key}'")
            elif stuff.stuff_name.startswith("_") and stuff.stuff_name != BATCH_ITEM_STUFF_NAME:
                log.warning(f"Stuff name '{stuff.stuff_name}' starts with '_', which is reserved for params")

        return self

    def pretty_print_summary(self):
        for stuff in self.root.values():
            content = stuff.content.rendered_plain()
            if len(content) > PRETTY_PRINT_MAX_LENGTH:
                content = content[:PRETTY_PRINT_MAX_LENGTH] + "..."
            pretty_print(content, title=f"{stuff.stuff_name} ({stuff.concept.code})")

    def make_deep_copy(self) -> Self:
        return self.model_copy(deep=True)

    def get_optional_stuff(self, name: str) -> Stuff | None:
        if named_stuff := self.root.get(name):
            return named_stuff
        if alias := self.aliases.get(name):
            return self.root.get(alias)
        return None

    def get_optional_main_stuff(self) -> Stuff | None:
        return self.get_optional_stuff(name=MAIN_STUFF_NAME)

    def get_main_stuff(self) -> Stuff:
        return self.get_stuff(name=MAIN_STUFF_NAME)

    def get_stuff(self, name: str) -> Stuff:
        if named_stuff := self.root.get(name):
            return named_stuff
        if alias := self.aliases.get(name):
            stuff = self.root.get(alias)
            if stuff is None:
                raise WorkingMemoryStuffNotFoundError(
                    variable_name=alias,
                    message=f"Alias '{alias}' points to a non-existent stuff '{name}'",
                )
            return stuff
        raise WorkingMemoryStuffNotFoundError(
            variable_name=name,
            message=f"Stuff '{name}' not found in working memory, valid keys are: {self.list_keys()}",
        )

    def get_stuffs(self, names: set[str]) -> list[Stuff]:
        the_stuffs: list[Stuff] = []
        for name in names:
            the_stuffs.append(self.get_stuff(name=name))
        return the_stuffs

    def get_existing_stuffs(self, names: set[str]) -> list[Stuff]:
        the_stuffs: list[Stuff] = []
        for name in names:
            if stuff := self.get_optional_stuff(name=name):
                the_stuffs.append(stuff)
        return the_stuffs

    def is_stuff_code_used(self, stuff_code: str) -> bool:
        return any(stuff.concept.code == stuff_code for stuff in self.root.values())

    def remove_stuff(self, name: str):
        self.root.pop(name, None)

    def remove_main_stuff(self):
        if MAIN_STUFF_NAME in self.root:
            del self.root[MAIN_STUFF_NAME]

    def set_stuff(self, name: str, stuff: Stuff):
        self.root[name] = stuff

    def is_stuff_exists(self, name: str) -> bool:
        return name in self.root or name in self.aliases

    def add_new_stuff(self, name: str, stuff: Stuff, aliases: list[str] | None = None):
        # TODO: Add unit tests for this method
        if self.is_stuff_code_used(stuff_code=stuff.stuff_code):
            msg = f"Stuff code '{stuff.stuff_code}' is already used by another stuff"
            raise WorkingMemoryConsistencyError(msg)
        if name in self.root or name in self.aliases:
            existing_stuff = self.get_stuff(name=name)
            if existing_stuff == stuff and name != TEST_DUMMY_NAME:
                log.warning(f"Key '{name}' already exists in WorkingMemory with the same stuff")
                return
            elif name != TEST_DUMMY_NAME:
                log.warning(f"Key '{name}' already exists in WorkingMemory and will be replaced by something different")
                log.verbose(f"Existing stuff: {existing_stuff}")
                log.verbose(f"New stuff: {stuff}")

        # it's a new stuff
        self.set_stuff(name=name, stuff=stuff)
        if aliases:
            for alias in aliases:
                self.set_alias(alias, name)

    def set_new_main_stuff(self, stuff: Stuff, name: str | None = None):
        # TODO: Add unit tests for this method
        if name:
            self.remove_main_stuff()
            self.add_new_stuff(name=name, stuff=stuff, aliases=[MAIN_STUFF_NAME])
            log.verbose(f"Setting new main stuff {name}: {stuff.concept.code} = '{stuff.short_desc}'")
        else:
            self.remove_alias_to_main_stuff()
            self.set_stuff(name=MAIN_STUFF_NAME, stuff=stuff)
            log.verbose(f"Setting new main stuff (unnamed): {stuff.concept.code} = '{stuff.short_desc}'")

    def set_alias(self, alias: str, target: str) -> None:
        """Add an alias pointing to a target name."""
        if alias == target:
            msg = f"Cannot create alias '{alias}' pointing to itself"
            raise WorkingMemoryConsistencyError(msg)
        if target not in self.root:
            msg = f"Cannot create alias to non-existent target '{target}'"
            raise WorkingMemoryConsistencyError(msg)
        self.aliases[alias] = target

    def add_alias(self, alias: str, target: str) -> None:
        """Add an alias pointing to a target name."""
        if alias in self.root:
            msg = f"Cannot add alias '{alias}' as it already exists"
            raise WorkingMemoryConsistencyError(msg)
        self.set_alias(alias=alias, target=target)

    def remove_alias(self, alias: str) -> None:
        """Remove an alias if it exists."""
        if alias in self.aliases:
            del self.aliases[alias]

    def remove_alias_to_main_stuff(self) -> None:
        """Remove the alias pointing to the main stuff if it exists."""
        self.remove_alias(alias=MAIN_STUFF_NAME)

    def get_aliases_for(self, target: str) -> list[str]:
        """Get all aliases pointing to a target name."""
        return [alias for alias, t in self.aliases.items() if t == target]

    def list_keys(self) -> list[str]:
        return list(self.root.keys()) + list(self.aliases.keys())

    ################################################################################################
    # ContextProviderAbstract
    ################################################################################################

    @override
    def generate_context(self) -> dict[str, Any]:
        artefact_dict: StuffArtefactDict = {}
        for name, stuff in self.root.items():
            artefact_dict[name] = stuff.make_artefact()
        for alias, target in self.aliases.items():
            artefact_dict[alias] = artefact_dict[target]
        return artefact_dict

    @override
    def get_typed_object_or_attribute(self, name: str, wanted_type: type[Any] | None = None, accept_list: bool = False) -> Any:
        """Retrieve a typed object or nested attribute from working memory.

        This method is primarily used to:

        - Find specific content types that require special handling when passed to LLMs,
          such as Images and Documents. These types need to be formatted differently
          in LLM prompts compared to plain text content.
        - Find the list of items to iterate over in a PipeBatch operation.

        Supports dot-notation for accessing nested attributes (e.g., "stuff_name.attribute.sub_attr").
        When the base stuff contains ListContent, extracts the specified attribute from each item.

        Args:
            name: The name of the stuff, optionally followed by dot-separated attribute path
                  (e.g., "my_stuff" or "my_stuff.nested_attr.deeper_attr")
            wanted_type: Optional type to validate the retrieved content against. If provided
                         and the content doesn't match, raises WorkingMemoryTypeError.
                         Commonly used with ImageContent, DocumentContent, etc.
            accept_list: If True, allows retrieval from ListContent stuffs. If False and the
                         content is ListContent, raises WorkingMemoryTypeError

        Returns:
            The retrieved content or attribute value. For ListContent with dot-notation,
            returns a flattened list of attribute values from each item.

        Raises:
            WorkingMemoryStuffNotFoundError: If the base stuff name doesn't exist
            WorkingMemoryStuffAttributeNotFoundError: If the attribute path is invalid
            WorkingMemoryTypeError: If type validation fails or ListContent is encountered
                                    when accept_list is False
        """
        # TODO: Refactor this method. In the python paradigm, we should not have those ".", but arrays with field names.
        if "." in name:
            parts = name.split(".", 1)  # Split only at the first dot
            base_name = parts[0]
            attr_path_str = parts[1]  # Keep the rest as a dot-separated string

            base_stuff = self.get_stuff(base_name)
            if isinstance(base_stuff.content, ListContent):
                if not accept_list:
                    raise WorkingMemoryTypeError(
                        variable_name=name,
                        message=f"Content of '{base_name}' is ListContent, but accept_list is False",
                    )
                the_items: list[Any] = []
                list_content: ListContent[Any] = base_stuff.content  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
                for item in list_content.items:
                    item_content = attrgetter(attr_path_str)(item)
                    # check type
                    if isinstance(item_content, list):
                        for item_content_item in item_content:  # pyright: ignore[reportUnknownVariableType]
                            if item_content_item is None:
                                continue
                            if wanted_type and not isinstance(item_content_item, wanted_type):
                                raise WorkingMemoryTypeError(
                                    variable_name=name,
                                    message=(
                                        f"Item of '{base_name}' is of type '{type(item_content_item).__name__}', "  # pyright: ignore[reportUnknownArgumentType]
                                        f"it should be '{wanted_type.__name__}'"
                                    ),
                                )
                            the_items.append(item_content_item)
                    else:
                        if item_content is None:
                            continue
                        if wanted_type and not isinstance(item_content, wanted_type):
                            raise WorkingMemoryTypeError(
                                variable_name=name,
                                message=(
                                    f"Item of item_content from '{base_name} [item] {attr_path_str}' is of type '{type(item_content).__name__}', "
                                    f"it should be '{wanted_type.__name__}'"
                                ),
                            )
                        the_items.append(item_content)
                return the_items
            try:
                stuff_content = attrgetter(attr_path_str)(base_stuff.content)
            except AttributeError as exc:
                raise WorkingMemoryStuffAttributeNotFoundError(
                    variable_name=name,
                    message=f"Stuff attribute not found in attribute path '{name}': {exc}",
                ) from exc

            # Sometimes, some stuff content are Optional, therefore can be None. So Do not impose a wanted type
            if stuff_content is not None and wanted_type is not None and not isinstance(stuff_content, wanted_type):
                raise WorkingMemoryTypeError(
                    variable_name=name,
                    message=f"Content at '{name}' is of type '{type(stuff_content).__name__}', it should be '{wanted_type.__name__}'",
                )

            return stuff_content

        top_level_content = self.get_stuff(name).content
        if isinstance(top_level_content, ListContent):
            top_level_content = cast("ListContent[Any]", top_level_content)
            if wanted_type is ListContent:
                return top_level_content.items
            else:
                if not accept_list:
                    raise WorkingMemoryTypeError(
                        variable_name=name,
                        message=f"Content of '{name}' is ListContent, but accept_list is False",
                    )
                top_level_content_items = self._get_typed_items_from_list_content(list_content=top_level_content, wanted_type=wanted_type)
                if top_level_content_items is None and wanted_type is not None:
                    raise WorkingMemoryTypeError(
                        variable_name=name,
                        message=f"Content of '{name}' is ListContent, but some of its items are not of type '{wanted_type.__name__}'",
                    )
                if top_level_content_items is not None and len(top_level_content_items) == 0:
                    raise WorkingMemoryTypeError(
                        variable_name=name,
                        message=f"Content of '{name}' is ListContent, but it is empty",
                    )
                return top_level_content_items

        if wanted_type is not None and not isinstance(top_level_content, wanted_type):
            raise WorkingMemoryTypeError(
                variable_name=name,
                message=f"Content of '{name}' is of type '{type(top_level_content).__name__}', it should be '{wanted_type.__name__}'",
            )

        return top_level_content

    def _get_typed_items_from_list_content(self, list_content: ListContent[Any], wanted_type: type[Any] | None) -> list[Any] | None:
        """Extract items from ListContent with optional type validation.

        Used to collect specific content types (e.g., Images, Documents) from lists
        that require special handling when passed to LLMs.

        Flattens nested lists and filters out None values. If any item fails type
        validation, returns None to indicate failure.

        Args:
            list_content: The ListContent to extract items from
            wanted_type: Optional type to validate each item against (e.g., ImageContent,
                         DocumentContent). If None, no type checking is performed

        Returns:
            A list of extracted items if all items pass type validation,
            or None if any item fails validation
        """
        the_items: list[Any] = []
        for item in list_content.items:
            item_content = item
            # check type
            if isinstance(item_content, list):
                for item_content_item in item_content:  # pyright: ignore[reportUnknownVariableType]
                    if item_content_item is None:
                        continue
                    if wanted_type and not isinstance(item_content_item, wanted_type):
                        return None
                    the_items.append(item_content_item)
            else:
                if item_content is None:
                    continue
                if wanted_type and not isinstance(item_content, wanted_type):
                    return None
                the_items.append(item_content)
        return the_items

    ################################################################################################
    # Stuff accessors
    ################################################################################################

    def get_stuff_as(self, name: str, content_type: type[StuffContentType]) -> StuffContentType:
        """Get stuff content as StuffContentType."""
        return self.get_stuff(name=name).content_as(content_type=content_type)

    def get_stuff_as_list(self, name: str, item_type: type[StuffContentType]) -> ListContent[StuffContentType]:
        """Get stuff content as ListContent with items of type StuffContentType.
        If the items are of possibly various types, use item_type=StuffContent.
        """
        return self.get_stuff(name=name).as_list_of_fixed_content_type(item_type=item_type)

    def get_list_stuff_first_item_as(self, name: str, item_type: type[StuffContentType]) -> StuffContentType:
        """Get stuff content as ListContent with items of type StuffContentType then return the first item."""
        return self.get_stuff_as_list(name=name, item_type=item_type).items[0]

    def get_stuff_as_text(self, name: str) -> TextContent:
        """Get stuff content as TextContent if applicable."""
        return self.get_stuff(name=name).as_text

    def get_stuff_as_str(self, name: str) -> str:
        """Get stuff content as string if applicable."""
        return self.get_stuff_as_text(name=name).text

    def get_stuff_as_image(self, name: str) -> ImageContent:
        """Get stuff content as ImageContent if applicable."""
        return self.get_stuff(name=name).as_image

    def get_stuff_as_text_and_image(self, name: str) -> TextAndImagesContent:
        """Get stuff content as TextAndImageContent if applicable."""
        return self.get_stuff(name=name).as_text_and_image

    def get_stuff_as_document(self, name: str) -> DocumentContent:
        """Get stuff content as DocumentContent if applicable."""
        return self.get_stuff(name=name).as_document

    def get_stuff_as_number(self, name: str) -> NumberContent:
        """Get stuff content as NumberContent if applicable."""
        return self.get_stuff(name=name).as_number

    def get_stuff_as_html(self, name: str) -> HtmlContent:
        """Get stuff content as HtmlContent if applicable."""
        return self.get_stuff(name=name).as_html

    def get_stuff_as_mermaid(self, name: str) -> MermaidContent:
        """Get stuff content as MermaidContent if applicable."""
        return self.get_stuff(name=name).as_mermaid

    ################################################################################################
    # Main stuff accessors
    ################################################################################################

    def main_stuff_as(self, content_type: type[StuffContentType]) -> StuffContentType:
        """Get main stuff content as StuffContentType."""
        return self.get_stuff_as(name=MAIN_STUFF_NAME, content_type=content_type)

    def main_stuff_as_list(self, item_type: type[StuffContentType]) -> ListContent[StuffContentType]:
        """Get main stuff content as ListContent with items of type StuffContentType.
        If the items are of possibly various types, use item_type=StuffContent.
        """
        return self.get_stuff_as_list(name=MAIN_STUFF_NAME, item_type=item_type)

    def main_list_stuff_first_item_as(self, item_type: type[StuffContentType]) -> StuffContentType:
        """Get main stuff content as first list item of type StuffContentType."""
        return self.get_list_stuff_first_item_as(name=MAIN_STUFF_NAME, item_type=item_type)

    @property
    def main_stuff_as_text(self) -> TextContent:
        """Get main stuff content as TextContent if applicable."""
        return self.get_stuff_as_text(name=MAIN_STUFF_NAME)

    @property
    def main_stuff_as_image(self) -> ImageContent:
        """Get main stuff content as ImageContent if applicable."""
        return self.get_stuff_as_image(name=MAIN_STUFF_NAME)

    @property
    def main_stuff_as_text_and_image(self) -> TextAndImagesContent:
        """Get main stuff content as TextAndImageContent if applicable."""
        return self.get_stuff_as_text_and_image(name=MAIN_STUFF_NAME)

    @property
    def main_stuff_as_number(self) -> NumberContent:
        """Get main stuff content as NumberContent if applicable."""
        return self.get_stuff_as_number(name=MAIN_STUFF_NAME)

    @property
    def main_stuff_as_html(self) -> HtmlContent:
        """Get main stuff content as HtmlContent if applicable."""
        return self.get_stuff_as_html(name=MAIN_STUFF_NAME)

    @property
    def main_stuff_as_mermaid(self) -> MermaidContent:
        """Get main stuff content as MermaidContent if applicable."""
        return self.get_stuff_as_mermaid(name=MAIN_STUFF_NAME)

    ################################################################################################
    # Serialization
    ################################################################################################

    def smart_dump(self) -> dict[str, Any]:
        """Serialize the working memory as a dictionary."""
        return self.model_dump(serialize_as_any=True)
