/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

use api::{MixBlendMode, units::PictureRect};
use crate::pattern::{PatternKind, PatternShaderInput};
use crate::renderer::BlendMode;
use crate::{spatial_tree::SpatialNodeIndex, render_task_graph::RenderTaskId, surface::SurfaceTileDescriptor, tile_cache::TileKey, renderer::GpuBufferAddress, FastHashMap};
use crate::gpu_types::QuadSegment;
use crate::prim_store::storage;
use crate::segment::EdgeMask;
use crate::transform::GpuTransformId;
use crate::visibility::PrimitiveDrawHeader;

/// A tightly packed command stored in a command buffer
#[cfg_attr(feature = "capture", derive(Serialize))]
#[cfg_attr(feature = "replay", derive(Deserialize))]
#[derive(Debug, Copy, Clone)]
pub struct Command(u32);

impl Command {
    /// Draw a simple primitive that needs prim instance index only.
    const CMD_DRAW_SIMPLE_PRIM: u32 = 0x00000000;
    /// Change the current spatial node.
    const CMD_SET_SPATIAL_NODE: u32 = 0x10000000;
    /// Draw a complex (3d-split) primitive, that has multiple GPU cache addresses.
    const CMD_DRAW_COMPLEX_PRIM: u32 = 0x20000000;
    /// Draw a primitive, that has a single GPU buffer addresses.
    const CMD_DRAW_INSTANCE: u32 = 0x30000000;
    /// Draw a generic quad primitive
    const CMD_DRAW_QUAD: u32 = 0x40000000;
    /// Set a list of variable-length segments
    const CMD_SET_SEGMENTS: u32 = 0x50000000;

    /// Bitmask for command bits of the command.
    const CMD_MASK: u32 = 0xf0000000;
    /// Bitmask for param bits of the command.
    const PARAM_MASK: u32 = 0x0fffffff;

    /// Encode drawing a simple primitive.
    fn draw_simple_prim(draw_index: storage::Index<PrimitiveDrawHeader>) -> Self {
        Command(Command::CMD_DRAW_SIMPLE_PRIM | draw_index.0)
    }

    /// Encode changing spatial node.
    fn set_spatial_node(spatial_node_index: SpatialNodeIndex) -> Self {
        Command(Command::CMD_SET_SPATIAL_NODE | spatial_node_index.0)
    }

    /// Encode a list of segments that follow
    fn set_segments(count: usize) -> Self {
        Command(Command::CMD_SET_SEGMENTS | count as u32)
    }

    /// Encode drawing a complex prim.
    fn draw_complex_prim(draw_index: storage::Index<PrimitiveDrawHeader>) -> Self {
        Command(Command::CMD_DRAW_COMPLEX_PRIM | draw_index.0)
    }

    fn draw_instance(draw_index: storage::Index<PrimitiveDrawHeader>) -> Self {
        Command(Command::CMD_DRAW_INSTANCE | draw_index.0)
    }

    /// Encode arbitrary data word.
    fn data(data: u32) -> Self {
        Command(data)
    }

    fn draw_quad(draw_index: storage::Index<PrimitiveDrawHeader>) -> Self {
        Command(Command::CMD_DRAW_QUAD | draw_index.0)
    }
}

bitflags! {
    /// Flags related to quad primitives
    #[repr(transparent)]
    #[cfg_attr(feature = "capture", derive(Serialize))]
    #[cfg_attr(feature = "replay", derive(Deserialize))]
    #[derive(Debug, Copy, PartialEq, Eq, Clone, PartialOrd, Ord, Hash)]
    pub struct QuadFlags : u8 {
        const IS_OPAQUE = 1 << 0;

        /// If true, the prim is 2d and axis-aligned in device space. The render task rect can
        /// cheaply be used as a device-space clip in the vertex shader.
        const APPLY_RENDER_TASK_CLIP = 1 << 1;

        /// If true, use segments for drawing the AA edges, to allow inner section to be opaque
        const USE_AA_SEGMENTS = 1 << 3;

        /// If true, render as a mask. This ignores the blue, green and alpha channels and replaces
        /// them with the red channel in the fragment shader. Used with multiply blending, on top
        /// of premultiplied alpha content, it has the effect of applying a mask to the content under ir.
        const IS_MASK = 1 << 4;
    }
}

