import random
from collections import deque

import gradio as gr
import gymnasium as gym
import numpy as np
from gymnasium import spaces


class WarehouseGridWorldEnv(gym.Env):
    metadata = {"render_modes": ["ansi"]}

    def __init__(self, width=10, height=10, obstacle_density=0.20, max_steps=150):
        super().__init__()
        self.width = width
        self.height = height
        self.obstacle_density = obstacle_density
        self.max_steps = max_steps

        self.action_space = spaces.Discrete(4)  # 0=UP, 1=RIGHT, 2=DOWN, 3=LEFT
        self.observation_space = spaces.Box(low=0.0, high=1.0, shape=(4,), dtype=np.float32)

        self.reset()

    def _obs(self):
        ax, ay = self.agent_pos
        gx, gy = self.goal_pos
        return np.array([
            ax / (self.width - 1),
            ay / (self.height - 1),
            gx / (self.width - 1),
            gy / (self.height - 1),
        ], dtype=np.float32)

    def _manhattan(self, a=None, b=None):
        a = self.agent_pos if a is None else a
        b = self.goal_pos if b is None else b
        return abs(a[0] - b[0]) + abs(a[1] - b[1])

    def _neighbors(self, cell):
        x, y = cell
        for dx, dy in [(0, -1), (1, 0), (0, 1), (-1, 0)]:
            nx, ny = x + dx, y + dy
            if 0 <= nx < self.width and 0 <= ny < self.height and (nx, ny) not in self.obstacles:
                yield (nx, ny)

    def _solvable(self, start, goal):
        q = deque([start])
        seen = {start}
        while q:
            cell = q.popleft()
            if cell == goal:
                return True
            for n in self._neighbors(cell):
                if n not in seen:
                    seen.add(n)
                    q.append(n)
        return False

    def _generate_solvable_maze(self):
        cells = [(x, y) for y in range(self.height) for x in range(self.width)]
        obstacle_count = int(len(cells) * self.obstacle_density)

        for _ in range(500):
            self.start_pos, self.goal_pos = random.sample(cells, 2)
            available = [c for c in cells if c not in {self.start_pos, self.goal_pos}]
            self.obstacles = set(random.sample(available, obstacle_count))
            if self._solvable(self.start_pos, self.goal_pos):
                return

        # Safe fallback: mostly empty maze if random attempts fail.
        self.start_pos, self.goal_pos = (0, 0), (self.width - 1, self.height - 1)
        self.obstacles = set()

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        self._generate_solvable_maze()
        self.agent_pos = self.start_pos
        self.visited = {self.agent_pos}
        self.steps = 0
        self.total_score = 0.0
        self.last_reward = 0.0
        self.reached_goal = False
        return self._obs(), self._info()

    def step(self, action):
        if self.reached_goal:
            return self._obs(), 0.0, True, False, self._info()

        moves = {
            0: (0, -1),  # UP
            1: (1, 0),   # RIGHT
            2: (0, 1),   # DOWN
            3: (-1, 0),  # LEFT
        }
        dx, dy = moves[int(action)]
        old_pos = self.agent_pos
        old_dist = self._manhattan()
        new_pos = (old_pos[0] + dx, old_pos[1] + dy)

        reward = -0.02  # small step cost
        terminated = False
        truncated = False

        outside = not (0 <= new_pos[0] < self.width and 0 <= new_pos[1] < self.height)
        blocked = new_pos in self.obstacles

        if outside or blocked:
            reward -= 1.0
            new_pos = old_pos
        else:
            self.agent_pos = new_pos
            new_dist = self._manhattan()
            if new_dist < old_dist:
                reward += 0.30
            elif new_dist > old_dist:
                reward -= 0.20
            else:
                reward -= 0.05

            if new_pos not in self.visited:
                reward += 0.10
                self.visited.add(new_pos)

            if new_pos == self.goal_pos:
                reward += 10.0
                self.reached_goal = True
                terminated = True

        self.steps += 1
        if self.steps >= self.max_steps and not terminated:
            reward -= 3.0
            truncated = True

        self.last_reward = float(reward)
        self.total_score += float(reward)
        return self._obs(), float(reward), terminated, truncated, self._info()

    def _info(self):
        return {
            "total_score": round(self.total_score, 2),
            "last_reward": round(self.last_reward, 2),
            "steps": self.steps,
            "agent_position": self.agent_pos,
            "goal_position": self.goal_pos,
            "manhattan_distance": self._manhattan(),
            "goal_reached": self.reached_goal,
        }

    def render_html(self):
        rows = []
        for y in range(self.height):
            cells = []
            for x in range(self.width):
                pos = (x, y)
                label = "."
                cls = "empty"
                inner = "."
                if pos in self.obstacles:
                    label, cls, inner = "X", "obstacle", "X"
                if pos == self.start_pos:
                    label, cls, inner = "S", "start", "S"
                if pos == self.goal_pos:
                    label, cls, inner = "G", "goal", "G"
                if pos == self.agent_pos:
                    cls += " agent-cell"
                    inner = f"<span class='agent'>A</span>"
                cells.append(f"<td class='{cls}' title='{label}'>{inner}</td>")
            rows.append("<tr>" + "".join(cells) + "</tr>")
        return "<table class='grid'>" + "".join(rows) + "</table>"


