//! On-Stack Replacement (OSR) for long-running loops.
//!
//! OSR allows the JVM to replace an executing interpreted method with a
//! JIT-compiled version mid-execution, without waiting for the method to
//! return. This is critical for methods with long-running loops that would
//! otherwise never be compiled.

use crate::class_file::MethodInfo;
use crate::jit::CompilationLevel;
use crate::memory::{StackFrame, Value};
use std::collections::HashMap;

/// OSR entry point: the bytecode offset at which OSR can be triggered
#[derive(Debug, Clone)]
pub struct OsrEntryPoint {
    /// Bytecode offset (typically a loop back-edge target)
    pub bytecode_offset: usize,
    /// Number of times this back-edge has been taken
    pub back_edge_count: u64,
    /// Whether a compiled version is available at this entry
    pub compiled: bool,
    /// Compilation level of the OSR stub
    pub level: CompilationLevel,
}

/// OSR trigger configuration
#[derive(Debug, Clone)]
pub struct OsrConfig {
    /// Number of back-edge iterations before triggering OSR compilation
    pub back_edge_threshold: u64,
    /// Whether OSR is enabled
    pub enabled: bool,
    /// Maximum number of OSR entry points per method
    pub max_entry_points: usize,
}

impl Default for OsrConfig {
    fn default() -> Self {
        Self {
            back_edge_threshold: 10_000,
            enabled: true,
            max_entry_points: 8,
        }
    }
}

/// State transfer for OSR: maps interpreter locals/stack to compiled frame
#[derive(Debug, Clone)]
pub struct OsrStateTransfer {
    /// Local variable values at the OSR point
    pub locals: Vec<Value>,
    /// Operand stack values at the OSR point
    pub stack: Vec<Value>,
    /// Program counter at the OSR point
    pub pc: usize,
}

impl OsrStateTransfer {
    /// Capture current interpreter state for OSR transfer
    pub fn capture(frame: &StackFrame, pc: usize) -> Self {
        Self {
            locals: frame.locals.clone(),
            stack: frame.stack.clone(),
            pc,
        }
    }
}

/// OSR manager: tracks back-edges and manages OSR compilation
pub struct OsrManager {
    /// Config
    config: OsrConfig,
    /// Back-edge counters: method_key -> (bytecode_offset -> count)
    back_edge_counters: HashMap<String, HashMap<usize, u64>>,
    /// OSR entry points: method_key -> Vec<OsrEntryPoint>
    entry_points: HashMap<String, Vec<OsrEntryPoint>>,
    /// Total OSR compilations triggered
    osr_compilations: u64,
    /// Total OSR transitions performed
    osr_transitions: u64,
}

impl OsrManager {
    pub fn new() -> Self {
        Self::with_config(OsrConfig::default())
    }

    pub fn with_config(config: OsrConfig) -> Self {
        Self {
            config,
            back_edge_counters: HashMap::new(),
            entry_points: HashMap::new(),
            osr_compilations: 0,
            osr_transitions: 0,
        }
    }

    /// Record a back-edge taken at `target_pc` in the given method.
    /// Returns `Some(target_pc)` if OSR should be triggered.
    pub fn record_back_edge(
        &mut self,
        class_name: &str,
        method_name: &str,
        target_pc: usize,
    ) -> Option<usize> {
        if !self.config.enabled {
            return None;
        }

        let key = format!("{}.{}", class_name, method_name);
        let counter = self
            .back_edge_counters
            .entry(key.clone())
            .or_default()
            .entry(target_pc)
            .or_insert(0);
        *counter += 1;

        if *counter >= self.config.back_edge_threshold {
            // Check if we already have an entry point here
            let entry_points = self.entry_points.entry(key).or_default();
            let already_compiled = entry_points
                .iter()
                .any(|ep| ep.bytecode_offset == target_pc && ep.compiled);

            if !already_compiled
                && entry_points.len() < self.config.max_entry_points
            {
                entry_points.push(OsrEntryPoint {
                    bytecode_offset: target_pc,
                    back_edge_count: *counter,
                    compiled: false,
                    level: CompilationLevel::Baseline,
                });
                return Some(target_pc);
            }
        }

        None
    }

    /// Mark an OSR entry point as compiled
    pub fn mark_compiled(
        &mut self,
        class_name: &str,
        method_name: &str,
        bytecode_offset: usize,
        level: CompilationLevel,
    ) {
        let key = format!("{}.{}", class_name, method_name);
        if let Some(entry_points) = self.entry_points.get_mut(&key) {
            for ep in entry_points.iter_mut() {
                if ep.bytecode_offset == bytecode_offset {
                    ep.compiled = true;
                    ep.level = level;
                    self.osr_compilations += 1;
                    break;
                }
            }
        }
    }

    /// Check if an OSR entry point is available at the given offset
    pub fn has_osr_entry(
        &self,
        class_name: &str,
        method_name: &str,
        bytecode_offset: usize,
    ) -> bool {
        let key = format!("{}.{}", class_name, method_name);
        self.entry_points
            .get(&key)
            .map(|eps| {
                eps.iter()
                    .any(|ep| ep.bytecode_offset == bytecode_offset && ep.compiled)
            })
            .unwrap_or(false)
    }

    /// Perform OSR transition: transfer state and return the new frame
    pub fn perform_osr_transition(
        &mut self,
        class_name: &str,
        method_name: &str,
        frame: &StackFrame,
        pc: usize,
    ) -> Option<OsrStateTransfer> {
        if !self.has_osr_entry(class_name, method_name, pc) {
            return None;
        }

        self.osr_transitions += 1;
        Some(OsrStateTransfer::capture(frame, pc))
    }

