use std::{fs::read, io::Cursor, num::NonZero, sync::LazyLock};

use anyhow::Context;
use bytemuck::{bytes_of, Pod, Zeroable};
use glam::{Mat4, Quat, Vec3};
use truth::{context::RootEmitter, io::BinReader, AnmFile, Game};
use wgpu::{
  include_wgsl, util::{BufferInitDescriptor, DeviceExt}, BindGroup, BindGroupDescriptor, BindGroupEntry, BindGroupLayoutDescriptor, BindGroupLayoutEntry, BlendState, Buffer, BufferUsages, Color, ColorTargetState, ColorWrites, Device, FragmentState, ImageCopyTexture, ImageDataLayout, LoadOp, Operations, Origin3d, PipelineLayoutDescriptor, PrimitiveState, Queue, RenderPassColorAttachment, RenderPassDescriptor, RenderPipeline, RenderPipelineDescriptor, ShaderModule, ShaderStages, SurfaceConfiguration, Texture, TextureAspect, TextureDescriptor, TextureUsages, VertexAttribute, VertexBufferLayout, VertexState
};
use winit::keyboard::KeyCode;

use crate::{
  anm::image::produce_image_from_entry, engine::{Engine, EngineState, UpdateContext}
};

pub struct AnmViewer {
  anm: AnmFile,
  vertices: Buffer,
  instances: Buffer,
  textures: Vec<Texture>,
  texture_bind_groups: Vec<BindGroup>,
  bind_group: BindGroup,
  pipeline: RenderPipeline,
  shader: ShaderModule,
  texture: u32,
  uniform_buffer: Buffer,
  uniform: Uniform,
}

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct Vertex {
  x: f32,
  y: f32,
  tx: f32,
  ty: f32,
}

const VERTICES: [Vertex; 4] = [
  Vertex {
    x: 0.,
    y: 0.,
    tx: 0.,
    ty: 0.,
  },
  Vertex {
    x: 1.,
    y: 0.,
    tx: 1.,
    ty: 0.,
  },
  Vertex {
    x: 0.,
    y: 1.,
    tx: 0.,
    ty: 1.,
  },
  Vertex {
    x: 1.,
    y: 1.,
    tx: 1.,
    ty: 1.,
  },
];

const INSTANCES: LazyLock<[Instance; 3]> = LazyLock::new(|| {
  [
    Instance {
      matrix: Mat4::from_scale(Vec3::new(256.0, 512.0, 1.0)),
    },
    Instance {
      matrix: Mat4::from_scale_rotation_translation(
        Vec3::new(256.0, 512.0, 1.0),
        Quat::IDENTITY,
        Vec3::new(0.0, 0.0, 0.0),
      ),
    },
    Instance {
      matrix: Mat4::from_scale_rotation_translation(
        Vec3::new(512.0, 256.0, 1.0),
        Quat::IDENTITY,
        Vec3::new(0.0, 0.0, 0.0),
      ),
    },
  ]
});

#[derive(Clone, Copy, Pod, Zeroable)]
#[repr(C)]
struct Instance {
  matrix: Mat4,
}

#[derive(Pod, Clone, Copy, Zeroable)]
#[repr(C, packed)]
struct Uniform {
  proj_matrix: Mat4,
}

