from typing import Any, Dict

from args import Args
from graph import State
from graph_builder import GraphBuilder


class Alfred:

    def __init__(self):
        print("Agent initialized.")

        self.graph_builder = GraphBuilder()
        self.agent_graph = self.graph_builder.build_agent_graph()

    def __call__(self, question: str) -> str:
        print(f"Agent received question (first 50 chars): {question[:50]}...")
        result = self.process_query(question)
        response = result["final_response"]
        print(f"Agent processed the response: {response}")

        return response

    def process_query(self, query: str) -> Dict[str, Any]:
        """
        Process a query through the agent graph.

        Args:
            query: The initial query to process

        Returns:
            The final state of the graph execution
        """
        initial_state: State = {
            "initial_query": query,
            "messages": [query],  # Manager's context
            "task_progress": [],  # Solver's context
            "audit_interval": Args.AlfredParams.AUDIT_INTERVAL,
            "manager_queries": 0,
            "solver_queries": 0,
            "max_interactions": Args.AlfredParams.MAX_INTERACTIONS,
            "max_solving_effort": Args.AlfredParams.MAX_SOLVING_EFFORT,
            "final_response": None
        }

        result = self.agent_graph.invoke(initial_state)
        return result