use std::any::Any;
use std::collections::{BTreeSet, VecDeque};
use std::sync::Arc;

use crate::builtins::trace_value;
use crate::instruction::{compare_operation, decode, function_counts, invalid};
use crate::{DecodedInstruction, Instruction, ObjectPayload, ObjectRef, Opcode, Value, VmError};

/// Local slots start with positional parameters. Captured locals occupy their
/// original slots as cells; free-variable cells follow all local slots.
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct CodeLayout {
    pub parameters: Vec<String>,
    pub locals: Vec<String>,
    pub cell_variables: Vec<String>,
    pub free_variables: Vec<String>,
    pub names: Vec<String>,
}

#[derive(Debug, Clone)]
pub struct EndorCode {
    bytecode: Vec<u8>,
    instructions: Vec<DecodedInstruction>,
    constants: Vec<Value>,
    layout: CodeLayout,
    stack_size: usize,
    name: String,
    filename: String,
    first_line: usize,
}

impl EndorCode {
    pub fn new(
        bytecode: Vec<u8>,
        constants: Vec<Value>,
        parameters: Vec<String>,
        free_variables: Vec<String>,
    ) -> Result<Self, VmError> {
        Self::with_layout(
            bytecode,
            constants,
            CodeLayout {
                locals: parameters.clone(),
                parameters,
                free_variables,
                ..CodeLayout::default()
            },
        )
    }

    pub fn from_instructions(
        instructions: &[Instruction],
        constants: Vec<Value>,
        layout: CodeLayout,
    ) -> Result<Self, VmError> {
        Self::with_layout(Instruction::encode_all(instructions)?, constants, layout)
    }

    pub fn with_layout(
        bytecode: Vec<u8>,
        constants: Vec<Value>,
        layout: CodeLayout,
    ) -> Result<Self, VmError> {
        validate_layout(&layout)?;
        let instructions = decode(&bytecode)?;
        let stack_size = verify(&instructions, constants.len(), &layout)?;
        Ok(Self {
            bytecode,
            instructions,
            constants,
            layout,
            stack_size,
            name: "<code>".into(),
            filename: "<endor>".into(),
            first_line: 1,
        })
    }

    pub fn with_location(
        mut self,
        name: impl Into<String>,
        filename: impl Into<String>,
        first_line: usize,
    ) -> Self {
        self.name = name.into();
        self.filename = filename.into();
        self.first_line = first_line;
        self
    }

    pub fn bytecode(&self) -> &[u8] {
        &self.bytecode
    }
    pub fn instructions(&self) -> &[DecodedInstruction] {
        &self.instructions
    }
    pub fn constants(&self) -> &[Value] {
        &self.constants
    }
    pub fn parameters(&self) -> &[String] {
        &self.layout.parameters
    }
    pub fn free_variables(&self) -> &[String] {
        &self.layout.free_variables
    }
    pub fn layout(&self) -> &CodeLayout {
        &self.layout
    }
    pub fn stack_size(&self) -> usize {
        self.stack_size
    }
    pub fn name(&self) -> &str {
        &self.name
    }
    pub fn filename(&self) -> &str {
        &self.filename
    }
    pub fn first_line(&self) -> usize {
        self.first_line
    }
    pub fn slot_count(&self) -> usize {
        self.layout.locals.len() + self.layout.free_variables.len()
    }
    /// Whether a local/free slot holds a managed closure cell.
    pub fn is_cell_slot(&self, slot: usize) -> bool {
        is_cell_slot(&self.layout, slot)
    }
    /// Name of a local/free slot used by validated instructions.
    ///
    /// # Panics
    /// Panics if the slot is outside this code's local/free-variable layout.
    pub fn slot_name(&self, slot: usize) -> &str {
        if slot < self.layout.locals.len() {
            &self.layout.locals[slot]
        } else {
            &self.layout.free_variables[slot - self.layout.locals.len()]
        }
    }
    /// Look up a logical instruction by its code-unit offset.
    pub fn instruction_at(&self, offset: usize) -> Result<DecodedInstruction, VmError> {
        instruction_index(&self.instructions, offset)
            .map(|index| self.instructions[index])
            .ok_or_else(|| {
                invalid(
                    offset,
                    "instruction pointer is not on an instruction boundary",
                )
            })
    }
}