bitflags! {
    /// Defines the space that a quad primitive is drawn in
    #[repr(transparent)]
    #[cfg_attr(feature = "capture", derive(Serialize))]
    #[cfg_attr(feature = "replay", derive(Deserialize))]
    #[derive(Debug, Copy, PartialEq, Eq, Clone, PartialOrd, Ord, Hash)]
    pub struct MaskFlags : i32 {
        const PRIM_SPACE = 1 << 0;
    }
}

/// The unpacked equivalent to a `Command`.
///
/// Each variant carries an `Index<PrimitiveDrawHeader>` identifying which
/// per-frame draw header to use. While `scratch.frame.draws` is identity-
/// indexed by `PrimitiveInstanceIndex.0`, the numerical value of the draw
/// index equals the prim instance index; a follow-up will make them
/// physically distinct when draws becomes push-per-draw.
#[cfg_attr(feature = "capture", derive(Serialize))]
pub enum PrimitiveCommand {
    Simple {
        draw_index: storage::Index<PrimitiveDrawHeader>,
    },
    Complex {
        draw_index: storage::Index<PrimitiveDrawHeader>,
        gpu_address: GpuBufferAddress,
    },
    Instance {
        draw_index: storage::Index<PrimitiveDrawHeader>,
        gpu_buffer_address: GpuBufferAddress,
    },
    Quad {
        pattern: PatternKind,
        pattern_input: PatternShaderInput,
        src_color_task_id: RenderTaskId,
        // TODO(gw): Used for bounding rect only, could possibly remove
        draw_index: storage::Index<PrimitiveDrawHeader>,
        gpu_buffer_address: GpuBufferAddress,
        transform_id: GpuTransformId,
        quad_flags: QuadFlags,
        edge_flags: EdgeMask,
        blend_mode: BlendMode,
    },
}

impl PrimitiveCommand {
    pub fn simple(
        draw_index: storage::Index<PrimitiveDrawHeader>,
    ) -> Self {
        PrimitiveCommand::Simple {
            draw_index,
        }
    }

    pub fn complex(
        draw_index: storage::Index<PrimitiveDrawHeader>,
        gpu_address: GpuBufferAddress,
    ) -> Self {
        PrimitiveCommand::Complex {
            draw_index,
            gpu_address,
        }
    }

    pub fn quad(
        pattern: PatternKind,
        pattern_input: PatternShaderInput,
        src_color_task_id: RenderTaskId,
        draw_index: storage::Index<PrimitiveDrawHeader>,
        gpu_buffer_address: GpuBufferAddress,
        transform_id: GpuTransformId,
        quad_flags: QuadFlags,
        edge_flags: EdgeMask,
        blend_mode: BlendMode,
    ) -> Self {
        PrimitiveCommand::Quad {
            pattern,
            pattern_input,
            src_color_task_id,
            draw_index,
            gpu_buffer_address,
            transform_id,
            quad_flags,
            edge_flags,
            blend_mode,
        }
    }

    pub fn instance(
        draw_index: storage::Index<PrimitiveDrawHeader>,
        gpu_buffer_address: GpuBufferAddress,
    ) -> Self {
        PrimitiveCommand::Instance {
            draw_index,
            gpu_buffer_address,
        }
    }
}


// Non-Advanced variants map to 0..=8; Advanced(mode) maps to 9 + mode as u32.
// MixBlendMode is repr(u8) with values 0..=16, so Advanced covers 9..=25.

