┌─────────────────────────────────────────────────────────────┐
│                     COALA AGENT                             │
├─────────────────────────────────────────────────────────────┤
│  MEMORY SYSTEM                                              │
│  ├── Working Memory (curto prazo): input, goals, cache      │
│  └── Long-Term Memory (longo prazo):                        │
│      ├── Episodic: experiências (ações + resultados)        │
│      ├── Semantic: fatos e conhecimento                     │
│      └── Procedural: como fazer (skills, workflows)         │
├─────────────────────────────────────────────────────────────┤
│  ACTION SPACE                                               │
│  ├── Internal: reasoning, retrieval, learning               │
│  └── External: tool_use, communication                      │
├─────────────────────────────────────────────────────────────┤
│  DECISION CYCLE                                             │
│  PLAN → EXECUTE → OBSERVE → LEARN → (repeat)               │
└─────────────────────────────────────────────────────────────┘

# Pseudocódigo simplificado
while not goal_achieved and iterations < max:
    # 1. PLAN: Propor e avaliar ações
    actions = plan_actions(working_memory, long_term_memory)
    best_action = select_best(actions)
    
    # 2. EXECUTE: Executar ação selecionada
    result = execute(best_action);
    
    # 3. OBSERVE: Processar resultado
    observation = observe(result);
    
    # 4. LEARN: Atualizar memórias
    learn(best_action, result, observation);
    
    # Atualizar working memory
    working_memory.update(observation);

class WorkingMemory:
    def __init__(self, max_tokens=8000):
        self.current_input = None      # Input do usuário
        self.active_goals = []         # Objetivos atuais
        self.intermediate_results = {} # Resultados de raciocínio
        self.retrieved_cache = {}      # Cache de memória recuperada
        self.max_tokens = max_tokens   # Limite de contexto

# Armazena: "Tentei X, resultado foi Y, lição Z"
episodic_memory.store_experience(
    action="chamar API com timeout=5s",
    result="TimeoutError",
    outcome="failure",
    lesson="API lenta requer timeout maior"
)

# Armazena: "API X requer autenticação OAuth2"
semantic_memory.store_fact(
    fact="API X usa OAuth2 com escopo 'read-write'",
    category="authentication",
    confidence=0.95
)

# Armazena: procedimentos e workflows
procedural_memory.store_procedure(
    name="debug_null_pointer",
    steps=["1. Check line in stack trace", 
           "2. Verify variable initialization",
           "3. Add null checks"],
    preconditions=["error is NullPointerException"]
)

def plan_actions(agent_state, goal):
    """
    Usa LLM para raciocinar sobre quais ações tomar.
    Retorna lista de ações candidatas ordenadas por prioridade.
    """
    prompt = f"""
    Contexto: {agent_state.working_memory}
    Objetivo: {goal}
    Memórias relevantes: {retrieve_similar_experiences(goal)}
    
    Quais ações devem ser tomadas? Liste em ordem de prioridade.
    Considere ações internas primeiro, depois externas.
    """
    
    response = llm.generate(prompt)
    return parse_actions(response)

def execute_action(action):
    if action.type == "internal":
        return execute_internal(action)
    else:
        return execute_external(action)

def execute_internal(action):
    if action.subtype == "reasoning":
        return llm.reason(action.prompt)
    elif action.subtype == "retrieval":
        return memory.retrieve(action.query)
    elif action.subtype == "learning":
        return memory.store(action.content)

def execute_external(action):
    if action.subtype == "tool_use":
        return call_tool(action.tool_name, action.parameters)
    elif action.subtype == "communication":
        return send_message(action.content)

def observe_result(raw_result):
    return {
        "content": raw_result,
        "success": is_successful(raw_result),
        "timestamp": now(),
        "relevance": assess_relevance(raw_result, goal)
    }

