use regex::Regex;
use std::env;
use std::fs;
use std::path::{Path, PathBuf};

fn main() {
    let watch_dir = "shaders/elements";
    for entry in fs::read_dir(watch_dir).unwrap() {
        let path = entry.unwrap().path();
        if path.is_file() {
            println!("cargo:rerun-if-changed={}", path.display());
        }
    }

    let crate_path = env::var("CARGO_MANIFEST_DIR").unwrap();

    let shader_path_base = Path::new(&crate_path).join("shaders");

    for entry in fs::read_dir(watch_dir).unwrap() {
        let path = entry.unwrap().path();
        if path.is_file() {
            println!("cargo:rerun-if-changed={}", path.display());
        }

        let merged_shader = merge_shader(&path.display().to_string(), &shader_path_base);

        let out_path = Path::new(&crate_path)
            .join("shaders")
            .join("elements")
            .join(format!(
                "{}_merged.wgsl",
                path.file_stem().unwrap().to_str().unwrap()
            ));

        fs::write(out_path, merged_shader).unwrap();
    }
}

fn merge_shader<'a>(shader: &'a str, base: &std::path::Path) -> String {
    const TOOLS: &'static str = r#"// This is a tool that merges the shader code with the shader code from the shader module.

struct UniformCommonTools {
    model_matrix: mat4x4f,
    view_matrix: mat4x4f,
    proj_matrix: mat4x4f,
    camera_x: f32,
    camera_y: f32,
    camera_z: f32,
    camera_target_x: f32,
    camera_target_y: f32,
    camera_target_z: f32,
    camera_up_x: f32,
    camera_up_y: f32,
    camera_up_z: f32,

    // camera_position: vec3f,
    // camera_front: vec3f,
    // camera_up: vec3f,
}

@group(0) @binding(0) var<uniform> common_tools: UniformCommonTools;
"#;

    let path = std::path::Path::new(shader);

    let mut visited_files = std::collections::HashSet::new();

    let mut result = process_file(&path, &mut visited_files, &base);

    // 将工具代码插入到 shader 代码的开头
    result.insert_str(0, TOOLS);

    return result;
}

fn process_file<'a, P: AsRef<std::path::Path> + 'a, B: AsRef<std::path::Path>>(
    file_path: P,
    visited_files: &mut std::collections::HashSet<std::path::PathBuf>,
    base_path: B,
) -> String {
    let base_path = base_path.as_ref();
    let file_path = file_path.as_ref();

    // 如果该文件已经处理过，则避免死循环
    if visited_files.contains(file_path) {
        return String::new(); // 返回空字符串表示已经处理过
    }

    visited_files.insert(file_path.to_path_buf());

    let content = std::fs::read_to_string(file_path).unwrap_or_else(|_| {
        panic!("Failed to read file: {}", file_path.display());
    });

    let re = Regex::new(r#"#include\s+\"([^\"]+\.wgsl)\";"#).unwrap();

    let result = re.replace_all(&content, |caps: &regex::Captures| {
        let included_file = &caps[1]; // 获取文件名

        let included_path = resolve_included_file(included_file, base_path.as_ref());

        let included_content = process_file(
            &included_path,
            visited_files,
            base_path.parent().unwrap_or(base_path),
        );

        included_content
    });

    result.to_string()
}

fn resolve_included_file(included_file: &str, base_path: &std::path::Path) -> std::path::PathBuf {
    let included_path = std::path::Path::new(included_file);

    if included_path.is_relative() {
        base_path.join(included_path)
    } else {
        included_path.to_path_buf()
    }
}