impl AnmViewer {
  pub fn new(device: &Device, queue: &Queue, config: &SurfaceConfiguration) -> Self {
    let anm = AnmFile::read_from_stream(
      &mut BinReader::from_reader(
        &RootEmitter::new_stderr(),
        "",
        Cursor::new(read("assets/title.anm").unwrap()),
      ),
      Game::Th11,
      true,
    )
    .unwrap();

    let mut textures = vec![];
    for entry in &anm.entries {
      println!(
        "entry {:?} {:?} {:?} {:?}",
        entry.img_width(),
        entry.img_height(),
        entry.img_format(),
        entry.path
      );
      if entry.img_data().is_none() {
        continue;
      }
      let image = produce_image_from_entry(&entry).with_context(|| format!("at {}", entry.path)).unwrap();

      let texture_size = wgpu::Extent3d {
        width: image.width as u32,
        height: image.height as u32,
        depth_or_array_layers: 1,
      };
      let texture = device.create_texture(&TextureDescriptor {
        format: wgpu::TextureFormat::Bgra8Unorm,
        dimension: wgpu::TextureDimension::D2,
        size: texture_size,
        mip_level_count: 1,
        sample_count: 1,
        usage: TextureUsages::TEXTURE_BINDING | TextureUsages::COPY_DST,
        label: None,
        view_formats: &[],
      });

      queue.write_texture(
        ImageCopyTexture {
          texture: &texture,
          mip_level: 0,
          origin: Origin3d::ZERO,
          aspect: TextureAspect::All,
        },
        &image.data,
        ImageDataLayout {
          bytes_per_row: Some(4 * image.width as u32),
          offset: 0,
          rows_per_image: Some(image.height as u32),
        },
        texture_size,
      );

      textures.push(texture);
    }

    let sampler = device.create_sampler(&wgpu::SamplerDescriptor {
      address_mode_u: wgpu::AddressMode::ClampToEdge,
      address_mode_v: wgpu::AddressMode::ClampToEdge,
      address_mode_w: wgpu::AddressMode::ClampToEdge,
      mag_filter: wgpu::FilterMode::Linear,
      min_filter: wgpu::FilterMode::Nearest,
      mipmap_filter: wgpu::FilterMode::Nearest,
      ..Default::default()
    });

    let texture_bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
      entries: &[BindGroupLayoutEntry {
        binding: 0,
        ty: wgpu::BindingType::Texture {
          sample_type: wgpu::TextureSampleType::Float { filterable: true },
          view_dimension: wgpu::TextureViewDimension::D2,
          multisampled: false,
        },
        count: None,
        visibility: ShaderStages::FRAGMENT,
      }],
      label: None,
    });

    let texture_bind_groups = textures
      .iter()
      .map(|texture| {
        device.create_bind_group(&BindGroupDescriptor {
          label: None,
          layout: &texture_bind_group_layout,
          entries: &[BindGroupEntry {
            binding: 0,
            resource: wgpu::BindingResource::TextureView(&texture.create_view(&Default::default())),
          }],
        })
      })
      .collect();

    let bind_group_layout = device.create_bind_group_layout(&BindGroupLayoutDescriptor {
      label: None,
      entries: &[
        BindGroupLayoutEntry {
          binding: 0,
          ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
          count: None,
          visibility: ShaderStages::VERTEX_FRAGMENT,
        },
        BindGroupLayoutEntry {
          binding: 1,
          count: None,
          visibility: ShaderStages::VERTEX,
          ty: wgpu::BindingType::Buffer {
            ty: wgpu::BufferBindingType::Uniform,
            has_dynamic_offset: false,
            min_binding_size: Some(NonZero::new(std::mem::size_of::<Uniform>() as u64).unwrap()),
          },
        },
      ],
    });

    let uniform = Uniform {
      proj_matrix: Mat4::orthographic_lh(0.0, 640.0, 480.0, 0.0, -1.0, 1.0),
    };
    let uniform_buffer = device.create_buffer_init(&BufferInitDescriptor {
      contents: bytes_of(&uniform),
      label: None,
      usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST,
    });
    let bind_group = device.create_bind_group(&BindGroupDescriptor {
      label: None,
      layout: &bind_group_layout,
      entries: &[
        BindGroupEntry {
          binding: 0,
          resource: wgpu::BindingResource::Sampler(&sampler),
        },
        BindGroupEntry {
          binding: 1,
          resource: uniform_buffer.as_entire_binding(),
        },
      ],
    });

    let vertices = device.create_buffer_init(&BufferInitDescriptor {
      contents: bytes_of(&VERTICES),
      usage: BufferUsages::VERTEX,
      label: None,
    });
    let instances = device.create_buffer_init(&BufferInitDescriptor {
      contents: bytes_of(&*INSTANCES),
      usage: BufferUsages::VERTEX,
      label: None,
    });
    let shader = device.create_shader_module(include_wgsl!("./game.wgsl"));

    let pipeline_layout = device.create_pipeline_layout(&PipelineLayoutDescriptor {
      label: None,
      bind_group_layouts: &[&texture_bind_group_layout, &bind_group_layout],
      push_constant_ranges: &[],
    });

    let pipeline = device.create_render_pipeline(&RenderPipelineDescriptor {
      label: Some("main pass"),
      layout: Some(&pipeline_layout),
      vertex: VertexState {
        buffers: &[VertexBufferLayout {
          attributes: &[
            VertexAttribute {
              format: wgpu::VertexFormat::Float32x4,
              offset: 0,
              shader_location: 0,
            },
            VertexAttribute {
              format: wgpu::VertexFormat::Float32x4,
              offset: 16,
              shader_location: 1,
            },
            VertexAttribute {
              format: wgpu::VertexFormat::Float32x4,
              offset: 32,
              shader_location: 2,
            },
            VertexAttribute {
              format: wgpu::VertexFormat::Float32x4,
              offset: 48,
              shader_location: 3,
            },
          ],
          array_stride: std::mem::size_of::<Instance>() as u64,
          step_mode: wgpu::VertexStepMode::Instance,
        }],
        entry_point: "vs_main",
        compilation_options: Default::default(),
        module: &shader,
      },
      primitive: PrimitiveState {
        topology: wgpu::PrimitiveTopology::TriangleStrip,
        ..Default::default()
      },
      depth_stencil: None,
      multisample: Default::default(),
      fragment: Some(FragmentState {
        targets: &[Some(ColorTargetState {
          format: config.format,
          blend: Some(BlendState::ALPHA_BLENDING),
          write_mask: ColorWrites::ALL,
        })],
        compilation_options: Default::default(),
        entry_point: "fs_main",
        module: &shader,
      }),
      multiview: Default::default(),
      cache: Default::default(),
    });

    Self {
      anm,
      vertices,
      instances,
      textures,
      texture_bind_groups,
      bind_group,
      pipeline,
      shader,
      uniform_buffer,
      texture: 0,
      uniform: uniform,
    }
  }
}

