from __future__ import annotations

import math
from dataclasses import dataclass

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


@dataclass
class Enemy:
    x: float
    y: float
    hp: float
    max_hp: float
    speed: float
    radius: float
    xp: float
    damage: float
    cooldown: float
    timer: float


@dataclass
class Shard:
    x: float
    y: float
    value: float


class TombfallEnv(gym.Env):
    metadata = {"render_modes": ["ansi"], "render_fps": 10}

    def __init__(self, render_mode: str | None = None, max_seconds: int = 360):
        super().__init__()
        self.render_mode = render_mode
        self.dt = 0.1
        self.max_steps = int(max_seconds / self.dt)
        self.world_limit = 1850.0
        self.action_space = spaces.MultiDiscrete([9, 3])
        self.observation_space = spaces.Box(low=-1.0, high=1.0, shape=(32,), dtype=np.float32)
        self.reset()

    def reset(self, seed: int | None = None, options: dict | None = None):
        super().reset(seed=seed)
        self.step_count = 0
        self.elapsed = 0.0
        self.wave = 1
        self.kills = 0
        self.score = 0.0
        self.xp = 0.0
        self.level = 1
        self.xp_next = 18.0
        self.player_x = 0.0
        self.player_y = 0.0
        self.hp = 118.0
        self.max_hp = 118.0
        self.shield = 38.0
        self.max_shield = 38.0
        self.fire_timer = 0.0
        self.melee_timer = 0.0
        self.dash_timer = 0.0
        self.spawn_timer = 0.0
        self.enemies: list[Enemy] = []
        self.shards: list[Shard] = []
        for _ in range(8):
            self._spawn_enemy()
        return self._observe(), self._info()

    def step(self, action):
        direction, skill = int(action[0]), int(action[1])
        move_x, move_y = self._direction(direction)
        reward = -0.002
        hp_before = self.hp + self.shield
        kills_before = self.kills
        xp_before = self.xp

        speed = 232.0
        if skill == 2 and self.dash_timer <= 0:
            speed = 640.0
            self.dash_timer = 3.2
            reward -= 0.01

        self.player_x = float(np.clip(self.player_x + move_x * speed * self.dt, -self.world_limit, self.world_limit))
        self.player_y = float(np.clip(self.player_y + move_y * speed * self.dt, -self.world_limit, self.world_limit))

        if skill == 1 and self.melee_timer <= 0:
            reward += self._melee_attack()
            self.melee_timer = 0.76

        self._auto_fire()
        self._update_enemies()
        self._collect_shards()
        self._spawn_routine()
        self._regen()

        self.elapsed += self.dt
        self.step_count += 1
        self.wave = min(6, int(self.elapsed // 60) + 1)
        self.fire_timer = max(0.0, self.fire_timer - self.dt)
        self.melee_timer = max(0.0, self.melee_timer - self.dt)
        self.dash_timer = max(0.0, self.dash_timer - self.dt)

        reward += (self.kills - kills_before) * 1.0
        reward += max(0.0, self.xp - xp_before) * 0.025
        reward -= max(0.0, hp_before - (self.hp + self.shield)) * 0.015

        terminated = self.hp <= 0
        truncated = self.step_count >= self.max_steps
        if truncated and not terminated:
            reward += 10.0
        return self._observe(), float(reward), terminated, truncated, self._info()

    def render(self):
        lines = [
            f"time={self.elapsed:05.1f}s wave={self.wave} kills={self.kills} score={self.score:.0f}",
            f"player hp={self.hp:.1f}/{self.max_hp:.0f} shield={self.shield:.1f}/{self.max_shield:.0f}",
            f"enemies={len(self.enemies)} shards={len(self.shards)} level={self.level}",
        ]
        return "\n".join(lines)

    def _observe(self):
        values = [
            self._scale(self.elapsed, 0, 360),
            self._scale(self.hp, 0, self.max_hp),
            self._scale(self.shield, 0, self.max_shield),
            self._scale(self.level, 1, 20),
            self._scale(len(self.enemies), 0, 40),
            self._scale(len(self.shards), 0, 30),
            self._scale(self.player_x, -self.world_limit, self.world_limit),
            self._scale(self.player_y, -self.world_limit, self.world_limit),
        ]

        for enemy in self._nearest(self.enemies, 4):
            values.extend(
                [
                    self._scale(enemy.x - self.player_x, -900, 900),
                    self._scale(enemy.y - self.player_y, -900, 900),
                    self._scale(enemy.hp, 0, enemy.max_hp),
                ]
            )

        while len(values) < 20:
            values.append(0.0)

        for shard in self._nearest(self.shards, 4):
            values.extend(
                [
                    self._scale(shard.x - self.player_x, -900, 900),
                    self._scale(shard.y - self.player_y, -900, 900),
                    self._scale(shard.value, 0, 45),
                ]
            )

        while len(values) < 32:
            values.append(0.0)
        return np.array(values[:32], dtype=np.float32)

    def _info(self):
        return {
            "elapsed": round(self.elapsed, 2),
            "wave": self.wave,
            "kills": self.kills,
            "score": round(self.score, 2),
            "hp": round(self.hp, 2),
            "shield": round(self.shield, 2),
            "enemies": len(self.enemies),
            "shards": len(self.shards),
        }

    def _direction(self, action: int):
        directions = [
            (0, 0),
            (0, -1),
            (0, 1),
            (-1, 0),
            (1, 0),
            (-1, -1),
            (1, -1),
            (-1, 1),
            (1, 1),
        ]
        x, y = directions[action] if 0 <= action < len(directions) else (0, 0)
        length = math.hypot(x, y) or 1.0
        return x / length, y / length

    def _auto_fire(self):
        if self.fire_timer > 0 or not self.enemies:
            return
        target = self._nearest(self.enemies, 1)[0]
        target.hp -= 18.0
        self.fire_timer = 1.0 / 3.55
        if target.hp <= 0:
            self._kill_enemy(target)

    def _melee_attack(self):
        reward = 0.0
        for enemy in list(self.enemies):
            if self._distance(enemy.x, enemy.y, self.player_x, self.player_y) < enemy.radius + 112:
                enemy.hp -= 58.0
                reward += 0.05
                if enemy.hp <= 0:
                    self._kill_enemy(enemy)
        return reward

    def _update_enemies(self):
        for enemy in list(self.enemies):
            dx = self.player_x - enemy.x
            dy = self.player_y - enemy.y
            dist = math.hypot(dx, dy) or 1.0
            enemy.x += dx / dist * enemy.speed * self.dt
            enemy.y += dy / dist * enemy.speed * self.dt
            enemy.timer -= self.dt
            if dist < enemy.radius + 28 and enemy.timer <= 0:
                self._damage_player(enemy.damage)
                enemy.timer = enemy.cooldown

    def _collect_shards(self):
        for shard in list(self.shards):
            dist = self._distance(shard.x, shard.y, self.player_x, self.player_y)
            if dist < 94:
                self.xp += shard.value
                self.score += shard.value * 0.1
                self.shards.remove(shard)
                if self.xp >= self.xp_next:
                    self.xp -= self.xp_next
                    self.level += 1
                    self.xp_next = self.xp_next * 1.24 + 8

    def _spawn_routine(self):
        self.spawn_timer -= self.dt
        if self.spawn_timer > 0:
            return
        self.spawn_timer = max(0.18, 1.05 - self.wave * 0.09 - self.elapsed * 0.0012)
        count = 2 if self.wave >= 2 else 1
        if self.wave >= 5:
            count = 3
        for _ in range(count):
            self._spawn_enemy()

    def _spawn_enemy(self):
        side = int(self.np_random.integers(0, 4))
        offset = 760.0
        if side == 0:
            x, y = self.player_x - offset, self.player_y + float(self.np_random.uniform(-420, 420))
        elif side == 1:
            x, y = self.player_x + offset, self.player_y + float(self.np_random.uniform(-420, 420))
        elif side == 2:
            x, y = self.player_x + float(self.np_random.uniform(-420, 420)), self.player_y - offset
        else:
            x, y = self.player_x + float(self.np_random.uniform(-420, 420)), self.player_y + offset

        scale = 1.0 + self.wave * 0.11 + self.elapsed * 0.0008
        base_hp = 28.0 if self.wave < 4 else 42.0
        self.enemies.append(
            Enemy(
                x=x,
                y=y,
                hp=base_hp * scale,
                max_hp=base_hp * scale,
                speed=92.0 + self.wave * 4,
                radius=13.0,
                xp=3.0 + self.wave,
                damage=10.0 + self.wave,
                cooldown=0.7,
                timer=float(self.np_random.uniform(0.2, 0.9)),
            )
        )

    def _kill_enemy(self, enemy: Enemy):
        if enemy in self.enemies:
            self.enemies.remove(enemy)
        self.kills += 1
        self.score += 1.0
        self.shards.append(Shard(enemy.x, enemy.y, enemy.xp))

    def _damage_player(self, amount: float):
        if self.shield > 0:
            taken = min(self.shield, amount)
            self.shield -= taken
            amount -= taken
        self.hp -= max(0.0, amount)

    def _regen(self):
        self.hp = min(self.max_hp, self.hp + 0.65 * self.dt)
        self.shield = min(self.max_shield, self.shield + 7.5 * self.dt)

    def _nearest(self, items, limit: int):
        return sorted(items, key=lambda item: self._distance(item.x, item.y, self.player_x, self.player_y))[:limit]

    @staticmethod
    def _distance(ax, ay, bx, by):
        return math.hypot(ax - bx, ay - by)

    @staticmethod
    def _scale(value, low, high):
        if high == low:
            return 0.0
        scaled = (float(value) - low) / (high - low) * 2.0 - 1.0
        return float(np.clip(scaled, -1.0, 1.0))