from typing import List, Dict, Any
from abc import ABC, abstractmethod

class Agent:
    def __init__(self, llm, tools, memory):
        self.llm = llm
        self.tools = tools
        self.memory = memory
        self.planning_module = PlanningModule(llm)

    def perceive(self, environment: Dict[str, Any]) -> Dict[str, Any]:
        """Observe and process environmental inputs."""
        observations = {
            'user_input': environment.get('input'),
            'context': self.memory.get_context(),
            'tool_results': environment.get('tool_results', [])
        }
        return observations

    def think(self, observations: Dict[str, Any]) -> Dict[str, Any]:
        """Plan actions based on observations."""
        plan = self.planning_module.plan(
            goal=observations['user_input'],
            context=observations['context'],
            available_tools=self.tools
        )
        return plan

    def act(self, plan: Dict[str, Any]) -> Dict[str, Any]:
        """Execute the planned action."""
        if plan['tool_use']:
            tool = self.get_tool(plan['tool_name'])
            result = tool.execute(plan['tool_args'])
            return {'action': 'tool_use', 'result': result}
        else:
            response = self.llm.generate(plan['prompt'])
            return {'action': 'generate', 'response': response}

    def observe(self, action_result: Dict[str, Any]):
        """Update memory based on action results."""
        self.memory.update(action_result)

class PlanningModule:
    def __init__(self, llm):
        self.llm = llm

    def plan(self, goal, context, available_tools):
        planning_prompt = f"""Given the goal: {goal}

Context: {context}

Available tools: {available_tools}

Create a step-by-step plan to achieve this goal.
Output as JSON with 'steps' and 'tool_use' fields."""

        plan = self.llm.generate(planning_prompt)
        return self.parse_plan(plan)

    def parse_plan(self, plan_text):
        # Parse LLM output into structured plan
        return {
            'steps': extract_steps(plan_text),
            'tool_use': check_tool_use(plan_text),
            'tool_name': extract_tool_name(plan_text),
            'tool_args': extract_tool_args(plan_text),
            'prompt': extract_prompt(plan_text)
        }

class ReActAgent(Agent):
    def __init__(self, llm, tools, memory):
        super().__init__(llm, tools, memory)
        self.max_iterations = 10

    def run(self, task):
        for i in range(self.max_iterations):
            # Thought
            thought = self.think(task)

            # Action
            if thought['action_type'] == 'finish':
                return thought['answer']

            action = self.act(thought)

            # Observation
            self.observe(action)

            # Update task with observation
            task = f"{task}\nObservation: {action['result']}"

        return "Max iterations reached"