use env_logger::init;
use log::{error, info};
use roto_codegen::generator::generate_rust_code;
use roto_codegen::google::protobuf::compiler::plugin::{
    CodeGeneratorRequest, CodeGeneratorResponseBuilder, code_generator_response::FileBuilder,
};
use roto_codegen::google::protobuf::descriptor::FileDescriptorSet;
// use roto_codegen::runtime::ProtoBuilder;
use std::io::{self, Read, Write};

fn main() {
    // Initialize logger to print to stderr
    init();

    if let Err(e) = run() {
        error!("Plugin error: {}", e);
        std::process::exit(1);
    }
}

fn run() -> std::result::Result<(), Box<dyn std::error::Error>> {
    // 1. Read CodeGeneratorRequest from stdin
    let mut stdin_buf = Vec::new();
    io::stdin().read_to_end(&mut stdin_buf)?;

    if stdin_buf.is_empty() {
        error!("Received empty CodeGeneratorRequest from stdin");
        return Err("Empty stdin".into());
    }

    let request = CodeGeneratorRequest::new(&stdin_buf)?;

    // 2. Process request and get response bytes
    let response_bytes = handle_request(&request)?;

    // 3. Write response to stdout
    io::stdout().write_all(&response_bytes)?;

    info!("Successfully wrote CodeGeneratorResponse to stdout");
    Ok(())
}

/// Core logic that transforms a CodeGeneratorRequest into a serialized CodeGeneratorResponse.
fn handle_request(
    request: &CodeGeneratorRequest,
) -> std::result::Result<Vec<u8>, Box<dyn std::error::Error>> {
    // 2. Construct a FileDescriptorSet from the request's proto_files
    let mut set_buf = Vec::new();
    for file_res in request.proto_file() {
        let (file_data, _) = file_res.map_err(|e| {
            error!("Failed to iterate proto_file: {:?}", e);
            e
        })?;

        // Tag 1, Length-delimited: (1 << 3) | 2 = 10
        set_buf.push(10);

        // Write length as varint
        let len = file_data.len() as u64;
        let mut len_buf = [0u8; 10];
        let len_size = roto_codegen::runtime::write_varint(len, &mut len_buf).map_err(|e| {
            error!("Failed to write varint length: {:?}", e);
            e
        })?;
        set_buf.extend_from_slice(&len_buf[..len_size]);

        // Write data
        set_buf.extend_from_slice(file_data);
    }

    let set = FileDescriptorSet::new(&set_buf)?;

    let files_to_generate: Vec<String> = request
        .file_to_generate()
        .filter_map(|res| {
            let (bytes, _) = res.ok()?;
            std::str::from_utf8(bytes).ok().map(|s| s.to_string())
        })
        .collect();

    // Generate the Rust code
    info!("Generating Rust code from descriptor set...");
    let generated_files = generate_rust_code(&set, Some(&files_to_generate), false);

    // Construct the response
    let mut response_buf = vec![0u8; 1024 * 1024 * 2]; // Allocate 2MB for response
    let mut resp_builder = CodeGeneratorResponseBuilder::builder(&mut response_buf);

    for (filename, content) in generated_files {
        let mut file_buf = vec![0u8; 1024 * 1024 * 2];
        let final_file = FileBuilder::builder(&mut file_buf)
            .name(&filename)?
            .content(&content)?
            .finish()
            .map_err(|e| {
                error!("Failed to build ResponseFile {}: {:?}", filename, e);
                e
            })?;
        resp_builder = resp_builder.file(final_file)?;
    }

    let final_response_slice = resp_builder.finish().map_err(|e| {
        error!("Failed to finish CodeGeneratorResponse: {:?}", e);
        e
    })?;

    // The finish() method returns a reference to the buffer.
    // We convert it to a Vec<u8> to return it from this function.
    Ok(final_response_slice.to_vec())
}

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

    #[test]
    fn test_handle_request_with_bin() {
        // Note: this test assumes request.bin exists in the current working directory
        // which is usually the project root during 'cargo test'.
        let request_path = "request.bin";
        if !std::path::Path::new(request_path).exists() {
            return; // Skip if file is not available in the environment
        }

        let data = fs::read(request_path).expect("Failed to read request.bin");
        let request =
            CodeGeneratorRequest::new(&data).expect("Failed to parse CodeGeneratorRequest");

        let result = handle_request(&request);

        assert!(
            result.is_ok(),
            "handle_request should succeed with request.bin"
        );
        let response = result.unwrap();
        assert!(
            !response.is_empty(),
            "The generated response should not be empty"
        );
    }
}
