from agentState import AgentState
from typing import Dict, Tuple
from city import City
from gridCellType import GridCellType
from agentActionType import AgentActionType
from agentOrientation import AgentOrientation
from queue import Queue, PriorityQueue
from turnCar import turn_left_orientation, turn_right_orientation

class Succ:
    state: AgentState
    action: AgentActionType
    cost: int
    predicted_cost: int

    def __init__(self, state: AgentState, action: AgentActionType, cost: int, predicted_cost: int) -> None:
        self.state = state
        self.action = action
        self.cost = cost
        self.predicted_cost = cost

class SuccList:
    succ_list: list[Succ]

    def __init__(self, succ_list: list[Succ]) -> None:
        self.succ_list = succ_list

    def __lt__(self, other):
        return self.succ_list[-1].predicted_cost < other.succ_list[-1].predicted_cost
    
    def __gt__(self, other):
        return self.succ_list[-1].predicted_cost > other.succ_list[-1].predicted_cost

def find_path_to_nearest_can(startState: AgentState, grid: Dict[Tuple[int, int], GridCellType], city: City) -> list[AgentActionType]:
    q: PriorityQueue[SuccList] = PriorityQueue()
    visited: list[AgentState] = []
    startStates: SuccList = SuccList([Succ(startState, AgentActionType.UNKNOWN, 0, _heuristics(startState.position, city))])
    q.put(startStates)
    while not q.empty():
        currently_checked = q.get()
        visited.append(currently_checked.succ_list[-1].state)
        if is_state_success(currently_checked.succ_list[-1].state, grid):
            return extract_actions(currently_checked)
        successors = succ(currently_checked.succ_list[-1], grid, city)
        for s in successors:
            already_visited = False
            for v in visited:
                if v.position[0] == s.state.position[0] and v.position[1] == s.state.position[1] and s.state.orientation == v.orientation:
                    already_visited = True
                    break
            if already_visited:
                continue
            if is_state_valid(s.state, grid):
                new_list = currently_checked.succ_list.copy()
                new_list.append(s)
                q.put(SuccList(new_list))
    return []

        
        
def extract_actions(successors: SuccList) -> list[AgentActionType]:
    output: list[AgentActionType] = []
    for s in successors.succ_list:
        if s.action != AgentActionType.UNKNOWN:
            output.append(s.action)
    return output

def succ(succ: Succ, grid: Dict[Tuple[int, int], GridCellType], city: City) -> list[Succ]:
    result: list[Succ] = []
    turn_left_cost = 1 + succ.cost
    result.append(Succ(AgentState(succ.state.position, turn_left_orientation(succ.state.orientation)), AgentActionType.TURN_LEFT, turn_left_cost, turn_left_cost + _heuristics(succ.state.position, city)))
    turn_right_cost = 1 + succ.cost
    result.append(Succ(AgentState(succ.state.position, turn_right_orientation(succ.state.orientation)), AgentActionType.TURN_RIGHT, turn_right_cost, turn_right_cost + _heuristics(succ.state.position, city)))
    state_succ = move_forward_succ(succ, city, grid)
    if state_succ != None:
        result.append(state_succ)
    return result

def move_forward_succ(succ: Succ, city: City, grid: Dict[Tuple[int, int], GridCellType]) -> Succ:
    position = get_next_cell(succ.state)
    if position == None:
        return None
    
    cost = get_cost_for_action(AgentActionType.MOVE_FORWARD, grid[position]) + succ.cost
    return Succ(AgentState(position, succ.state.orientation), AgentActionType.MOVE_FORWARD, cost, cost + _heuristics(position, city))


def get_next_cell(state: AgentState) -> Tuple[int, int]:
    if state.orientation == AgentOrientation.UP:
        if state.position[1] - 1 < 1:
            return None
        return (state.position[0], state.position[1] - 1)
    if state.orientation == AgentOrientation.DOWN:
        if state.position[1] + 1 > 27:
            return None
        return (state.position[0], state.position[1] + 1)
    if state.orientation == AgentOrientation.LEFT:
        if state.position[0] - 1 < 1:
            return None
        return (state.position[0] - 1, state.position[1])
    if state.position[0] + 1 > 27:
        return None
    return (state.position[0] + 1, state.position[1])

def is_state_success(state: AgentState, grid: Dict[Tuple[int, int], GridCellType]) -> bool:
    next_cell = get_next_cell(state)
    try:
        return grid[next_cell] == GridCellType.GARBAGE_CAN
    except:
        return False

def get_cost_for_action(action: AgentActionType, cell_type: GridCellType) -> int:
    if action == AgentActionType.TURN_LEFT or action == AgentActionType.TURN_RIGHT:
        return 1
    if cell_type == GridCellType.SPEED_BUMP:
        if action == AgentActionType.MOVE_FORWARD:
            return 10
    if action == AgentActionType.MOVE_FORWARD:
        return 3


def is_state_valid(state: AgentState, grid: Dict[Tuple[int, int], GridCellType]) -> bool:
    try: 
        return grid[state.position] == GridCellType.STREET_HORIZONTAL or grid[state.position] == GridCellType.STREET_VERTICAL or grid[state.position] == GridCellType.SPEED_BUMP
    except:
        return False
    
def _heuristics(position: Tuple[int, int], city: City):
    min_distance: int = 300
    found_nonvisited: bool = False
    for can in city.cans:
        if can.is_visited:
            continue
        found_nonvisited = True
        distance = 3 * (abs(position[0] - can.position[0]) + abs(position[1] - can.position[1]))
        if distance < min_distance:
            min_distance = distance
    if found_nonvisited:
        return min_distance
    return -1