jvmrs 0.1.2

//! Built-in runtime operations (println, string concatenation, etc.).

use crate::class_file::{ClassFile, ConstantPoolEntry};
use crate::error::{to_runtime_error_enum, RuntimeError};
use crate::memory::{StackFrame, Value};
use byteorder::{BigEndian, ReadBytesExt};
use std::io::Cursor;

use super::Interpreter;
use super::InterpreterResult;

/// Parse the BootstrapMethods attribute to find the method at the given index
fn get_bootstrap_method(class: &ClassFile, index: u16) -> Option<(u16, Vec<u16>)> {
    for attr in &class.attributes {
        let name = class.get_string(attr.attribute_name_index)?;
        if name == "BootstrapMethods" {
            let mut cursor = Cursor::new(&attr.info);
            let num_bootstrap_methods = cursor.read_u16::<BigEndian>().ok()?;
            if index >= num_bootstrap_methods {
                return None;
            }
            for i in 0..num_bootstrap_methods {
                let method_ref = cursor.read_u16::<BigEndian>().ok()?;
                let num_args = cursor.read_u16::<BigEndian>().ok()?;
                let mut args = Vec::with_capacity(num_args as usize);
                for _ in 0..num_args {
                    args.push(cursor.read_u16::<BigEndian>().ok()?);
                }
                if i == index {
                    return Some((method_ref, args));
                }
            }
        }
    }
    None
}

/// Resolve a CONSTANT_MethodHandle to class name and method name
fn resolve_method_handle(class: &ClassFile, handle_index: u16) -> Option<(String, String)> {
    let entry = class.constant_pool.get(handle_index as usize)?;
    match entry {
        ConstantPoolEntry::ConstantMethodHandle {
            reference_index, ..
        } => {
            let ref_entry = class.constant_pool.get(*reference_index as usize)?;
            match ref_entry {
                ConstantPoolEntry::ConstantMethodref {
                    class_index,
                    name_and_type_index,
                }
                | ConstantPoolEntry::ConstantInterfaceMethodref {
                    class_index,
                    name_and_type_index,
                } => {
                    let class_name = class.get_class_name_from_index(*class_index)?;
                    let nt_entry = class.constant_pool.get(*name_and_type_index as usize)?;
                    match nt_entry {
                        ConstantPoolEntry::ConstantNameAndType { name_index, .. } => {
                            let method_name = class.get_string(*name_index)?;
                            Some((class_name, method_name))
                        }
                        _ => None,
                    }
                }
                _ => None,
            }
        }
        _ => None,
    }
}

/// Get a constant value as string from the constant pool
fn get_constant_value_as_string(class: &ClassFile, index: u16) -> String {
    if let Some(entry) = class.constant_pool.get(index as usize) {
        match entry {
            ConstantPoolEntry::ConstantString { string_index } => {
                class.get_string(*string_index).unwrap_or_default()
            }
            ConstantPoolEntry::ConstantInteger { bytes } => bytes.to_string(),
            ConstantPoolEntry::ConstantFloat { bytes } => bytes.to_string(),
            ConstantPoolEntry::ConstantLong { bytes } => bytes.to_string(),
            ConstantPoolEntry::ConstantDouble { bytes } => bytes.to_string(),
            ConstantPoolEntry::ConstantClass { name_index } => {
                class.get_string(*name_index).unwrap_or_default()
            }
            ConstantPoolEntry::ConstantUtf8 { bytes } => String::from_utf8_lossy(bytes).to_string(),
            _ => format!("Constant#{}", index),
        }
    } else {
        String::new()
    }
}

impl Interpreter {
    /// Native println implementation (PrintStream.println)
    pub(crate) fn native_println(&mut self, value: Value) -> InterpreterResult {
        match value {
            Value::Int(i) => println!("{}", i),
            Value::Float(f) => println!("{}", f),
            Value::Long(l) => println!("{}", l),
            Value::Double(d) => println!("{}", d),
            Value::Reference(addr) => {
                if let Some(obj) = self.memory.heap.get_object(addr) {
                    if obj.class_name == "java/lang/String" {
                        if let Some(string_data) = &obj.string_data {
                            println!("{}", string_data);
                        } else {
                            println!("[String]");
                        }
                    } else {
                        println!("[{}]", obj.class_name);
                    }
                } else {
                    println!("null");
                }
            }
            Value::ArrayRef(addr) => println!("[Array@{}]", addr),
            Value::Null => println!("null"),
            _ => println!("{:?}", value),
        }
        Ok(())
    }

