from typing_extensions import (
    TypedDict,
    List
)
from typing import (
    Optional,
)
from dataclasses import dataclass, field
from typing import Optional

from langgraph.graph import MessagesState

from ea4all.src.ea4all_gra.data import (
    ListRequirement, 
    ListObjective, 
    UserJourney, 
    StakeholderList, 
    BusinessCapability, 
    LandscapeAsIs, 
    CapabilityGap, 
    Principles
)

# Optional, the InputState is a restricted version of the State that is used to
# define a narrower interface to the outside world vs. what is maintained
# internally.
class InputState(MessagesState):
    """Represents the input state for the agent.

    This class defines the structure of the input state, which includes
    the messages exchanged between the user and the agent. It serves as
    a restricted version of the full State, providing a narrower interface
    to the outside world compared to what is maintained internally.
    """

    """Attributes:
        business_query: a business requirement is the starting point of the TOGAF process
    """
    #business_query: Optional[Annotated[List[str], Field(
    #    description="A business requirement is the starting point of the TOGAF process."), operator.add]] 
    #business_query: str
    pass

class OutputState(TypedDict):
    """Represents te output state for the agent."""
    vision_target: Optional[str]
    architecture_runway: Optional[str]

@dataclass(kw_only=True)
class OverallState(InputState, OutputState):
    """Represents the state of a Togaf system."""

    """
    Attributes:
    - user_feedback: used to capture additional information needed from the user by the graph
    - business_query: a business requirement is the starting point of the TOGAF process
    - query_status (Optional[bool]): Indicates the status of the query. Default value is False.
    - messages (Optional[Annotated[list[AnyMessage], add_messages]]): A list of messages associated with the state.
    - stakeholder (Optional[StakeholderList]): Represents the list of stakeholders.
    - principles (Optional[Principles]): Represents the principles of the Togaf system.
    - requirement (Optional[ListRequirement]): Represents the list of requirements.
    - intent (Optional[ListObjective]): Represents the list of objectives.
    - userjourney (Optional[UserJourney]): Represents the user journey of the Togaf system.
    - biz_capability (Optional[BusinessCapability]): Represents the business capability of the Togaf system.
    - landscape_asis (Optional[List[str]]): Represents the list of landscape as-is.
    - identified_asis (Optional[LandscapeAsIs]): Represents the identified landscape as-is.
    - landscape_gap (Optional[CapabilityGap]): Represents the capability gap of the landscape.
    - vision_target (Optional[str]): Represents the vision target of the Togaf system.
    - architecture_runway (Optional[str]): Represents the architecture runway of the Togaf system.
    - next (Optional[str]): Represents the next step in the Togaf system.
    """

    business_query: str
    query_status: Optional[bool]
    stakeholder: Optional[StakeholderList]
    principles: Optional[Principles]
    requirement: Optional[ListRequirement]
    intent: Optional[ListObjective]
    userjourney: Optional[UserJourney]
    biz_capability: Optional[BusinessCapability]
    landscape_asis: Optional[List[str]]
    identified_asis: Optional[LandscapeAsIs]
    landscape_gap: Optional[CapabilityGap]
    next: Optional[str]