impl EngineState for AnmViewer {
  fn update(&mut self, context: &UpdateContext) {
    if context.keys.is_key_pressed(KeyCode::ArrowLeft) {
      self.texture = self.texture.saturating_sub(1);
    } else if context.keys.is_key_pressed(KeyCode::ArrowRight) {
      self.texture = u32::min(self.texture + 1, self.anm.entries.len() as u32 - 1);
      log::info!("increased texture {}", self.texture)
    }
  }
  fn render(&self, state: &Engine, surface_view: &Texture, encoder: &mut wgpu::CommandEncoder) {
    state.queue.write_buffer(&self.uniform_buffer, 0, bytes_of(&self.uniform));
    let mut pass = encoder.begin_render_pass(&RenderPassDescriptor {
      color_attachments: &[Some(RenderPassColorAttachment {
        view: &surface_view.create_view(&Default::default()),
        resolve_target: None,
        ops: Operations {
          load: LoadOp::Clear(Color {
            r: 100. / 255.,
            g: 149. / 255.,
            b: 237. / 255.,
            a: 1.0,
          }),
          ..Default::default()
        },
      })],
      ..Default::default()
    });

    pass.set_pipeline(&self.pipeline);
    pass.set_vertex_buffer(0, self.instances.slice(..));
    pass.set_bind_group(0, &self.texture_bind_groups[self.texture as usize], &[]);
    pass.set_bind_group(1, &self.bind_group, &[]);
    pass.draw(0..4, 0..INSTANCES.len() as u32);
  }
}
