use crate::{
    all::ForestKind,
    block::block_from_structure,
    column::ColumnGen,
    util::{RandomPerm, Sampler, UnitChooser},
    Canvas, CONFIG,
};
use common::{
    assets::AssetHandle,
    terrain::{Block, BlockKind, structure::{Structure, StructureBlock, StructuresGroup}},
    vol::ReadVol,
};
use hashbrown::HashMap;
use lazy_static::lazy_static;
use std::{f32, ops::Range};
use vek::*;
use rand::prelude::*;

lazy_static! {
    static ref OAKS: AssetHandle<StructuresGroup> = Structure::load_group("oaks");
    static ref OAK_STUMPS: AssetHandle<StructuresGroup> = Structure::load_group("oak_stumps");
    static ref PINES: AssetHandle<StructuresGroup> = Structure::load_group("pines");
    static ref PALMS: AssetHandle<StructuresGroup> = Structure::load_group("palms");
    static ref ACACIAS: AssetHandle<StructuresGroup> = Structure::load_group("acacias");
    static ref BAOBABS: AssetHandle<StructuresGroup> = Structure::load_group("baobabs");
    static ref FRUIT_TREES: AssetHandle<StructuresGroup> = Structure::load_group("fruit_trees");
    static ref BIRCHES: AssetHandle<StructuresGroup> = Structure::load_group("birch");
    static ref MANGROVE_TREES: AssetHandle<StructuresGroup> =
        Structure::load_group("mangrove_trees");
    static ref QUIRKY: AssetHandle<StructuresGroup> = Structure::load_group("quirky");
    static ref QUIRKY_DRY: AssetHandle<StructuresGroup> = Structure::load_group("quirky_dry");
    static ref SWAMP_TREES: AssetHandle<StructuresGroup> = Structure::load_group("swamp_trees");
}

static MODEL_RAND: RandomPerm = RandomPerm::new(0xDB21C052);
static UNIT_CHOOSER: UnitChooser = UnitChooser::new(0x700F4EC7);
static QUIRKY_RAND: RandomPerm = RandomPerm::new(0xA634460F);

#[allow(clippy::if_same_then_else)]
pub fn apply_trees_to(canvas: &mut Canvas) {
    // TODO: Get rid of this
    enum TreeModel {
        Structure(Structure),
        Procedural(ProceduralTree),
    }

    struct Tree {
        pos: Vec3<i32>,
        model: TreeModel,
        seed: u32,
        units: (Vec2<i32>, Vec2<i32>),
    }

    let mut tree_cache = HashMap::new();

    let info = canvas.info();
    canvas.foreach_col(|canvas, wpos2d, col| {
        let trees = info.land().get_near_trees(wpos2d);

        for (tree_wpos, seed) in trees {
            let tree = if let Some(tree) = tree_cache.entry(tree_wpos).or_insert_with(|| {
                let col = ColumnGen::new(info.land()).get((tree_wpos, info.index()))?;

                let is_quirky = QUIRKY_RAND.chance(seed, 1.0 / 500.0);

                // Ensure that it's valid to place a *thing* here
                if col.alt < col.water_level
                    || col.spawn_rate < 0.9
                    || col.water_dist.map(|d| d < 8.0).unwrap_or(false)
                    || col.path.map(|(d, _, _, _)| d < 12.0).unwrap_or(false)
                {
                    return None;
                }

                // Ensure that it's valid to place a tree here
                if !is_quirky && ((seed.wrapping_mul(13)) & 0xFF) as f32 / 256.0 > col.tree_density
                {
                    return None;
                }

                Some(Tree {
                    pos: Vec3::new(tree_wpos.x, tree_wpos.y, col.alt as i32),
                    model: 'model: {
                        let models: AssetHandle<_> = if is_quirky {
                            if col.temp > CONFIG.desert_temp {
                                *QUIRKY_DRY
                            } else {
                                *QUIRKY
                            }
                        } else {
                            match col.forest_kind {
                                ForestKind::Oak if QUIRKY_RAND.chance(seed + 1, 1.0 / 16.0) => {
                                    *OAK_STUMPS
                                },
                                ForestKind::Oak if QUIRKY_RAND.chance(seed + 2, 1.0 / 20.0) => {
                                    *FRUIT_TREES
                                },
                                ForestKind::Palm => *PALMS,
                                ForestKind::Acacia => *ACACIAS,
                                ForestKind::Baobab => *BAOBABS,
                                // ForestKind::Oak => *OAKS,
                                ForestKind::Oak => {
                                    break 'model TreeModel::Procedural(ProceduralTree::generate(TreeConfig::OAK, seed));
                                },
                                ForestKind::Pine => *PINES,
                                ForestKind::Birch => *BIRCHES,
                                ForestKind::Mangrove => *MANGROVE_TREES,
                                ForestKind::Swamp => *SWAMP_TREES,
                            }
                        };

                        let models = models.read();
                        TreeModel::Structure(models[(MODEL_RAND.get(seed.wrapping_mul(17)) / 13) as usize % models.len()]
                            .clone())
                    },
                    seed,
                    units: UNIT_CHOOSER.get(seed),
                })
            }) {
                tree
            } else {
                continue;
            };

            let bounds = match &tree.model {
                TreeModel::Structure(s) => s.get_bounds(),
                TreeModel::Procedural(t) => t.get_bounds().map(|e| e as i32),
            };

            let rpos2d = (wpos2d - tree.pos.xy())
                .map2(Vec2::new(tree.units.0, tree.units.1), |p, unit| {
                    unit * p
                })
                .sum();
            if !Aabr::from(bounds).contains_point(rpos2d) {
                // Skip this column
                continue;
            }

            let mut is_top = true;
            let mut is_leaf_top = true;
            for z in (bounds.min.z..bounds.max.z).rev() {
                let wpos = Vec3::new(wpos2d.x, wpos2d.y, tree.pos.z + z);
                let model_pos = Vec3::from(
                    (wpos - tree.pos)
                        .xy()
                        .map2(Vec2::new(tree.units.0, tree.units.1), |rpos, unit| {
                            unit * rpos
                        })
                        .sum(),
                ) + Vec3::unit_z() * (wpos.z - tree.pos.z);
                block_from_structure(
                    info.index(),
                    if let Some(block) = match &tree.model {
                        TreeModel::Structure(s) => s.get(model_pos).ok().copied(),
                        TreeModel::Procedural(t) => Some(match t.is_branch_or_leaves_at(model_pos.map(|e| e as f32 + 0.5)) {
                            (true, _) => StructureBlock::Normal(Rgb::new(60, 30, 0)),
                            (_, true) => StructureBlock::TemperateLeaves,
                            (_, _) => StructureBlock::None,
                        }),
                    } {
                        block
                    } else {
                        break
                    },
                    wpos,
                    tree.pos.xy(),
                    tree.seed,
                    col,
                    Block::air,
                )
                .map(|block| {
                    // Add a snow covering to the block above under certain circumstances
                    if col.snow_cover
                        && ((block.kind() == BlockKind::Leaves && is_leaf_top)
                            || (is_top && block.is_filled()))
                    {
                        canvas.set(
                            wpos + Vec3::unit_z(),
                            Block::new(BlockKind::Snow, Rgb::new(210, 210, 255)),
                        );
                    }
                    canvas.set(wpos, block);
                    is_leaf_top = false;
                    is_top = false;
                })
                .unwrap_or_else(|| {
                    is_leaf_top = true;
                });
            }
        }
    });
}

