class OptimizedGoalGenerator:
    def __init__(self):
        self.all_goals = None
        self.need_config = {
            'health': {'threshold': 40, 'weight': 1.0, 'name': 'Health'},
            'fullness': {'threshold': 40, 'weight': 1.1, 'name': 'Food'},
            'hydration': {'threshold': 40, 'weight': 1.1, 'name': 'Water'},
            'wakefulness': {'threshold': 70, 'weight': 1.1, 'name': 'Sleep'}
        }

        self.crafting_recipes = {
            'wood_pickaxe': {'requirements': {'wood': 1}, 'priority': 1.2},
            'stone_pickaxe': {'requirements': {'wood': 1, 'stone': 1}, 'priority': 1.5},
            'iron_pickaxe': {'requirements': {'wood': 1, 'stone': 1, 'coal': 1, 'iron': 1}, 'priority': 1.8},
            'wood_sword': {'requirements': {'wood': 1}, 'priority': 1.1},
            'stone_sword': {'requirements': {'wood': 1, 'stone': 1}, 'priority': 1.3},
            'iron_sword': {'requirements': {'wood': 1, 'stone': 1, 'coal': 1, 'iron': 1}, 'priority': 1.6}
        }

        self.resource_info = {
            'wood': {'tool': None, 'priority': 1.1},
            'stone': {'tool': 'wood_pickaxe', 'priority': 1.4},
            'coal': {'tool': 'wood_pickaxe', 'priority': 1.5},
            'iron': {'tool': 'stone_pickaxe', 'priority': 1.7},
            'diamond': {'tool': 'iron_pickaxe', 'priority': 2.0},
            'water': {'tool': None, 'priority': 1.0},
            'sapling': {'tool': None, 'priority': 0.9}
        }

    def _parse_observation(self, text_obs):
        obs = {
            'environment': set(),
            'inventory': {},
            'status': {},
            'dangers': set(),
            'edibles': set(),
            'resources': set()
        }

        sections = [s.strip() for s in text_obs.split('\n') if s.strip()]

        if sections and sections[0].startswith('You see:'):
            seen_items = sections[0][8:].strip().split(',')
            for item in seen_items:
                item = item.strip()
                if item:
                    obs['environment'].add(item)
                    if item in ['zombie', 'skeletons']:
                        obs['dangers'].add(item)
                    elif item in ['cows', 'ripe plant']:
                        obs['edibles'].add(item)
                    elif item in ['tree', 'stone', 'coal', 'iron', 'diamond', 'water', 'sapling']:
                        obs['resources'].add(item)

        if len(sections) > 1 and sections[1].startswith('Inventory:'):
            inv_items = sections[1][10:].strip()
            if inv_items:
                for item_pair in inv_items.split(','):
                    if ':' in item_pair:
                        name, val = item_pair.split(':', 1)
                        obs['inventory'][name.strip()] = int(val.strip())

        if len(sections) > 2 and sections[2].startswith('Status:'):
            status_items = sections[2][7:].strip().split(',')
            for stat in status_items:
                if ':' in stat:
                    name, val = stat.split(':', 1)
                    name = name.strip().lower()
                    if '%' in val:
                        val = val.replace('%', '')
                    obs['status'][name] = int(val.strip())

        return obs

    def _evaluate_needs(self, status):
        need_scores = []
        for need, config in self.need_config.items():
            current = status.get(need, 100)
            if current < config['threshold']:
                severity = (config['threshold'] - current) / config['threshold']
                need_scores.append({
                    'type': need,
                    'name': config['name'],
                    'score': severity * config['weight'],
                    'value': current
                })
        return sorted(need_scores, key=lambda x: -x['score'])

    def _generate_survival_goals(self, needs, obs):
        goals = []

        need_handlers = {
            'health': {
                'condition': lambda need, obs: obs['dangers'],
                'goal': lambda need, obs: {
                    'goal': 'Avoid enemies to recover health',
                    'priority': 1.2 + need['score'],
                    'type': 'safety'
                },
                'else_goal': lambda need, obs: {
                    'goal': 'Rest to recover health',
                    'priority': 1.1 + need['score'],
                    'type': 'health'
                }
            },
            'fullness': {
                'condition': lambda need, obs: 'ripe plant' in obs['edibles'],
                'goal': lambda need, obs: {
                    'goal': 'Harvest ripe plant for food',
                    'priority': 1.2 + need['score'],
                    'type': 'food'
                },
                'else_condition': lambda need, obs: 'cows' in obs['edibles'],
                'else_goal': lambda need, obs: {
                    'goal': 'Hunt cows for meat' if obs['inventory'].get('wood sword',
                                                                          0) > 0 else 'Approach cows for possible food',
                    'priority': 1.2 + need['score'] if obs['inventory'].get('wood sword', 0) > 0 else 1.2 + need[
                        'score'],
                    'type': 'food'
                },
                'else_else_goal': lambda need, obs: {
                    'goal': 'Find cows',
                    'priority': 1.0 + need['score'],
                    'type': 'food'
                }
            },
            'hydration': {
                'condition': lambda need, obs: 'water' in obs['resources'],
                'goal': lambda need, obs: {
                    'goal': 'Drink water',
                    'priority': 1.1 + need['score'],
                    'type': 'water'
                },
                'else_goal': lambda need, obs: {
                    'goal': 'Find water',
                    'priority': 1.0 + need['score'],
                    'type': 'water'
                }
            },
            'wakefulness': {
                'goal': lambda need, obs: {
                    'goal': 'Sleep',
                    'priority': 1.0 + need['score'],
                    'type': 'sleep'
                }
            }
        }

        for need in needs:
            need_type = need['type']
            if need_type in need_handlers:
                handler = need_handlers[need_type]
                if 'condition' in handler and handler['condition'](need, obs):
                    goals.append(handler['goal'](need, obs))
                elif 'else_condition' in handler and handler['else_condition'](need, obs):
                    goals.append(handler['else_goal'](need, obs))
                elif 'else_else_goal' in handler:
                    goals.append(handler['else_else_goal'](need, obs))
                elif 'goal' in handler:
                    goals.append(handler['goal'](need, obs))

        return goals

    def _generate_progression_goals(self, obs):
        goals = []
        inv = obs['inventory']

        has_table = 'table' in obs['environment']
        wood_count = inv.get('wood', 0)

        if not has_table and wood_count >= 2:
            goals.append({
                'goal': 'Place crafting table',
                'priority': 1.5,
                'type': 'crafting'
            })
        elif not has_table:
            goals.append({
                'goal': f'Collect wood for table ({2 - wood_count} needed)',
                'priority': 1.3,
                'type': 'gather'
            })

        # 工具制作行动
        for tool, recipe in self.crafting_recipes.items():
            if inv.get(tool, 0) == 0:
                can_craft = all(inv.get(res, 0) >= amt for res, amt in recipe['requirements'].items())
                if can_craft and has_table:
                    goals.append({
                        'goal': f'Craft {tool.replace("_", " ")}',
                        'priority': recipe['priority'],
                        'type': 'crafting'
                    })

        for resource in obs['resources']:
            if resource in self.resource_info:
                info = self.resource_info[resource]
                has_tool = info['tool'] is None or inv.get(info['tool'], 0) > 0

                if has_tool:
                    goal = {
                        'goal': f'Collect {resource}',
                        'priority': info['priority'],
                        'type': 'gather'
                    }

                    for tool, recipe in self.crafting_recipes.items():
                        if resource in recipe['requirements'] and inv.get(tool, 0) == 0:
                            goal['priority'] *= 1.2
                            break

                    goals.append(goal)
                else:
                    goals.append({
                        'goal': f'Acquire {info["tool"]} to collect {resource}',
                        'priority': info['priority'] * 0.8,
                        'type': 'prepare'
                    })

        return goals

    def determine_goal(self, text_obs):
        self.all_goals = None

        obs = self._parse_observation(text_obs)
        needs = self._evaluate_needs(obs['status'])

        survival_goals = self._generate_survival_goals(needs, obs)
        progression_goals = self._generate_progression_goals(obs)

        self.all_goals = survival_goals + progression_goals

        best_goal = max(self.all_goals, key=lambda x: x['priority'])

        if (best_goal['type'] == 'explore' and
                'tree' in obs['resources'] and
                obs['inventory'].get('wood', 0) < 2):
            return {
                'goal': 'Collect wood for basic tools',
                'priority': 1.2,
                'type': 'gather'
            }
        return best_goal

    def get_augmentation_more_goal_text(self, text_obs):
        _ = self.determine_goal(text_obs)
        all_goals = [
            {
                'goal': 'Explore the environment',
                'priority': 1,
                'type': 'explore'
            },
            {
                'goal': 'Collect wood for basic tools',
                'priority': 0.8,
                'type': 'gather'
            },
            {
                'goal': 'Place crafting table',
                'priority': 0.5,
                'type': 'crafting'
            }
        ]
        for i in range(3-len(self.all_goals)):
            self.all_goals.append(all_goals[i])
        top_three_goals = sorted(self.all_goals, key=lambda x: x['priority'], reverse=True)[:3]
        priority = top_three_goals[0]['priority'] + top_three_goals[1]['priority'] + top_three_goals[2]['priority']
        top_three_goals[0]['priority'] /= priority
        top_three_goals[1]['priority'] /= priority
        top_three_goals[2]['priority'] /= priority
        return top_three_goals