fn encode_blend_mode(blend_mode: BlendMode) -> u32 {
    match blend_mode {
        BlendMode::None => 0,
        BlendMode::Alpha => 1,
        BlendMode::PremultipliedAlpha => 2,
        BlendMode::PremultipliedDestOut => 3,
        BlendMode::SubpixelDualSource => 4,
        BlendMode::MultiplyDualSource => 5,
        BlendMode::Screen => 6,
        BlendMode::Exclusion => 7,
        BlendMode::PlusLighter => 8,
        BlendMode::Advanced(mode) => 9 + mode as u32,
    }
}

fn decode_blend_mode(val: u32) -> BlendMode {
    match val {
        0 => BlendMode::None,
        1 => BlendMode::Alpha,
        2 => BlendMode::PremultipliedAlpha,
        3 => BlendMode::PremultipliedDestOut,
        4 => BlendMode::SubpixelDualSource,
        5 => BlendMode::MultiplyDualSource,
        6 => BlendMode::Screen,
        7 => BlendMode::Exclusion,
        8 => BlendMode::PlusLighter,
        _ => BlendMode::Advanced(unsafe { std::mem::transmute::<u8, MixBlendMode>((val - 9) as u8) }),
    }
}

/// A list of commands describing how to draw a primitive list.
#[cfg_attr(feature = "capture", derive(Serialize))]
#[cfg_attr(feature = "replay", derive(Deserialize))]
pub struct CommandBuffer {
    /// The encoded drawing commands.
    commands: Vec<Command>,
    /// Cached current spatial node.
    current_spatial_node_index: SpatialNodeIndex,
}

impl CommandBuffer {
    /// Construct a new cmd buffer.
    pub fn new() -> Self {
        CommandBuffer {
            commands: Vec::new(),
            current_spatial_node_index: SpatialNodeIndex::INVALID,
        }
    }

    /// Push a list of segments in to the cmd buffer
    pub fn set_segments(
        &mut self,
        segments: &[QuadSegment],
    ) {
        self.commands.push(Command::set_segments(segments.len()));
        for segment in segments {
            self.commands.push(Command::data(segment.task_id.index));
            self.commands.push(Command::data(segment.task_id.sub_rect_index as u32));
        }
    }

    /// Add a primitive to the command buffer.
    pub fn add_prim(
        &mut self,
        prim_cmd: &PrimitiveCommand,
        spatial_node_index: SpatialNodeIndex,
    ) {
        if self.current_spatial_node_index != spatial_node_index {
            self.commands.push(Command::set_spatial_node(spatial_node_index));
            self.current_spatial_node_index = spatial_node_index;
        }

        self.add_cmd(prim_cmd);
    }

    /// Add a cmd to the command buffer.
    pub fn add_cmd(
        &mut self,
        prim_cmd: &PrimitiveCommand,
    ) {
        match *prim_cmd {
            PrimitiveCommand::Simple { draw_index } => {
                self.commands.push(Command::draw_simple_prim(draw_index));
            }
            PrimitiveCommand::Complex { draw_index, gpu_address } => {
                self.commands.push(Command::draw_complex_prim(draw_index));
                self.commands.push(Command::data(gpu_address.as_u32()));
            }
            PrimitiveCommand::Instance { draw_index, gpu_buffer_address } => {
                self.commands.push(Command::draw_instance(draw_index));
                self.commands.push(Command::data(gpu_buffer_address.as_u32()));
            }
            PrimitiveCommand::Quad { pattern, pattern_input, draw_index, gpu_buffer_address, transform_id, quad_flags, edge_flags, src_color_task_id, blend_mode } => {
                self.commands.push(Command::draw_quad(draw_index));
                self.commands.push(Command::data(pattern as u32));
                self.commands.push(Command::data(pattern_input.0 as u32));
                self.commands.push(Command::data(pattern_input.1 as u32));
                self.commands.push(Command::data(src_color_task_id.index));
                self.commands.push(Command::data(src_color_task_id.sub_rect_index as u32));
                self.commands.push(Command::data(gpu_buffer_address.as_u32()));
                self.commands.push(Command::data(transform_id.0));
                self.commands.push(Command::data((quad_flags.bits() as u32) << 16 | edge_flags.bits() as u32));
                self.commands.push(Command::data(encode_blend_mode(blend_mode)));
            }
        }
    }