/// A type that specifies the generation properties of a tree.
pub struct TreeConfig {
    /// Length of trunk, also scales other branches.
    pub trunk_len: f32,
    /// Radius of trunk, also scales other branches.
    pub trunk_radius: f32,
    // The scale that child branch lengths should be compared to their parents
    pub branch_child_len: f32,
    // The scale that child branch radii should be compared to their parents
    pub branch_child_radius: f32,
    /// 0 - 1 (0 = chaotic, 1 = straight).
    pub straightness: f32,
    /// Maximum number of branch layers (not including trunk).
    pub max_depth: usize,
    /// The number of branches that form from each branch.
    pub splits: usize,
    /// The range of proportions along a branch at which a split into another branch might occur.
    /// This value is clamped between 0 and 1, but a wider range may bias the results towards branch ends.
    pub split_range: Range<f32>,
}

impl TreeConfig {
    pub const OAK: Self = Self {
        trunk_len: 12.0,
        trunk_radius: 3.0,
        branch_child_len: 0.8,
        branch_child_radius: 0.6,
        straightness: 0.5,
        max_depth: 4,
        splits: 3,
        split_range: 0.5..1.5,
    };
}

// TODO: Rename this to `Tree` when the name conflict is gone
pub struct ProceduralTree {
    branches: Vec<Branch>,
    trunk_idx: usize,
}

impl ProceduralTree {
    /// Generate a new tree using the given configuration and seed.
    pub fn generate(config: TreeConfig, seed: u32) -> Self {
        let mut rng = RandomPerm::new(seed);

        let mut this = Self {
            branches: Vec::new(),
            trunk_idx: 0, // Gets replaced later
        };

        // Add the tree trunk (and sub-branches) recursively
        let (trunk_idx, _) = this.add_branch(
            &config,
            // Our trunk starts at the origin...
            Vec3::zero(),
            // ...and has a roughly upward direction
            Vec3::new(rng.gen_range(-1.0, 1.0), rng.gen_range(-1.0, 1.0), 5.0).normalized(),
            config.trunk_len,
            config.trunk_radius,
            0,
            None,
            &mut rng,
        );
        this.trunk_idx = trunk_idx;

        this
    }

