from dataclasses import dataclass, field
from typing import TypedDict, Annotated, Sequence
from enum import Enum

from langchain_core.messages import BaseMessage
from langgraph.graph.message import add_messages


class AgentRole(str, Enum):
    MASTER = "master"
    SCHEDULE_PLANNER = "schedule_planner"
    EXECUTOR = "executor"
    LIBRARIAN = "librarian"
    ANALYST = "analyst"


@dataclass
class ConversationTurn:
    role: str  # "user" | "assistant"
    content: str
    agent: AgentRole | None = None
    model: str | None = None


class AgentState(TypedDict):
    # Core message history with add_messages reducer
    messages: Annotated[list[BaseMessage], add_messages]
    
    # Session identifiers
    user_id: str
    conversation_id: str

    # Agent routing state
    current_agent: str | None
    next_step: str | None  # For explicit graph routing
    
    # Traceability
    agent_trace: list[str]
    
    # Task & Entity Tracking (Business Logic)
    pending_tasks: list[dict]
    completed_tasks: list[dict]
    created_entities: list[dict]

    # Context summaries (for long-term or cross-agent context)
    knowledge_context: str | None
    schedule_context_summary: str | None
    analysis_report: str | None

    # Output control
    final_response: str | None
    
    # Memory & Environment
    memory_context: str | None
    current_datetime_context: str | None
    
    # Configuration
    user_llm_config: dict | None
    provider_capabilities: dict | None


def initial_state(user_id: str, conversation_id: str) -> AgentState:
    return AgentState(
        messages=[],
        user_id=user_id,
        conversation_id=conversation_id,
        current_agent=AgentRole.MASTER.value,
        next_step=None,
        agent_trace=[AgentRole.MASTER.value],
        pending_tasks=[],
        completed_tasks=[],
        created_entities=[],
        knowledge_context=None,
        schedule_context_summary=None,
        analysis_report=None,
        final_response=None,
        memory_context=None,
        current_datetime_context=None,
        user_llm_config=None,
        provider_capabilities=None,
    )