class AdvancedGoalGenerator(OptimizedGoalGenerator):
    def __init__(self):
        super().__init__()
        self.all_goals = None
        # Define need-priority configuration
        self.need_config = {
            'health': {'threshold': 40, 'weight': 1.0, 'name': 'Health'},
            'fullness': {'threshold': 40, 'weight': 1.0, 'name': 'Food'},
            'hydration': {'threshold': 40, 'weight': 1.0, 'name': 'Water'},
            'wakefulness': {'threshold': 70, 'weight': 1.0, 'name': 'Sleep'}
        }

        # Tool-crafting tree
        self.crafting_recipes = {
            'wood pickaxe': {'requirements': {'wood': 1}, 'priority': 1.8, 'achievement': 'make_wood_pickaxe'},
            'stone pickaxe': {'requirements': {'wood': 1, 'stone': 1}, 'priority': 1.7, 'achievement': 'make_stone_pickaxe'},
            'iron pickaxe': {'requirements': {'wood': 1, 'stone': 1, 'coal': 1, 'iron': 1}, 'priority': 1.6, 'achievement': 'make_iron_pickaxe'},
            'wood sword': {'requirements': {'wood': 1}, 'priority': 1.7, 'achievement': 'make_wood_sword'},
            'stone sword': {'requirements': {'wood': 1, 'stone': 1}, 'priority': 1.6, 'achievement': 'make_stone_sword'},
            'iron sword': {'requirements': {'wood': 1, 'stone': 1, 'coal': 1, 'iron': 1.5}, 'priority': 1.1, 'achievement': 'make_iron_sword'}
        }

        # Resource-collection configuration
        self.resource_info = {
            'wood': {'tool': None, 'priority': 1.1, 'achievement': 'collect_wood'},
            'stone': {'tool': 'wood pickaxe', 'priority': 1.4, 'achievement': 'collect_stone'},
            'coal': {'tool': 'wood pickaxe', 'priority': 1.5, 'achievement': 'collect_coal'},
            'iron': {'tool': 'stone pickaxe', 'priority': 1.7, 'achievement': 'collect_iron'},
            'diamond': {'tool': 'iron pickaxe', 'priority': 2.0, 'achievement': 'collect_diamond'},
            'water': {'tool': None, 'priority': 1.0, 'achievement': 'collect_drink'},
            'sapling': {'tool': None, 'priority': 0.9, 'achievement': 'collect_sapling'},
            'food': {'tool': None, 'priority': 1.3, 'achievement': 'eat_plant'},
            'ripe plant': {'tool': None, 'priority': 1.4, 'achievement': 'eat_plant'},
            'cow': {'tool': 'wood sword', 'priority': 1.5, 'achievement': 'eat_cow'}
        }

        # Item-placement configuration
        self.placement_info = {
            'table': {'requirements': {'wood': 2}, 'priority': 1.2, 'achievement': 'place_table'},
            'furnace': {'requirements': {'stone': 1}, 'priority': 1.4, 'achievement': 'place_furnace'}, # stone 1
            'plant': {'requirements': {'sapling': 1}, 'priority': 1.3, 'achievement': 'place_plant'},
            'stone': {'requirements': {'stone': 1}, 'priority': 1.1, 'achievement': 'place_stone'}
        }

        # Achievement tracking
        self.achievements = {
            'coal': False,
            'diamond': False,
            'collect_drink': False,
            'iron': False,
            'sapling': False,
            'stone': False,
            'wood': False,
            'defeat_skeleton': False,
            'defeat_zombie': False,
            'eat_cow': False,
            'eat_plant': False,
            'iron pickaxe': False,
            'iron sword': False,
            'stone pickaxe': False,
            'stone sword': False,
            'wood pickaxe': False,
            'wood sword': False,
            'place_furnace': False,
            'place_plant': False,
            'place_stone': False,
            'place_table': False,
            'wake_up': False
        }

    def _parse_observation(self, text_obs):
        """Enhanced observation parsing"""
        obs = {
            'environment': set(),
            'inventory': {},
            'status': {},
            'dangers': set(),
            'edibles': set(),
            'resources': set(),
            'placeables': set()
        }

        # Split observation text into segments
        sections = [s.strip() for s in text_obs.split('\n') if s.strip()]

        # Parse environment objects
        if sections and sections[0].startswith('You see:'):
            seen_items = sections[0][8:].strip().split(',')
            for item in seen_items:
                item = item.strip()
                if item:
                    obs['environment'].add(item)
                    if item in ['zombie', 'skeleton']:
                        obs['dangers'].add(item)
                    elif item in ['cow', 'ripe plant']:
                        obs['edibles'].add(item)
                    elif item in ['tree', 'stone', 'coal', 'iron', 'diamond', 'water', 'sapling', 'grass']:
                        if item == 'tree':
                            item = 'wood'
                        obs['resources'].add(item)
                    elif item in ['crafting table', 'furnace', 'plant', 'placed_stone']:
                        obs['placeables'].add(item)

        # Parse inventory
        if len(sections) > 1 and sections[1].startswith('Inventory:'):
            inv_items = sections[1][10:].strip()
            if inv_items:
                for item_pair in inv_items.split(','):
                    if ':' in item_pair:
                        name, val = item_pair.split(':', 1)
                        self.achievements[name] = True
                        obs['inventory'][name.strip()] = int(val.strip())

        # Parse state
        if len(sections) > 2 and sections[2].startswith('Status:'):
            status_items = sections[2][7:].strip().split(',')
            for stat in status_items:
                if ':' in stat:
                    name, val = stat.split(':', 1)
                    name = name.strip().lower()
                    if '%' in val:
                        val = val.replace('%', '')
                    obs['status'][name] = int(val.strip())

        return obs

    def _evaluate_needs(self, status):
        """Dynamic need evaluation"""
        need_scores = []
        for need, config in self.need_config.items():
            current = status.get(need, 100)
            if current < config['threshold']:
                severity = (config['threshold'] - current) / config['threshold']
                need_scores.append({
                    'type': need,
                    'name': config['name'],
                    'score': severity * config['weight'],
                    'value': current
                })
        return sorted(need_scores, key=lambda x: -x['score'])

    def _generate_survival_goals(self, needs, obs):
        """Survival-oriented action generation"""
        goals = []

        # Define need-handling logic
        need_handlers = {
            'health': {
                'condition': lambda need, obs: obs['dangers'],
                'goal': lambda need, obs: {
                    'goal': 'Avoid enemies to recover health',
                    'priority': 1.2 + need['score'],
                    'type': 'safety'
                },
                'else_goal': lambda need, obs: {
                    'goal': 'Rest to recover health',
                    'priority': 1.1 + need['score'],
                    'type': 'health'
                }
            },
            'fullness': {
                'condition': lambda need, obs: 'ripe plant' in obs['edibles'],
                'goal': lambda need, obs: {
                    'goal': 'Harvest ripe plant for food',
                    'priority': 1.2 + need['score'],
                    'type': 'food',
                    'achievement': 'eat_plant'
                },
                'else_condition': lambda need, obs: 'cow' in obs['edibles'] and obs['inventory'].get('wood sword', 0) > 0,
                'else_goal': lambda need, obs: {
                    'goal': 'Hunt cows for meat',
                    'priority': 1.2 + need['score'],
                    'type': 'food',
                    'achievement': 'eat_cow'
                },
                'else_else_goal': lambda need, obs: {
                    'goal': 'Find cows' if obs['inventory'].get('wood_sword', 0) > 0 else 'Find food sources',
                    'priority': 1.2 + need['score'],
                    'type': 'food'
                }
            },
            'hydration': {
                'condition': lambda need, obs: 'water' in obs['resources'],
                'goal': lambda need, obs: {
                    'goal': 'Drink water',
                    'priority': 1.1 + need['score'],
                    'type': 'water',
                    'achievement': 'collect_drink'
                },
                'else_goal': lambda need, obs: {
                    'goal': 'Find water',
                    'priority': 1.0 + need['score'],
                    'type': 'water'
                }
            },
            'wakefulness': {
                'goal': lambda need, obs: {
                    'goal': 'Sleep',
                    'priority': 1.0 + need['score'],
                    'type': 'sleep',
                    'achievement': 'wake_up'
                }
            }
        }

        # Iterate over needs and generate actions
        for need in needs:
            need_type = need['type']
            if need_type in need_handlers:
                handler = need_handlers[need_type]
                if 'condition' in handler and handler['condition'](need, obs):
                    goals.append(handler['goal'](need, obs))
                elif 'else_condition' in handler and handler['else_condition'](need, obs):
                    goals.append(handler['else_goal'](need, obs))
                elif 'else_else_goal' in handler:
                    goals.append(handler['else_else_goal'](need, obs))
                elif 'goal' in handler:
                    goals.append(handler['goal'](need, obs))

        return goals

    def _generate_progression_goals(self, obs):
        """Development-oriented action generation"""
        goals = []
        inv = obs['inventory']
        env = obs['environment']

        # Crafting-bench actions
        has_table = 'crafting table' in env
        wood_count = inv.get('wood', 0)

        if not has_table and wood_count >= 2:
            goals.append({
                'goal': 'Place crafting table',
                'priority': 1.5,
                'type': 'crafting',
                'achievement': 'place_table'
            })
        elif not has_table:
            goals.append({
                'goal': f'Collect wood for table',
                'priority': 1.3,
                'type': 'gather'
            })

        # Furnace-related actions
        has_furnace = 'furnace' in env
        stone_count = inv.get('stone', 0)

        if not has_furnace and stone_count >= 1 and has_table:
            goals.append({
                'goal': 'Place furnace',
                'priority': 1.4,
                'type': 'crafting',
                'achievement': 'place_furnace'
            })
        elif not has_furnace and stone_count < 1 and has_table and "wood pickaxe" in inv:
            goals.append({
                'goal': f'Collect stone for furnace',
                'priority': 1.4,
                'type': 'gather'
            })

        # Tool-crafting actions
        for tool, recipe in self.crafting_recipes.items():
            if inv.get(tool, 0) == 0:
                can_craft = all(inv.get(res, 0) >= amt for res, amt in recipe['requirements'].items())
                if can_craft and has_table:
                    if 'iron' in tool and has_furnace:
                        goals.append({
                            'goal': f'Craft {tool}',
                            'priority': recipe['priority'],
                            'type': 'crafting',
                            'achievement': recipe['achievement']
                        })
                    elif 'iron' not in tool:
                        goals.append({
                            'goal': f'Craft {tool}',
                            'priority': recipe['priority'],
                            'type': 'crafting',
                            'achievement': recipe['achievement']
                        })
                else:
                    # Add sub-goals for obtaining required materials
                    for res, amt in recipe['requirements'].items():
                        if inv.get(res, 0) < amt and (self.resource_info[res]['tool'] is None or self.resource_info[res]['tool'] in inv):
                            goals.append({
                                'goal': f'Collect {res} for {tool}',
                                'priority': recipe['priority'],
                                'type': 'prepare'
                            })

        # Resource-collection actions
        for resource in obs['resources']:
            if resource in self.resource_info and (self.resource_info[resource]['tool'] is None or self.resource_info[resource]['tool'] in inv):
                info = self.resource_info[resource]
                has_tool = info['tool'] is None or inv.get(info['tool'], 0) > 0

                if has_tool:
                    goal = {
                        'goal': f'Collect {resource}',
                        'priority': info['priority'],
                        'type': 'gather',
                        'achievement': info.get('achievement', None)
                    }
                    goals.append(goal)
                else:
                    goals.append({
                        'goal': f'Acquire {info["tool"]} to collect {resource}',
                        'priority': info['priority'] * 0.8,
                        'type': 'prepare'
                    })

        # Item-placement actions
        for item, info in self.placement_info.items():
            if f'placed_{item}' not in env and item != 'table':  # 已经放置的不再考虑
                has_resources = all(inv.get(res, 0) >= amt for res, amt in info['requirements'].items())
                if has_resources:
                    goals.append({
                        'goal': f'Place {item}',
                        'priority': info['priority'],
                        'type': 'placement',
                        'achievement': info['achievement']
                    })
                else:
                    for res, amt in info['requirements'].items():
                        if inv.get(res, 0) < amt and (self.resource_info[res]['tool'] is None or self.resource_info[res]['tool'] in inv):
                            goals.append({
                                'goal': f'Collect {res} for placing {item}',
                                'priority': info['priority'] * 0.7,
                                'type': 'prepare'
                            })

        # Combat-related achievements
        if 'zombie' in obs['dangers']:
            if inv.get('wood sword', 0) > 0 or inv.get('stone sword', 0) > 0 or inv.get('iron sword', 0) > 0:
                goals.append({
                    'goal': 'Defeat zombie',
                    'priority': 1.7,
                    'type': 'combat',
                    'achievement': 'defeat_zombie'
                })
            else:
                goals.append({
                    'goal': 'Avoid zombie until better armed',
                    'priority': 2,
                    'type': 'safety'
                })

        if 'skeleton' in obs['dangers']:
            if inv.get('wood sword', 0) > 0 or inv.get('stone sword', 0) > 0 or inv.get('iron sword', 0) > 0:
                goals.append({
                    'goal': 'Defeat skeleton',
                    'priority': 1.7,
                    'type': 'combat',
                    'achievement': 'defeat_skeleton'
                })
            else:
                goals.append({
                    'goal': 'Avoid skeleton until better armed',
                    'priority': 2,
                    'type': 'safety'
                })

        # Plant crops
        if 'sapling' in inv and 'plant' not in env:
            goals.append({
                'goal': 'Place plant',
                'priority': 1.5 + inv.get('sapling', 0) * 0.1,
                'type': 'farming',
                'achievement': 'place_plant'
            })

        return goals

    def determine_goal(self, text_obs):
        self.all_goals = None

        """Core goal-decision method"""
        obs = self._parse_observation(text_obs)
        needs = self._evaluate_needs(obs['status'])

        # Generate all possible actions
        survival_goals = self._generate_survival_goals(needs, obs)
        progression_goals = self._generate_progression_goals(obs)

        # Merge and filter actions
        self.all_goals = survival_goals + progression_goals
        # self.all_goals = progression_goals

        # Choose the best action
        self.all_goals.append({
                'goal': 'Explore the environment',
                'priority': 0.5,
                'type': 'explore'
            })
        best_goal = max(self.all_goals, key=lambda x: x['priority'])

        return best_goal

    def get_augmentation_more_goal_text(self, text_obs):
        _ = self.determine_goal(text_obs)
        all_goals = [
            {
                'goal': 'Explore the environment',
                'priority': 1,
                'type': 'explore'
            },
            {
                'goal': 'Collect wood for basic tools',
                'priority': 0.8,
                'type': 'gather'
            },
            {
                'goal': 'Place crafting table',
                'priority': 0.5,
                'type': 'crafting'
            }
        ]
        for i in range(3-len(self.all_goals)):
            self.all_goals.append(all_goals[i])
        top_three_goals = sorted(self.all_goals, key=lambda x: x['priority'], reverse=True)[:3]

        priority = top_three_goals[0]['priority'] + top_three_goals[1]['priority'] + top_three_goals[2]['priority']
        top_three_goals[0]['priority'] /= priority
        top_three_goals[1]['priority'] /= priority
        top_three_goals[2]['priority'] /= priority
        return top_three_goals, self.all_goals