    /// Iterate the command list, calling a provided closure for each primitive draw command.
    pub fn iter_prims<F>(
        &self,
        f: &mut F,
    ) where F: FnMut(&PrimitiveCommand, SpatialNodeIndex, &[RenderTaskId]) {
        let mut current_spatial_node_index = SpatialNodeIndex::INVALID;
        let mut cmd_iter = self.commands.iter();
        // TODO(gw): Consider pre-allocating this / Smallvec if it shows up in profiles.
        let mut segments = Vec::new();

        while let Some(cmd) = cmd_iter.next() {
            let command = cmd.0 & Command::CMD_MASK;
            let param = cmd.0 & Command::PARAM_MASK;

            match command {
                Command::CMD_DRAW_SIMPLE_PRIM => {
                    let draw_index = storage::Index::from_u32(param);
                    let cmd = PrimitiveCommand::simple(draw_index);
                    f(&cmd, current_spatial_node_index, &[]);
                }
                Command::CMD_SET_SPATIAL_NODE => {
                    current_spatial_node_index = SpatialNodeIndex(param);
                }
                Command::CMD_DRAW_COMPLEX_PRIM => {
                    let draw_index = storage::Index::from_u32(param);
                    let data = cmd_iter.next().unwrap();
                    let gpu_address = GpuBufferAddress::from_u32(data.0);
                    let cmd = PrimitiveCommand::complex(
                        draw_index,
                        gpu_address,
                    );
                    f(&cmd, current_spatial_node_index, &[]);
                }
                Command::CMD_DRAW_QUAD => {
                    let draw_index = storage::Index::from_u32(param);
                    let pattern = PatternKind::from_u32(cmd_iter.next().unwrap().0);
                    let pattern_input = PatternShaderInput(
                        cmd_iter.next().unwrap().0 as i32,
                        cmd_iter.next().unwrap().0 as i32,
                    );
                    let src_color_task_id = RenderTaskId {
                        index: cmd_iter.next().unwrap().0,
                        sub_rect_index: cmd_iter.next().unwrap().0 as u16
                    };
                    let data = cmd_iter.next().unwrap();
                    let transform_id = GpuTransformId(cmd_iter.next().unwrap().0);
                    let bits = cmd_iter.next().unwrap().0;
                    let quad_flags = QuadFlags::from_bits((bits >> 16) as u8).unwrap();
                    let edge_flags = EdgeMask::from_bits((bits & 0xff) as u8).unwrap();
                    let blend_mode = decode_blend_mode(cmd_iter.next().unwrap().0);
                    let gpu_buffer_address = GpuBufferAddress::from_u32(data.0);
                    let cmd = PrimitiveCommand::quad(
                        pattern,
                        pattern_input,
                        src_color_task_id,
                        draw_index,
                        gpu_buffer_address,
                        transform_id,
                        quad_flags,
                        edge_flags,
                        blend_mode,
                    );
                    f(&cmd, current_spatial_node_index, &segments);
                    segments.clear()
                }
                Command::CMD_DRAW_INSTANCE => {
                    let draw_index = storage::Index::from_u32(param);
                    let data = cmd_iter.next().unwrap();
                    let gpu_buffer_address = GpuBufferAddress::from_u32(data.0);
                    let cmd = PrimitiveCommand::instance(
                        draw_index,
                        gpu_buffer_address,
                    );
                    f(&cmd, current_spatial_node_index, &[]);
                }
                Command::CMD_SET_SEGMENTS => {
                    let count = param;
                    for _ in 0 .. count {
                        segments.push(
                            RenderTaskId {
                                index: cmd_iter.next().unwrap().0,
                                sub_rect_index: cmd_iter.next().unwrap().0 as u16,
                            }
                        );
                    }
                }
                _ => {
                    unreachable!();
                }
            }
        }
    }
}

