pub mod modules;
pub mod templates;

use crate::ir::SrcGenIR;
use crate::{GeneratedRust, SrcGenBackend};
use proc_macro2::TokenStream;
use quote::ToTokens;
use syn::{File, Item, parse_file};

const MODULE_ITEM_CHUNK_SIZE: usize = 256;
const MODULE_CHUNK_TRIGGER_ITEMS: usize = MODULE_ITEM_CHUNK_SIZE * 2;
const MODULE_BYTE_CHUNK_SIZE: usize = 256 * 1024;
const MODULE_CHUNK_TRIGGER_BYTES: usize = MODULE_BYTE_CHUNK_SIZE * 2;
const CHUNKABLE_MODULES: &[&str] = &[
    "prelude.rs",
    "targets.rs",
    "validators_property.rs",
    "validators_node.rs",
    "inference.rs",
    "report.rs",
];

pub fn generate_modules(ir: &SrcGenIR, backend: SrcGenBackend) -> Result<GeneratedRust, String> {
    let files = modules::generate_module_files(ir, backend)?;
    let files = chunk_large_modules(files)?;
    let binary_files = if modules::run::has_planned_runtime_fallback(ir) {
        vec![(
            "shape_ir.bin".to_string(),
            ir.embedded_shape_ir_bincode.clone(),
        )]
    } else {
        Vec::new()
    };
    let root = build_root_module();
    Ok(GeneratedRust {
        root,
        files,
        binary_files,
    })
}

fn build_root_module() -> String {
    [
        "// Code generated by shacl-srcgen-compiler. DO NOT EDIT.",
        "#![allow(unused_variables)]",
        "#![allow(dead_code)]",
        "#![allow(unused_mut)]",
        "#![allow(unused_imports)]",
        "#![allow(unused_comparisons)]",
        "include!(\"prelude.rs\");",
        "include!(\"paths.rs\");",
        "include!(\"targets.rs\");",
        "include!(\"validators_property.rs\");",
        "include!(\"validators_node.rs\");",
        "include!(\"inference.rs\");",
        "include!(\"report.rs\");",
        "include!(\"run.rs\");",
        "",
    ]
    .join("\n")
}

pub(crate) fn render_tokens_as_module(tokens: TokenStream) -> Result<String, String> {
    let syntax = syn::parse_file(&tokens.to_string())
        .map_err(|err| format!("failed to parse generated tokens: {err}"))?;
    Ok(prettyplease::unparse(&syntax))
}

fn chunk_large_modules(files: Vec<(String, String)>) -> Result<Vec<(String, String)>, String> {
    let mut out: Vec<(String, String)> = Vec::new();
    for (name, content) in files {
        if !CHUNKABLE_MODULES.contains(&name.as_str()) {
            out.push((name, content));
            continue;
        }
        out.extend(chunk_module_file(
            &name,
            &content,
            MODULE_CHUNK_TRIGGER_ITEMS,
            MODULE_ITEM_CHUNK_SIZE,
            MODULE_CHUNK_TRIGGER_BYTES,
            MODULE_BYTE_CHUNK_SIZE,
        )?);
    }
    Ok(out)
}

fn chunk_module_file(
    module_name: &str,
    content: &str,
    chunk_trigger_items: usize,
    chunk_item_size: usize,
    chunk_trigger_bytes: usize,
    chunk_byte_size: usize,
) -> Result<Vec<(String, String)>, String> {
    let syntax = parse_file(content)
        .map_err(|err| format!("generated module {module_name} failed to parse: {err}"))?;
    if !syntax.attrs.is_empty() {
        return Ok(vec![(module_name.to_string(), content.to_string())]);
    }

    let stem = module_name
        .strip_suffix(".rs")
        .ok_or_else(|| format!("module name missing .rs suffix: {module_name}"))?;

    let chunk_items: Vec<Vec<Item>> = if syntax.items.len() > chunk_trigger_items {
        split_items_by_count(syntax.items, chunk_item_size)
    } else if content.len() > chunk_trigger_bytes {
        // Keep include! chunks syntax-safe: never split inside a Rust item.
        if syntax.items.len() <= 1 {
            return Ok(vec![(module_name.to_string(), content.to_string())]);
        }
        split_items_by_approx_bytes(syntax.items, chunk_byte_size)
    } else {
        return Ok(vec![(module_name.to_string(), content.to_string())]);
    };

    let mut chunk_files: Vec<(String, String)> = Vec::with_capacity(chunk_items.len());
    for (idx, items) in chunk_items.into_iter().enumerate() {
        let chunk_syntax = File {
            shebang: None,
            attrs: Vec::new(),
            items,
        };
        let chunk_content = prettyplease::unparse(&chunk_syntax);
        let chunk_name = format!("{stem}/chunk_{idx:04}.rs");
        chunk_files.push((chunk_name, chunk_content));
    }

    let mut wrapper = String::from("// Code generated by shacl-srcgen-compiler. DO NOT EDIT.\n");
    for idx in 0..chunk_files.len() {
        wrapper.push_str(&format!("include!(\"{stem}/chunk_{idx:04}.rs\");\n"));
    }

    let mut out = Vec::with_capacity(1 + chunk_files.len());
    out.push((module_name.to_string(), wrapper));
    out.extend(chunk_files);
    Ok(out)
}