    /// Handle invokedynamic instruction (e.g. string concatenation)
    pub(crate) fn handle_invokedynamic(
        &mut self,
        class: &ClassFile,
        frame: &mut StackFrame,
        index: usize,
    ) -> InterpreterResult {
        let cp_entry = class.constant_pool.get(index).ok_or_else(|| {
            to_runtime_error_enum(RuntimeError::IllegalArgument(format!(
                "Constant pool index {} out of bounds",
                index
            )))
        })?;

        let (bootstrap_index, name_and_type_index) = match cp_entry {
            ConstantPoolEntry::ConstantInvokeDynamic {
                bootstrap_method_attr_index,
                name_and_type_index,
            } => (bootstrap_method_attr_index, name_and_type_index),
            _ => {
                return Err(to_runtime_error_enum(RuntimeError::IllegalArgument(
                    "Expected InvokeDynamic constant".to_string(),
                )));
            }
        };

        let name_and_type_entry = class.constant_pool.get(*name_and_type_index as usize);
        let (_method_name, descriptor) = match name_and_type_entry {
            Some(ConstantPoolEntry::ConstantNameAndType {
                name_index,
                descriptor_index,
            }) => (
                class.get_string(*name_index).unwrap_or_default(),
                class.get_string(*descriptor_index).unwrap_or_default(),
            ),
            _ => {
                return Err(to_runtime_error_enum(RuntimeError::IllegalArgument(
                    "Invalid NameAndType in invokedynamic".to_string(),
                )));
            }
        };

        // Resolve bootstrap method
        let (bsm_ref, bsm_args) =
            get_bootstrap_method(class, *bootstrap_index).ok_or_else(|| {
                to_runtime_error_enum(RuntimeError::IllegalArgument(format!(
                    "Invalid bootstrap method index {}",
                    bootstrap_index
                )))
            })?;

        let (bsm_class, bsm_name) = resolve_method_handle(class, bsm_ref).ok_or_else(|| {
            to_runtime_error_enum(RuntimeError::IllegalArgument(
                "Could not resolve bootstrap method handle".to_string(),
            ))
        })?;

        if bsm_class == "java/lang/invoke/StringConcatFactory"
            && bsm_name == "makeConcatWithConstants"
        {
            let param_count = crate::interpreter::descriptor::count_parameters(&descriptor);
            let mut stack_args = Vec::with_capacity(param_count);
            for _ in 0..param_count {
                stack_args.push(frame.pop()?);
            }
            stack_args.reverse();

            let mut result = String::new();

            if bsm_args.is_empty() {
                // Simple concatenation (unlikely for makeConcatWithConstants, but fallback)
                for arg in stack_args {
                    let s = self.value_to_string(&arg);
                    result.push_str(&s);
                }
            } else {
                // Recipe based
                // bsm_args[0] is recipe
                let recipe_index = bsm_args[0];
                let recipe = get_constant_value_as_string(class, recipe_index);

                let mut arg_idx = 0;
                let mut const_idx = 0;
                // constants start from bsm_args[1]

                for c in recipe.chars() {
                    if c == '\u{1}' {
                        // TAG_ARG
                        if arg_idx < stack_args.len() {
                            let s = self.value_to_string(&stack_args[arg_idx]);
                            result.push_str(&s);
                            arg_idx += 1;
                        }
                    } else if c == '\u{2}' {
                        // TAG_CONST
                        if const_idx < bsm_args.len() - 1 {
                            let const_val_idx = bsm_args[1 + const_idx];
                            let s = get_constant_value_as_string(class, const_val_idx);
                            result.push_str(&s);
                            const_idx += 1;
                        }
                    } else {
                        result.push(c);
                    }
                }

                // If there are remaining args (unlikely if recipe matches), append them?
                // Spec says recipe must match. We assume it matches.
            }

            let addr = self.memory.heap.allocate_string(result);
            frame.push(Value::Reference(addr))?;
        } else {
            return Err(to_runtime_error_enum(RuntimeError::Unimplemented(format!(
                "Unsupported invokedynamic bootstrap method: {}.{}",
                bsm_class, bsm_name
            ))));
        }

        Ok(())
    }
}