    // Recursively add a branch (with sub-branches) to the tree's branch graph, returning the index and AABB of the
    // branch. This AABB gets propagated down to the parent and is used later during sampling to cull the branches to
    // be sampled.
    fn add_branch(
        &mut self,
        config: &TreeConfig,
        start: Vec3<f32>,
        dir: Vec3<f32>,
        branch_len: f32,
        branch_radius: f32,
        depth: usize,
        sibling_idx: Option<usize>,
        rng: &mut impl Rng,
    ) -> (usize, Aabb<f32>) {
        let end = start + dir * branch_len;
        let line = LineSegment3 { start, end };
        let wood_radius = branch_radius;
        let leaf_radius = if depth == config.max_depth { rng.gen_range(3.0, 5.0) } else { 0.0 };

        // The AABB that covers this branch, along with wood and leaves that eminate from it
        let mut aabb = Aabb {
            min: Vec3::partial_min(start, end) - wood_radius.max(leaf_radius),
            max: Vec3::partial_max(start, end) + wood_radius.max(leaf_radius),
        };

        let mut child_idx = None;
        // Don't add child branches if we're already enough layers into the tree
        if depth < config.max_depth {
            for _ in 0..config.splits {
                // Choose a point close to the branch to act as the target direction for the branch to grow in
                let tgt = Lerp::lerp(
                    start,
                    end,
                    rng.gen_range(config.split_range.start, config.split_range.end).clamped(0.0, 1.0),
                ) + Vec3::<f32>::zero().map(|_| rng.gen_range(-1.0, 1.0));
                // Start the branch at the closest point to the target
                let branch_start = line.projected_point(tgt);
                // Now, interpolate between the target direction and the parent branch's direction to find a direction
                let branch_dir = Lerp::lerp(tgt - branch_start, dir, config.straightness).normalized();

                let (branch_idx, branch_aabb) = self.add_branch(
                    config,
                    branch_start,
                    branch_dir,
                    branch_len * config.branch_child_len,
                    branch_radius * config.branch_child_radius,
                    depth + 1,
                    child_idx,
                    rng,
                );
                child_idx = Some(branch_idx);
                // Parent branches AABBs include the AABBs of child branches to allow for culling during sampling
                aabb.expand_to_contain(branch_aabb);
            }
        }

        let idx = self.branches.len(); // Compute the index that this branch is going to have
        self.branches.push(Branch {
            line,
            wood_radius,
            leaf_radius,
            aabb,
            sibling_idx,
            child_idx,
        });

        (idx, aabb)
    }

    /// Get the bounding box that covers the tree (all branches and leaves)
    pub fn get_bounds(&self) -> Aabb<f32> {
        self.branches[self.trunk_idx].aabb
    }

    // Recursively search for branches or leaves by walking the tree's branch graph.
    fn is_branch_or_leaves_at_inner(&self, pos: Vec3<f32>, branch_idx: usize) -> (bool, bool) {
        let branch = &self.branches[branch_idx];
        // Always probe the sibling branch, since our AABB doesn't include its bounds (it's not one of our children)
        let branch_or_leaves = branch.sibling_idx
            .map(|idx| Vec2::from(self.is_branch_or_leaves_at_inner(pos, idx)))
            .unwrap_or_default();

        // Only continue probing this sub-graph of the tree if the sample position falls within its AABB
        if branch.aabb.contains_point(pos) {
            (branch_or_leaves
                // Probe this branch
                | Vec2::from(branch.is_branch_or_leaves_at(pos))
                // Probe the children of this branch
                | branch.child_idx
                    .map(|idx| Vec2::from(self.is_branch_or_leaves_at_inner(pos, idx)))
                    .unwrap_or_default())
                .into_tuple()
        } else {
            branch_or_leaves.into_tuple()
        }
    }

    /// Determine whether there are either branches or leaves at the given position in the tree.
    #[inline(always)]
    pub fn is_branch_or_leaves_at(&self, pos: Vec3<f32>) -> (bool, bool) {
        self.is_branch_or_leaves_at_inner(pos, self.trunk_idx)
    }
}

// Branches are arranged in a graph shape. Each branch points to both its first child (if any) and also to the next
// branch in the list of child branches associated with the parent. This means that the entire tree is laid out in a
// walkable graph where each branch refers only to two other branches. As a result, walking the tree is simply a case
// of performing double recursion.
struct Branch {
    line: LineSegment3<f32>,
    wood_radius: f32,
    leaf_radius: f32,
    aabb: Aabb<f32>,

    sibling_idx: Option<usize>,
    child_idx: Option<usize>,
}

impl Branch {
    /// Determine whether there are either branches or leaves at the given position in the branch.
    pub fn is_branch_or_leaves_at(&self, pos: Vec3<f32>) -> (bool, bool) {
        // fn finvsqrt(x: f32) -> f32 {
        //     let y = f32::from_bits(0x5f375a86 - (x.to_bits() >> 1));
        //     y * (1.5 - ( x * 0.5 * y * y ))
        // }

        let p_d2 = self.line.projected_point(pos).distance_squared(pos);

        if p_d2 < self.wood_radius.powi(2) {
            (true, false)
        } else {
            (false, p_d2 < self.leaf_radius.powi(2))
        }
    }
}