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


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


@dataclass
class AgentInfo:
    name: str
    role: AgentRole
    description: str


@dataclass
class ToolCall:
    tool: str
    args: dict
    result: str | None = None


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


def turn_to_message(turn: ConversationTurn) -> HumanMessage:
    return HumanMessage(content=turn.content)


def message_to_turn(msg, agent: AgentRole | None = None) -> ConversationTurn:
    msg_type = getattr(msg, "type", None) or getattr(msg, "role", "assistant")
    return ConversationTurn(
        role="user" if msg_type in ("human", "user") else "assistant",
        content=msg.content,
        agent=agent,
        model=getattr(msg, "model", None),
    )


class AgentState(TypedDict):
    messages: Annotated[list, None]
    user_id: str
    conversation_id: str

    # Agent routing
    current_agent: AgentRole
    active_agents: list[AgentRole]

    # Task tracking
    pending_tasks: list[dict]
    completed_tasks: list[dict]

    # Tool usage
    tool_calls: list[ToolCall]
    last_tool_result: str | None

    # Knowledge context
    knowledge_context: str | None
    graph_context: str | None

    # Planning
    plan: str | None
    plan_steps: list[dict]

    # Analysis
    analysis_report: str | None

    # Output control
    final_response: str | None
    should_respond: bool

    # Memory context (injected at start of each conversation)
    memory_context: str | None

    # User LLM config (for using user-configured models)
    user_llm_config: 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,
        active_agents=[AgentRole.MASTER],
        pending_tasks=[],
        completed_tasks=[],
        tool_calls=[],
        last_tool_result=None,
        knowledge_context=None,
        graph_context=None,
        plan=None,
        plan_steps=[],
        analysis_report=None,
        final_response=None,
        should_respond=True,
        memory_context=None,
        user_llm_config=None,
    )