fn split_items_by_count(items: Vec<Item>, chunk_item_size: usize) -> Vec<Vec<Item>> {
    let mut out: Vec<Vec<Item>> = Vec::new();
    let mut items = items;
    while !items.is_empty() {
        let take = chunk_item_size.min(items.len());
        out.push(items.drain(0..take).collect());
    }
    out
}

fn split_items_by_approx_bytes(items: Vec<Item>, chunk_byte_size: usize) -> Vec<Vec<Item>> {
    let mut out: Vec<Vec<Item>> = Vec::new();
    let mut current: Vec<Item> = Vec::new();
    let mut current_bytes = 0usize;

    for item in items {
        let item_bytes = item.to_token_stream().to_string().len();
        if !current.is_empty() && current_bytes.saturating_add(item_bytes) > chunk_byte_size {
            out.push(current);
            current = Vec::new();
            current_bytes = 0;
        }
        current_bytes = current_bytes.saturating_add(item_bytes);
        current.push(item);
    }

    if !current.is_empty() {
        out.push(current);
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn chunk_module_file_splits_large_item_module() {
        let mut src = String::new();
        for idx in 0..20 {
            src.push_str(&format!("fn generated_{idx}() {{}}\n"));
        }

        let files = chunk_module_file("validators_node.rs", &src, 4, 3, usize::MAX, usize::MAX)
            .expect("chunking should succeed");
        assert_eq!(
            files.first().map(|f| f.0.as_str()),
            Some("validators_node.rs")
        );
        assert!(
            files[0]
                .1
                .contains("include!(\"validators_node/chunk_0000.rs\");")
        );
        assert!(
            files
                .iter()
                .any(|(name, _)| name == "validators_node/chunk_0001.rs")
        );
    }

    #[test]
    fn chunk_module_file_splits_large_byte_module() {
        let mut src = String::from("fn giant_a() {\n");
        for _ in 0..200 {
            src.push_str("    let _x = 42;\n");
        }
        src.push_str("}\n");
        src.push_str("fn giant_b() {\n");
        for _ in 0..200 {
            src.push_str("    let _x = 42;\n");
        }
        src.push_str("}\n");

        let files = chunk_module_file(
            "validators_property.rs",
            &src,
            usize::MAX,
            usize::MAX,
            128,
            64,
        )
        .expect("chunking should succeed");
        assert_eq!(
            files.first().map(|f| f.0.as_str()),
            Some("validators_property.rs")
        );
        assert!(
            files[0]
                .1
                .contains("include!(\"validators_property/chunk_0000.rs\");")
        );
        assert!(
            files
                .iter()
                .any(|(name, _)| name == "validators_property/chunk_0001.rs")
        );
    }

    #[test]
    fn chunk_module_file_does_not_split_single_large_item_module() {
        let mut src = String::from("fn giant_only() {\n");
        for _ in 0..400 {
            src.push_str("    let _x = 42;\n");
        }
        src.push_str("}\n");

        let files = chunk_module_file("prelude.rs", &src, usize::MAX, usize::MAX, 128, 64)
            .expect("chunking should succeed");
        assert_eq!(files.len(), 1);
        assert_eq!(files[0].0, "prelude.rs");
    }
}