def learn(action, result, observation):
    # Armazenar experiência
    episodic_memory.store_experience(
        action=action,
        result=result,
        outcome="success" if observation["success"] else "failure"
    )
    
    # Extrair e armazenar fatos (se sucesso)
    if observation["success"]:
        facts = extract_facts(result)
        for fact in facts:
            semantic_memory.store_fact(fact)

class CoALAAgent:
    def __init__(self, llm_client):
        self.llm = llm_client
        self.working_memory = WorkingMemory()
        self.episodic = EpisodicMemory()
        self.semantic = SemanticMemory()
        self.procedural = ProceduralMemory()
        self.max_iterations = 10
    
    def run(self, user_input, goal=None):
        # Setup
        self.working_memory.current_input = user_input
        if goal:
            self.working_memory.active_goals.append(goal)
        
        # Ciclo principal
        for iteration in range(self.max_iterations):
            # 1. PLAN
            actions = self.plan()
            action = self.select_action(actions)
            
            # 2. EXECUTE
            result = self.execute(action)
            
            # 3. OBSERVE
            observation = self.observe(result)
            
            # 4. LEARN
            self.learn(action, result, observation)
            
            # Check completion
            if self.is_complete():
                break
        
        # Cleanup
        response = self.generate_response()
        self.working_memory.clear()
        return response
    
    def plan(self):
        # Recuperar contexto relevante
        context = self.working_memory.intermediate_results
        similar = self.episodic.retrieve_similar(
            self.working_memory.current_input
        )
        
        # Usar LLM para propor ações
        prompt = build_planning_prompt(context, similar)
        response = self.llm.generate(prompt)
        return self.parse_actions(response)
    
    def select_action(self, actions):
        # Seleciona primeira ação válida (simplificado)
        for action in actions:
            if self.check_preconditions(action):
                return action
        return actions[0] if actions else None
    
    def execute(self, action):
        if action.is_internal:
            return self.execute_internal(action)
        else:
            return self.execute_external(action)
    
    def observe(self, result):
        return {
            "content": result,
            "success": self.is_success(result),
            "timestamp": datetime.now()
        }
    
    def learn(self, action, result, observation):
        self.episodic.store_experience(
            action=action.description,
            result=result,
            outcome="success" if observation["success"] else "failure"
        )
        
        if observation["success"]:
            facts = self.extract_facts(result)
            for fact in facts:
                self.semantic.store_fact(fact)

# Exemplo: Decay de memórias
def retrieve_with_decay(self, query, k=5):
    results = self.retrieve(query, k*2)
    
    # Aplicar decay exponencial baseado na idade
    now = datetime.now()
    scored = []
    for result in results:
        age_days = (now - result.timestamp).days
        decay = 0.95 ** age_days
        relevance = calculate_relevance(result, query)
        scored.append((relevance * decay, result))
    
    scored.sort(reverse=True)
    return [r for _, r in scored[:k]]

agent = CoALAAgent(llm_client=openai_client)
agent.register_tool("web_search", search_function)

result = agent.run(
    user_input="Quais são os principais papers sobre RAG em 2024?",
    goal="Encontrar e sintetizar papers relevantes"
)

# Ciclo:
# 1. Recupera buscas similares passadas (episodic)
# 2. Busca conhecimento sobre RAG (semantic)
# 3. Executa web_search (external action)
# 4. Sintetiza resposta (internal reasoning)
# 5. Armazena experiência (learning)

class DebugAgent(CoALAAgent):
    def __init__(self):
        super().__init__()
        # Procedimentos pré-definidos
        self.procedural.store_procedure(
            name="debug_null_pointer",
            steps=["Check line", "Verify init", "Add checks"],
            preconditions=["NullPointerException"]
        )
    
    def run(self, error_log):
        # Recupera procedimento adequado
        procedure = self.procedural.retrieve(error_log)
        
        if procedure:
            return self.follow_procedure(procedure, error_log)
        else:
            # Explora solução nova
            return super().run(error_log)