    /// Scan bytecode for back-edge targets (loop headers)
    pub fn find_back_edges(method: &MethodInfo) -> Vec<usize> {
        let bytecode = Self::extract_bytecode(method);
        let mut back_edges = Vec::new();
        let mut pc = 0usize;

        while pc < bytecode.len() {
            let opcode = bytecode[pc];
            pc += 1;

            match opcode {
                // Conditional branches
                0x99..=0xa4 => {
                    if pc + 1 < bytecode.len() {
                        let offset =
                            (((bytecode[pc] as i8) as i16) << 8) | (bytecode[pc + 1] as i16);
                        pc += 2;
                        // Negative offset = backward branch = loop back-edge
                        if offset < 0 {
                            let target = (pc as i64 - 3 + offset as i64) as usize;
                            back_edges.push(target);
                        }
                    }
                }
                // goto
                0xa7 => {
                    if pc + 1 < bytecode.len() {
                        let offset =
                            (((bytecode[pc] as i8) as i16) << 8) | (bytecode[pc + 1] as i16);
                        pc += 2;
                        if offset < 0 {
                            let target = (pc as i64 - 3 + offset as i64) as usize;
                            back_edges.push(target);
                        }
                    }
                }
                // goto_w
                0xc8 => {
                    pc += 4;
                }
                // Instructions with 1 operand byte
                0x10 | 0x12 | 0x15..=0x19 | 0x36..=0x3a | 0xa9 | 0xbc => {
                    pc += 1;
                }
                // Instructions with 2 operand bytes
                0x11 | 0x13 | 0x14 | 0x84 | 0xb2..=0xb8 | 0xbb | 0xbd | 0xc0 | 0xc1
                | 0xc6 | 0xc7 => {
                    pc += 2;
                }
                // invokeinterface / invokedynamic
                0xb9 | 0xba => {
                    pc += 4;
                }
                // multianewarray
                0xc5 => {
                    pc += 3;
                }
                _ => {}
            }
        }

        back_edges.sort_unstable();
        back_edges.dedup();
        back_edges
    }

    fn extract_bytecode(method: &MethodInfo) -> Vec<u8> {
        for attr in &method.attributes {
            if attr.info.len() >= 8 {
                let code_len = ((attr.info[4] as usize) << 24)
                    | ((attr.info[5] as usize) << 16)
                    | ((attr.info[6] as usize) << 8)
                    | (attr.info[7] as usize);
                if let Some(code) = attr.info.get(8..8 + code_len) {
                    return code.to_vec();
                }
            }
        }
        Vec::new()
    }

    /// Get OSR statistics
    pub fn stats(&self) -> OsrStats {
        let total_entry_points: usize =
            self.entry_points.values().map(|v| v.len()).sum();
        let compiled_entry_points: usize = self
            .entry_points
            .values()
            .flat_map(|v| v.iter())
            .filter(|ep| ep.compiled)
            .count();

        OsrStats {
            methods_with_back_edges: self.back_edge_counters.len(),
            total_entry_points,
            compiled_entry_points,
            osr_compilations: self.osr_compilations,
            osr_transitions: self.osr_transitions,
        }
    }
}

impl Default for OsrManager {
    fn default() -> Self {
        Self::new()
    }
}

/// OSR statistics
#[derive(Debug, Clone)]
pub struct OsrStats {
    pub methods_with_back_edges: usize,
    pub total_entry_points: usize,
    pub compiled_entry_points: usize,
    pub osr_compilations: u64,
    pub osr_transitions: u64,
}

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

    #[test]
    fn test_osr_manager_creation() {
        let manager = OsrManager::new();
        assert!(manager.config.enabled);
        assert_eq!(manager.osr_compilations, 0);
    }

    #[test]
    fn test_back_edge_recording() {
        let mut manager = OsrManager::with_config(OsrConfig {
            back_edge_threshold: 3,
            enabled: true,
            max_entry_points: 4,
        });

        // First two back-edges should not trigger OSR
        assert!(manager.record_back_edge("Test", "loop", 10).is_none());
        assert!(manager.record_back_edge("Test", "loop", 10).is_none());
        // Third should trigger
        assert_eq!(
            manager.record_back_edge("Test", "loop", 10),
            Some(10)
        );
    }

    #[test]
    fn test_osr_mark_compiled() {
        let mut manager = OsrManager::with_config(OsrConfig {
            back_edge_threshold: 1,
            enabled: true,
            max_entry_points: 4,
        });

        manager.record_back_edge("Test", "loop", 5);
        manager.mark_compiled("Test", "loop", 5, CompilationLevel::Baseline);

        assert!(manager.has_osr_entry("Test", "loop", 5));
    }

    #[test]
    fn test_osr_disabled() {
        let mut manager = OsrManager::with_config(OsrConfig {
            back_edge_threshold: 1,
            enabled: false,
            max_entry_points: 4,
        });

        let result = manager.record_back_edge("Test", "loop", 5);
        assert!(result.is_none());
    }

    #[test]
    fn test_osr_state_transfer_capture() {
        let mut frame = StackFrame::new(4, 8, "test".to_string());
        frame.locals[0] = Value::Int(42);
        frame.push(Value::Int(100)).unwrap();

        let transfer = OsrStateTransfer::capture(&frame, 10);
        assert_eq!(transfer.pc, 10);
        assert_eq!(transfer.locals[0], Value::Int(42));
        assert_eq!(transfer.stack[0], Value::Int(100));
    }
}
