from langchain_core.output_parsers import BaseOutputParser
from orchestrator_maze_implementation.state.maze_state import MazeState
from typing import Dict, List, Tuple, Optional
import re

class MazeViewParser(BaseOutputParser):
    """Parse get_current_view() output into structured data"""
    
    def parse(self, text: str) -> Dict:
        lines = text.strip().split('\n')
        
        # Extract position
        position_match = re.search(r'Current position: \((\d+), (\d+)\)', text)
        position = (int(position_match.group(1)), int(position_match.group(2))) if position_match else None
        
        # Extract possible moves
        moves_match = re.search(r"Possible moves: \[(.*?)\]", text)
        possible_moves = []
        if moves_match:
            moves_str = moves_match.group(1).replace("'", "").replace('"', '')
            possible_moves = [move.strip() for move in moves_str.split(',') if move.strip()]
        
        # Extract at_exit status
        at_exit = "At exit: True" in text
        
        return {
            "position": position,
            "possible_moves": possible_moves,
            "at_exit": at_exit,
            "raw_view": text
        }

class MazeMoveParser(BaseOutputParser):
    """Parse movement tool output"""
    
    def parse(self, text: str) -> Dict:
        success = "Success" in text
        new_position = None
        
        if success:
            pos_match = re.search(r'New position: \((\d+), (\d+)\)', text)
            if pos_match:
                new_position = (int(pos_match.group(1)), int(pos_match.group(2)))
        
        return {
            "success": success,
            "new_position": new_position,
            "message": text
        }

class MazeStateExtractor:
    """Extract current maze state for efficient context provision"""
    
    @staticmethod
    def extract_current_state(state: MazeState, agent_id: str) -> Dict:
        """Extract only essential current state information"""
        maze_wrapper = state["maze_wrappers"][agent_id]
        
        return {
            "current_position": maze_wrapper.get_agent_position(),
            "possible_moves": maze_wrapper.get_possible_moves(),
            "at_exit": maze_wrapper.is_at_exit(),
            "visited_positions": list(maze_wrapper.move_history),
            "marked_dead_ends": list(maze_wrapper.get_marked_dead_ends()),
            "step_in_plan": state.get("step_index", 0),
            "current_step": state["plan"][state.get("step_index", 0)] if state.get("step_index", 0) < len(state["plan"]) else "Complete"
        }