import engine
import json
import random

class MycoController:
    def __init__(self):
        self.position = [1, 1]  # 3x3 Grid (0-2)
        self.history = []

    def get_agent_decision(self, current_mushroom, collection):
        """
        Queries Gemma for a decision. 
        Note: We use engine._llm but provide a specific 'Agent' prompt.
        """
        prompt = f"""
        Current Position: {self.position}
        Mushroom in clearing: {'Yes' if current_mushroom else 'No'}
        Collection count: {len(collection)}
        
        Decide your next action: 'move', 'search', 'study', 'collect', or 'wait'.
        If 'move', provide target coordinate (e.g., [1, 2]).
        
        Respond ONLY in JSON format:
        {{"action": "...", "target": [x, y], "thought": "..."}}
        """
        
        # We reuse the existing engine._llm functionality
        response = engine._llm(prompt) 
        
        # Simple parser
        try:
            # Extract JSON from potential markdown text
            start = response.find("{")
            end = response.rfind("}") + 1
            return json.loads(response[start:end])
        except:
            return {"action": "wait", "target": None, "thought": "Thinking..."}

    def run_tick(self, current_mushroom, collection):
        """
        This is the heartbeat of the AI. 
        It decides and then executes via engine functions.
        """
        decision = self.get_agent_decision(current_mushroom, collection)
        action = decision.get("action")
        
        result = {"action_taken": action, "thought": decision.get("thought")}
        
        # EXECUTION LAYER (calling existing engine functions)
        if action == "move":
            self.position = decision.get("target", self.position)
        
        elif action == "search":
            # We call the engine function directly
            mushroom, current, history = engine.discover_mushroom(collection)
            result["data"] = {"mushroom": mushroom, "current": current}
            
        elif action == "collect":
            if current_mushroom:
                coll, hist = engine.collect_current(current_mushroom, collection, self.history)
                result["data"] = {"collection": coll}
        
        return result