env = WarehouseGridWorldEnv()


def scoreboard():
    info = env._info()
    return f"""
### Scoreboard
- **Total score:** {info['total_score']}
- **Last reward:** {info['last_reward']}
- **Steps:** {info['steps']} / {env.max_steps}
- **Agent position:** {info['agent_position']}
- **Goal position:** {info['goal_position']}
- **Manhattan distance:** {info['manhattan_distance']}
- **Goal reached:** {info['goal_reached']}
"""


def update(action=None):
    if action is not None:
        env.step(action)
    return env.render_html(), scoreboard()


def reset_game():
    env.reset()
    return env.render_html(), scoreboard()


CSS = """
.grid { border-collapse: collapse; margin: 8px 0; }
.grid td {
    width: 42px; height: 42px; text-align: center; vertical-align: middle;
    border: 1px solid #888; font-family: monospace; font-size: 20px; font-weight: 700;
}
.empty { background: #f7f7f7; color: #999; }
.start { background: #1976d2; color: white; }
.goal { background: #2e7d32; color: white; }
.obstacle { background: #222; color: white; }
.agent-cell { position: relative; }
.agent {
    display: inline-flex; align-items: center; justify-content: center;
    width: 30px; height: 30px; border-radius: 50%;
    background: #d32f2f; color: white; font-size: 16px; font-family: Arial, sans-serif;
}
#hidden_action { display: none; }
"""

JS = """
<script>
document.addEventListener('keydown', function(e) {
  const map = {ArrowUp: 'up_btn', ArrowRight: 'right_btn', ArrowDown: 'down_btn', ArrowLeft: 'left_btn'};
  const id = map[e.key];
  if (id) {
    e.preventDefault();
    const btn = document.getElementById(id);
    if (btn) btn.click();
  }
});
</script>
"""

with gr.Blocks(css=CSS, title="Warehouse GridWorld") as demo:
    gr.Markdown("# Warehouse GridWorld Navigation")
    gr.HTML(JS)
    with gr.Row():
        grid = gr.HTML(env.render_html())
        board = gr.Markdown(scoreboard())

    with gr.Row():
        up = gr.Button("⬆ Up", elem_id="up_btn")
        right = gr.Button("➡ Right", elem_id="right_btn")
        down = gr.Button("⬇ Down", elem_id="down_btn")
        left = gr.Button("⬅ Left", elem_id="left_btn")
        reset = gr.Button("Reset / Randomize")

    up.click(lambda: update(0), outputs=[grid, board])
    right.click(lambda: update(1), outputs=[grid, board])
    down.click(lambda: update(2), outputs=[grid, board])
    left.click(lambda: update(3), outputs=[grid, board])
    reset.click(reset_game, outputs=[grid, board])

if __name__ == "__main__":
    demo.launch()