import logging
from typing import TypedDict
from langgraph.graph import StateGraph, END
from smolagents import ToolCallingAgent, LiteLLMModel
from tools import tools
import yaml
import importlib.resources

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Define the state for our agent graph
class AgentState(TypedDict):
    messages: list
    question: str
    answer: str | None

class AgentNode:
    def __init__(self):
        # Load default prompt templates
        prompt_templates = yaml.safe_load(
            importlib.resources.files("smolagents.prompts").joinpath("toolcalling_agent.yaml").read_text()
        )
        
        # Log the default system prompt
        logger.info("Default system prompt:")
        logger.info("-" * 80)
        logger.info(prompt_templates["system_prompt"])
        logger.info("-" * 80)
        
#         # Define our custom system prompt
#         custom_system_prompt = "..."

#         # Update the system prompt in the loaded templates
#         prompt_templates["system_prompt"] = custom_system_prompt
        
        # Log our custom system prompt
        # logger.info("Custom system prompt:")
        # logger.info("-" * 80)
        # logger.info(custom_system_prompt)
        # logger.info("-" * 80)

        # In"itialize the model and agent
        self.model = LiteLLMModel(
            model="ollama/codellama",
            temperature=0.0,
            max_tokens=4096,
            top_p=0.9,
            frequency_penalty=0.0,
            presence_penalty=0.0,
            stop=["Observation:"],
        )
        
        self.agent = ToolCallingAgent(
            model=self.model,
            prompt_templates=prompt_templates,
            tools=tools
        )

    def __call__(self, state: AgentState) -> AgentState:
        try:
            # Process the question through the agent
            result = self.agent.run(state["question"])
            
            # Update the state with the answer
            state["answer"] = result
            return state
            
        except Exception as e:
            logger.error(f"Error in agent node: {str(e)}", exc_info=True)
            state["answer"] = f"Error: {str(e)}"
            return state

def build_agent_graph():
    # Create the graph
    graph = StateGraph(AgentState)
    
    # Add the agent node
    graph.add_node("agent", AgentNode())
    
    # Add edges
    graph.add_edge("agent", END)
    
    # Set the entry point
    graph.set_entry_point("agent")
    
    # Compile the graph
    return graph.compile()

# Create an instance of the compiled graph
agent_graph = build_agent_graph()