from __future__ import annotations

from collections import defaultdict

import numpy as np
import wgpu

from engine.camera import OrthoCamera
from engine.character import Character
from engine.defs import DefRegistry
from engine.map import Map, rotate_point
from engine.render.pipeline import QuadPipeline
from engine.render.texture import TextureCache
from engine.tile import TILE_LAYERS
from engine.world import World

CLEAR_COLOR = {"r": 0.05, "g": 0.05, "b": 0.08, "a": 1.0}
CHARACTER_SLOT_ORDER = ("head", "torso", "left_arm", "right_arm", "left_leg", "right_leg")


class Renderer:
    def __init__(self, device, pipeline: QuadPipeline, texture_cache: TextureCache, registry: DefRegistry):
        self.device = device
        self.pipeline = pipeline
        self.texture_cache = texture_cache
        self.registry = registry
        self._texture_bind_groups: dict[str, object] = {}

    def _bind_group_for(self, relative_path: str):
        if relative_path not in self._texture_bind_groups:
            view = self.texture_cache.get_view(relative_path)
            self._texture_bind_groups[relative_path] = self.pipeline.make_texture_bind_group(view)
        return self._texture_bind_groups[relative_path]

    def render_frame(self, context, world: World, camera: OrthoCamera) -> None:
        scale, offset = camera.scale_offset()
        self.pipeline.update_camera(scale, offset)

        buckets: dict[str, list[float]] = defaultdict(list)
        if world.player is not None and world.player.position is not None:
            active_map = world.maps[world.player.position.map_id]
            self._draw_map(active_map, (0, 0, 0), 0, world, buckets)

        texture = context.get_current_texture()
        view = texture.create_view()
        encoder = self.device.create_command_encoder()
        pass_encoder = encoder.begin_render_pass(
            color_attachments=[
                {
                    "view": view,
                    "clear_value": CLEAR_COLOR,
                    "load_op": wgpu.LoadOp.clear,
                    "store_op": wgpu.StoreOp.store,
                }
            ]
        )
        pass_encoder.set_pipeline(self.pipeline.pipeline)
        pass_encoder.set_bind_group(0, self.pipeline.camera_bind_group)
        pass_encoder.set_vertex_buffer(0, self.pipeline.quad_vertex_buffer)
        pass_encoder.set_index_buffer(self.pipeline.quad_index_buffer, wgpu.IndexFormat.uint16)

        for relative_path, flat_instances in buckets.items():
            instances = np.array(flat_instances, dtype=np.float32).reshape(-1, 4)
            instance_buffer = self.pipeline.make_instance_buffer(instances)
            pass_encoder.set_bind_group(1, self._bind_group_for(relative_path))
            pass_encoder.set_vertex_buffer(1, instance_buffer)
            pass_encoder.draw_indexed(self.pipeline.quad_index_count, len(instances))

        pass_encoder.end()
        self.device.queue.submit([encoder.finish()])

    def _draw_map(
        self, m: Map, offset: tuple[int, int, int], rotation: int, world: World, buckets: dict[str, list[float]]
    ) -> None:
        for z in range(m.depth):
            for layer in TILE_LAYERS:
                for y in range(m.height):
                    for x in range(m.width):
                        tile = m.get_tile(x, y, z)
                        def_id = tile.layer_id(layer)
                        if def_id is None:
                            continue
                        layer_def = self.registry.get_tile_layer_def(layer, def_id)
                        if layer_def is None or not layer_def.texture:
                            continue
                        rx, ry = rotate_point(x, y, m.width, m.height, rotation)
                        buckets[layer_def.texture].extend(
                            (float(offset[0] + rx), float(offset[1] + ry), 1.0, 1.0)
                        )

        for entity in world.entities_on_map(m.map_id):
            if isinstance(entity, Character):
                self._draw_character(entity, offset, rotation, m.width, m.height, buckets)

        for embedding in m.embeddings:
            arx, ary = rotate_point(embedding.anchor[0], embedding.anchor[1], m.width, m.height, rotation)
            child_offset = (offset[0] + arx, offset[1] + ary, offset[2] + embedding.anchor[2])
            child_rotation = (rotation + embedding.rotation) % 4
            self._draw_map(embedding.child_map, child_offset, child_rotation, world, buckets)

    def _draw_character(
        self,
        character: Character,
        offset: tuple[int, int, int],
        rotation: int,
        map_width: int,
        map_height: int,
        buckets: dict[str, list[float]],
    ) -> None:
        assert character.position is not None
        rx, ry = rotate_point(character.position.x, character.position.y, map_width, map_height, rotation)
        base_x = offset[0] + rx
        base_y = offset[1] + ry
        resolved = character.resolved_body_parts()
        row_h = 1.0 / len(CHARACTER_SLOT_ORDER)
        for i, slot in enumerate(CHARACTER_SLOT_ORDER):
            part = resolved.get(slot)
            if part is None:
                continue
            part_def = self.registry.get_body_part_def(part.part_def_id)
            if not part_def.texture:
                continue
            buckets[part_def.texture].extend(
                (base_x + 0.1, base_y + i * row_h + row_h * 0.05, 0.8, row_h * 0.9)
            )