/// Abstracts whether a command buffer is being built for a tiled (picture cache)
/// or simple (child surface).
#[cfg_attr(feature = "capture", derive(Serialize))]
#[cfg_attr(feature = "replay", derive(Deserialize))]
pub enum CommandBufferBuilderKind {
    Tiled {
        // TODO(gw): It might be worth storing this as a 2d-array instead
        //           of a hash map if it ever shows up in profiles. This is
        //           slightly complicated by the sub_slice_index in the
        //           TileKey structure - could have a 2 level array?
        tiles: FastHashMap<TileKey, SurfaceTileDescriptor>,
    },
    Simple {
        render_task_id: RenderTaskId,
        root_task_id: Option<RenderTaskId>,
        dirty_rect: PictureRect,
    },
    Invalid,
}

#[cfg_attr(feature = "capture", derive(Serialize))]
#[cfg_attr(feature = "replay", derive(Deserialize))]
pub struct CommandBufferBuilder {
    pub kind: CommandBufferBuilderKind,

    /// If a command buffer establishes a sub-graph, then at the end of constructing
    /// the surface, the parent surface is supplied as an input dependency, and the
    /// parent surface gets a duplicated (existing) task with the same location, and
    /// with the sub-graph output as an input dependency.
    pub establishes_sub_graph: bool,

    /// If this surface builds a sub-graph, it will mark a task in the filter sub-graph
    /// as a resolve source for the input from the parent surface.
    pub resolve_source: Option<RenderTaskId>,

    /// List of render tasks that depend on the task that will be created for this builder.
    pub extra_dependencies: Vec<RenderTaskId>,
}

impl CommandBufferBuilder {
    pub fn empty() -> Self {
        CommandBufferBuilder {
            kind: CommandBufferBuilderKind::Invalid,
            establishes_sub_graph: false,
            resolve_source: None,
            extra_dependencies: Vec::new(),
        }
    }

    /// Construct a tiled command buffer builder.
    pub fn new_tiled(
        tiles: FastHashMap<TileKey, SurfaceTileDescriptor>,
    ) -> Self {
        CommandBufferBuilder {
            kind: CommandBufferBuilderKind::Tiled {
                tiles,
            },
            establishes_sub_graph: false,
            resolve_source: None,
            extra_dependencies: Vec::new(),
        }
    }

    /// Construct a simple command buffer builder.
    pub fn new_simple(
        render_task_id: RenderTaskId,
        establishes_sub_graph: bool,
        root_task_id: Option<RenderTaskId>,
        dirty_rect: PictureRect,
    ) -> Self {
        CommandBufferBuilder {
            kind: CommandBufferBuilderKind::Simple {
                render_task_id,
                root_task_id,
                dirty_rect,
            },
            establishes_sub_graph,
            resolve_source: None,
            extra_dependencies: Vec::new(),
        }
    }
}

// Index into a command buffer stored in a `CommandBufferList`.
#[cfg_attr(feature = "capture", derive(Serialize))]
#[cfg_attr(feature = "replay", derive(Deserialize))]
#[derive(Debug, Copy, Clone)]
pub struct CommandBufferIndex(pub u32);

// Container for a list of command buffers that are built for a frame.
pub struct CommandBufferList {
    cmd_buffers: Vec<CommandBuffer>,
}

impl CommandBufferList {
    pub fn new() -> Self {
        CommandBufferList {
            cmd_buffers: Vec::new(),
        }
    }

    pub fn create_cmd_buffer(
        &mut self,
    ) -> CommandBufferIndex {
        let index = CommandBufferIndex(self.cmd_buffers.len() as u32);
        self.cmd_buffers.push(CommandBuffer::new());
        index
    }

    pub fn get(&self, index: CommandBufferIndex) -> &CommandBuffer {
        &self.cmd_buffers[index.0 as usize]
    }

    pub fn get_mut(&mut self, index: CommandBufferIndex) -> &='color:red'>mut CommandBuffer {
        &mut self.cmd_buffers[index.0 as usize]
    }
}