/// Heap-allocated, immutable and shareable executable. Function objects hold a
/// managed reference to this object rather than embedding a copy of its code.
#[derive(Debug, Clone)]
pub struct EndorCodeObject {
    code: Arc<EndorCode>,
}

impl EndorCodeObject {
    pub fn new(code: EndorCode) -> Self {
        Self {
            code: Arc::new(code),
        }
    }
    pub fn code(&self) -> &EndorCode {
        &self.code
    }
    /// Share immutable executable data without copying constants or wordcode.
    /// This Rust Arc does not retain heap objects referenced by constants;
    /// callers must keep the managed code object rooted while using them.
    pub fn shared_code(&self) -> Arc<EndorCode> {
        Arc::clone(&self.code)
    }
}

impl ObjectPayload for EndorCodeObject {
    fn reference_assignment_only(&self) -> bool {
        true
    }
    fn as_any(&self) -> &dyn Any {
        self
    }
    fn kind_name(&self) -> &'static str {
        "code"
    }
    fn trace(&self, visit: &mut dyn FnMut(ObjectRef)) {
        for value in self.code.constants() {
            trace_value(value, visit);
        }
    }
    fn repr(&self, _identity: ObjectRef) -> Result<String, VmError> {
        Ok(format!(
            "<code {} at {}:{}>",
            self.code.name(),
            self.code.filename(),
            self.code.first_line()
        ))
    }
}

fn validate_layout(layout: &CodeLayout) -> Result<(), VmError> {
    for names in [
        &layout.parameters,
        &layout.locals,
        &layout.cell_variables,
        &layout.free_variables,
        &layout.names,
    ] {
        let mut seen = BTreeSet::new();
        for name in names {
            if name.is_empty() || !seen.insert(name) {
                return Err(VmError::InvalidArgument(
                    "empty or duplicate code variable name",
                ));
            }
        }
    }
    if !layout.locals.starts_with(&layout.parameters) {
        return Err(VmError::InvalidArgument(
            "parameter names must lead local slots",
        ));
    }
    if layout
        .cell_variables
        .iter()
        .any(|name| !layout.locals.contains(name))
        || layout
            .free_variables
            .iter()
            .any(|name| layout.locals.contains(name))
    {
        return Err(VmError::InvalidArgument(
            "invalid local/cell/free-variable layout",
        ));
    }
    Ok(())
}

fn is_cell_slot(layout: &CodeLayout, slot: usize) -> bool {
    if slot >= layout.locals.len() {
        slot < layout.locals.len() + layout.free_variables.len()
    } else {
        layout.cell_variables.contains(&layout.locals[slot])
    }
}

fn instruction_index(instructions: &[DecodedInstruction], offset: usize) -> Option<usize> {
    instructions
        .binary_search_by_key(&offset, |instruction| instruction.offset)
        .ok()
}

pub(crate) fn jump_target(instruction: DecodedInstruction) -> Result<usize, VmError> {
    let argument = instruction.instruction.argument as usize;
    let target = if instruction.instruction.opcode == Opcode::JumpBackward {
        instruction.next_offset.checked_sub(argument)
    } else {
        instruction.next_offset.checked_add(argument)
    };
    target.ok_or_else(|| invalid(instruction.offset, "jump offset overflow"))
}

