from typing import Annotated, Dict, List, Any, Optional, Tuple
from typing_extensions import TypedDict
from pydantic import Field
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage
from orchestrator_maze_implementation.utils.maze_utils.maze_wrapper import MazeWrapper


def overwrite_dict_only_present_in_new(state1: dict, state2: dict) -> dict:
    """Overwrite the state1 with the state2, only if the key is present in the state2"""
    updated_state = dict(state1)
    #print(f"og state: {updated_state}", f"state2: {state2}")
    for key, value in state2.items():
        updated_state[key] = value
    #print(f"Updated state: {updated_state}")
    return updated_state

def overwrite_scalar_only_present_in_new(old_value: Any, new_value: Any) -> Any:
    """Overwrite old_value with new_value only if new_value is provided (not None)."""
    if new_value is not None:  # "Present" means not None in updates
        #print(f"og value: {old_value}", f"new_value: {new_value}")
        updated_value = new_value
        #print(f"Updated value: {updated_value}")
        return updated_value
    return old_value

def overwrite_list_only_present_in_new(old_list: Optional[List], new_list: Optional[List]) -> Optional[List]:
    """Overwrite old_list with new_list only if new_list is provided (not None)."""
    if new_list is not None:
        return new_list
    return old_list

class MazeState(TypedDict):
    """Holds state mapping agent ids to their maze wrappers"""
    maze_wrappers: Annotated[dict[str, MazeWrapper], overwrite_dict_only_present_in_new] #state management for agents in each maze_wrapper
    messages: List[AnyMessage] #top-level message list. Purpose tbd but langchain native -> könnte cache für orchestrator sein
    system_task: str
    all_agents: List[str]
    plan: List[str] #plan = list of subtasks divided into steps
    plan_completed: bool
    step_index: int
    turn_count: Annotated[int, overwrite_scalar_only_present_in_new]  # Counter for number of full turns completed

########################################################################################
# MAZE VICTORY CONDITIONS
########################################################################################

    maze_exit_found: bool 
    exit_position: List[str]
    winning_agent: Optional[str]
    turn_complete: Optional[bool]

########################################################################################
# AGENT STATE MANAGEMENT (Agent Memory)
########################################################################################

    current_agent: Annotated[str, overwrite_scalar_only_present_in_new]
    agent_messages: Annotated[Dict[str, List[AnyMessage]], overwrite_dict_only_present_in_new] # e.g., {"agent_1": [msgs]}
    known_openings: Annotated[Optional[Dict[str, List[str]]], overwrite_dict_only_present_in_new]     # Structure: {"agent_1": ["(1,2)-north", "(3,4)-east", "(1,2)-south"]} - unexplored directions
    # agent_positions: Dict[str, Tuple[int, int]]
    # previously_visited_tiles: Dict[str, List[Tuple[int, int]]]
    
    # Backtracking state management
    agent_backtracking_state: Annotated[Optional[Dict[str, Dict[str, Any]]], overwrite_dict_only_present_in_new]  # {agent_id: {is_backtracking: bool, target_position: Tuple, path: List}}

########################################################################################
# ORCHESTRATOR STATE MANAGEMENT
########################################################################################
    
    orchestrator_guidance: Optional[Dict[str, str]]
    shared_knowledge: Optional[Dict[str, Any]] #shared_knowledge.discovered_cells -> converts previously_visited_tiles per agent into one mapping object.

########################################################################################
# FREE ENERGY BENCHMARKING
########################################################################################

    free_energy_metrics: Annotated[Optional[Dict[str, Any]], overwrite_dict_only_present_in_new]  # Latest FE calculation results
    entropy_history: Annotated[Optional[List[Dict[str, Any]]], overwrite_list_only_present_in_new]  # Historical FE data