/// Validate every operand, then propagate stack heights through reachable
/// control flow. Bad jumps, underflow, mismatched joins and fallthrough fail
/// before execution; loops must return to the same stack height.
fn verify(
    instructions: &[DecodedInstruction],
    constants: usize,
    layout: &CodeLayout,
) -> Result<usize, VmError> {
    for instruction in instructions {
        let Instruction { opcode, argument } = instruction.instruction;
        let index = argument as usize;
        let valid = match opcode {
            Opcode::LoadConst => index < constants,
            Opcode::LoadFast | Opcode::StoreFast | Opcode::DeleteFast => {
                index < layout.locals.len() && !is_cell_slot(layout, index)
            }
            Opcode::LoadDeref | Opcode::StoreDeref | Opcode::DeleteDeref | Opcode::LoadClosure => {
                is_cell_slot(layout, index)
            }
            Opcode::LoadGlobal
            | Opcode::StoreGlobal
            | Opcode::DeleteGlobal
            | Opcode::LoadAttr
            | Opcode::StoreAttr => index < layout.names.len(),
            Opcode::CompareOp => compare_operation(argument).is_some(),
            Opcode::Copy => argument >= 1,
            Opcode::Swap => argument >= 2,
            Opcode::JumpForward
            | Opcode::JumpBackward
            | Opcode::PopJumpIfFalse
            | Opcode::PopJumpIfTrue => {
                instruction_index(instructions, jump_target(*instruction)?).is_some()
            }
            Opcode::ExtendedArg => false,
            _ => true,
        };
        if !valid {
            return Err(invalid(
                instruction.offset,
                "invalid operand or jump target",
            ));
        }
    }
    let mut heights = vec![None; instructions.len()];
    heights[0] = Some(0_usize);
    let mut pending = VecDeque::from([0]);
    let mut maximum = 0;
    while let Some(index) = pending.pop_front() {
        let instruction = instructions[index];
        let height = heights[index].unwrap();
        let Instruction { opcode, argument } = instruction.instruction;
        let argument = argument as usize;
        let (required, popped, pushed) = match opcode {
            Opcode::LoadConst
            | Opcode::LoadFast
            | Opcode::LoadGlobal
            | Opcode::LoadDeref
            | Opcode::LoadClosure
            | Opcode::PushNull => (0, 0, 1),
            Opcode::StoreFast
            | Opcode::StoreGlobal
            | Opcode::PopTop
            | Opcode::StoreDeref
            | Opcode::PopJumpIfFalse
            | Opcode::PopJumpIfTrue => (1, 1, 0),
            Opcode::ToBool | Opcode::UnaryNot | Opcode::LoadAttr => (1, 1, 1),
            Opcode::CompareOp => (2, 2, 1),
            Opcode::StoreAttr => (2, 2, 0),
            Opcode::Copy => (argument, 0, 1),
            Opcode::Swap => (argument, 0, 0),
            Opcode::Call => {
                let consumed = argument
                    .checked_add(2)
                    .ok_or_else(|| invalid(instruction.offset, "call operand overflow"))?;
                (consumed, consumed, 1)
            }
            Opcode::BuildMap => {
                let consumed = argument
                    .checked_mul(2)
                    .ok_or_else(|| invalid(instruction.offset, "map operand overflow"))?;
                (consumed, consumed, 1)
            }
            Opcode::MakeFunction => {
                let (defaults, closure) = function_counts(instruction.instruction.argument);
                let consumed = defaults + closure + 1;
                (consumed, consumed, 1)
            }
            Opcode::BuildClass => (4, 4, 1),
            Opcode::ReturnValue => (1, 1, 0),
            _ => (0, 0, 0),
        };
        if height < required {
            return Err(invalid(instruction.offset, "evaluation stack underflow"));
        }
        if opcode == Opcode::ReturnValue && height != 1 {
            return Err(invalid(
                instruction.offset,
                "RETURN_VALUE must leave an empty frame stack",
            ));
        }
        let next_height = height - popped + pushed;
        maximum = maximum.max(height).max(next_height);
        let successors: Vec<usize> = match opcode {
            Opcode::ReturnValue => Vec::new(),
            Opcode::JumpForward | Opcode::JumpBackward => vec![jump_target(instruction)?],
            Opcode::PopJumpIfFalse | Opcode::PopJumpIfTrue => {
                vec![jump_target(instruction)?, instruction.next_offset]
            }
            _ => vec![instruction.next_offset],
        };
        for successor in successors {
            let target = instruction_index(instructions, successor).ok_or_else(|| {
                invalid(instruction.offset, "execution falls off the code object")
            })?;
            match heights[target] {
                Some(existing) if existing != next_height => {
                    return Err(invalid(
                        successor,
                        "control-flow paths have different stack heights",
                    ));
                }
                Some(_) => {}
                None => {
                    heights[target] = Some(next_height);
                    pending.push_back(target);
                }
            }
        }
    }
    Ok(maximum)
}