use endor_core::{
    CodeLayout, DecodedInstruction, EndorCode, EndorNumber, Instruction, ObjectRef, Opcode, Value,
    Vm, VmError,
};
use frontend::ast::{AstRoot, Decl, DeclKind, Expr, ExprKind, Stmt, StmtKind, UnaryOp};
use frontend::{Diagnostic, DiagnosticKind, SourceUnit, Span};

use crate::{NameBinding, Scope, ScopeKind, SymbolTable, diagnostic};

/// Validated, heap-independent executable template. Nested code constants
/// remain templates until load(), so no ObjectRef is shared across VM heaps.
#[derive(Debug, Clone)]
pub struct CompiledCode {
    code: EndorCode,
    children: Vec<(usize, CompiledCode)>,
}

impl CompiledCode {
    pub fn instructions(&self) -> &[DecodedInstruction] {
        self.code.instructions()
    }
    pub fn layout(&self) -> &CodeLayout {
        self.code.layout()
    }
    pub fn name(&self) -> &str {
        self.code.name()
    }
    pub fn filename(&self) -> &str {
        self.code.filename()
    }
    pub fn first_line(&self) -> usize {
        self.code.first_line()
    }
    pub fn children(&self) -> impl Iterator<Item = &CompiledCode> {
        self.children.iter().map(|(_, child)| child)
    }

    /// Return one owned code-object root; release it after execution. The
    /// parent retains its nested code constants. Temporary child roots are
    /// released on both success and failure, including partial recursive loads.
    pub fn load(&self, vm: &mut Vm) -> Result<ObjectRef, VmError> {
        let mut roots = Vec::new();
        let result = (|| {
            let mut constants = self.code.constants().to_vec();
            for (index, child) in &self.children {
                let reference = child.load(vm)?;
                roots.push(reference);
                constants[*index] = Value::Object(reference);
            }
            let code = EndorCode::with_layout(
                self.code.bytecode().to_vec(),
                constants,
                self.code.layout().clone(),
            )?
            .with_location(
                self.code.name(),
                self.code.filename(),
                self.code.first_line(),
            );
            vm.new_code(code)
        })();
        for reference in roots {
            vm.release(reference)
                .expect("loader owns each temporary code root");
        }
        result
    }
}

fn error(span: Span, message: impl Into<String>) -> Diagnostic {
    diagnostic(DiagnosticKind::Compilation, span, message)
}

pub(crate) fn generate(
    root: &AstRoot,
    source: &SourceUnit,
    symbols: &SymbolTable,
) -> Result<CompiledCode, Diagnostic> {
    let mut emitter = Emitter::new(&symbols.root, source);
    match root {
        AstRoot::Module(program) => {
            for declaration in &program.declarations {
                emitter.declaration(declaration, symbols)?;
            }
            emitter.constant(Value::Nil)?;
            emitter.simple(Opcode::ReturnValue);
        }
        AstRoot::Expression(expression) => {
            emitter.expression(expression)?;
            emitter.simple(Opcode::ReturnValue);
        }
        AstRoot::Interactive(input) => {
            for statement in &input.statements {
                if let StmtKind::Expression(expression) = &statement.kind {
                    emitter.expression(expression)?;
                } else {
                    emitter.statement(statement, symbols)?;
                    emitter.constant(Value::Nil)?;
                }
            }
            emitter.simple(Opcode::ReturnValue);
        }
    }
    emitter.finish()
}

struct Emitter<'a> {
    scope: &'a Scope,
    source: &'a SourceUnit,
    instructions: Vec<Instruction>,
    constants: Vec<Value>,
    children: Vec<(usize, CompiledCode)>,
}

impl<'a> Emitter<'a> {
    fn new(scope: &'a Scope, source: &'a SourceUnit) -> Self {
        Self {
            scope,
            source,
            instructions: Vec::new(),
            constants: Vec::new(),
            children: Vec::new(),
        }
    }

    fn emit(&mut self, opcode: Opcode, argument: usize) -> Result<(), Diagnostic> {
        let argument = u32::try_from(argument)
            .map_err(|_| error(self.scope.span, "bytecode operand exceeds u32"))?;
        self.instructions.push(Instruction::new(opcode, argument));
        Ok(())
    }

    fn simple(&mut self, opcode: Opcode) {
        self.instructions.push(Instruction::simple(opcode));
    }

    fn constant(&mut self, value: Value) -> Result<(), Diagnostic> {
        let index = self.constants.len();
        self.constants.push(value);
        self.emit(Opcode::LoadConst, index)
    }

    fn name_index(&self, name: &str) -> usize {
        self.scope
            .layout
            .names
            .iter()
            .position(|candidate| candidate == name)
            .expect("analysis interns every emitted name")
    }

    fn declaration(&mut self, declaration: &Decl, symbols: &SymbolTable) -> Result<(), Diagnostic> {
        let DeclKind::Function(function) = &declaration.kind;
        let scope = symbols
            .functions
            .iter()
            .find(|scope| scope.span == function.span && scope.name == function.name.text)
            .expect("analysis creates each function scope");
        let mut body = Emitter::new(scope, self.source);
        for statement in &function.body.statements {
            body.statement(statement, symbols)?;
        }
        if !matches!(
            function
                .body
                .statements
                .last()
                .map(|statement| &statement.kind),
            Some(StmtKind::Return(_))
        ) {
            body.constant(Value::Nil)?;
            body.simple(Opcode::ReturnValue);
        }
        let child = body.finish()?;
        let index = self.constants.len();
        self.children.push((index, child));
        self.constant(Value::Nil)?; // Replaced by a managed code object during load().
        self.instructions.push(Instruction::make_function(0, 0));
        self.emit(Opcode::StoreGlobal, self.name_index(&function.name.text))
    }

    fn store_binding(&mut self, name: &str, value: &Expr) -> Result<(), Diagnostic> {
        self.expression(value)?;
        if self.scope.kind == ScopeKind::Function {
            let index = self
                .scope
                .layout
                .locals
                .iter()
                .position(|local| local == name)
                .expect("analysis allocates binding and assignment targets");
            self.emit(Opcode::StoreFast, index)?;
        } else {
            self.emit(Opcode::StoreGlobal, self.name_index(name))?;
        }
        Ok(())
    }

    fn statement(&mut self, statement: &Stmt, symbols: &SymbolTable) -> Result<(), Diagnostic> {
        match &statement.kind {
            StmtKind::Binding(binding) => self.store_binding(&binding.name.text, &binding.value)?,
            StmtKind::Assignment(assignment) => {
                self.store_binding(&assignment.name.text, &assignment.value)?
            }
            StmtKind::Return(statement) => {
                if let Some(value) = &statement.value {
                    self.expression(value)?;
                } else {
                    self.constant(Value::Nil)?;
                }
                self.simple(Opcode::ReturnValue);
            }
            StmtKind::Expression(expression) => {
                self.expression(expression)?;
                self.simple(Opcode::PopTop);
            }
            StmtKind::Declaration(declaration) => self.declaration(declaration, symbols)?,
        }
        Ok(())
    }

    fn expression(&mut self, expression: &Expr) -> Result<(), Diagnostic> {
        match &expression.kind {
            ExprKind::String(value) => self.constant(Value::Text(value.clone()))?,
            ExprKind::Number(spelling) => self.number(spelling, expression.span)?,
            ExprKind::Name(name) => {
                let reference = self
                    .scope
                    .references
                    .iter()
                    .find(|reference| reference.span == expression.span && &reference.name == name)
                    .expect("analysis resolves every name expression");
                let (opcode, index) = match reference.binding {
                    NameBinding::Local(index) => (Opcode::LoadFast, index),
                    NameBinding::Global(index) => (Opcode::LoadGlobal, index),
                };
                self.emit(opcode, index)?;
            }
            ExprKind::Call { callee, arguments } => {
                self.expression(callee)?;
                self.simple(Opcode::PushNull);
                for argument in arguments {
                    self.expression(argument)?;
                }
                self.emit(Opcode::Call, arguments.len())?;
            }
            ExprKind::Member { object, name } => {
                self.expression(object)?;
                self.emit(Opcode::LoadAttr, self.name_index(&name.text))?;
            }
            ExprKind::Unary { .. } => {
                // Literal sign folding preserves arbitrarily large integers
                // and decimal spelling; this is not general numeric arithmetic.
                let Some(spelling) = signed_literal(expression, false) else {
                    return Err(error(
                        expression.span,
                        "numeric unary operations are not implemented; only signed number literals are supported",
                    ));
                };
                self.number(&spelling, expression.span)?;
            }
            ExprKind::Binary { .. } => {
                return Err(error(
                    expression.span,
                    "numeric arithmetic bytecode is not implemented yet",
                ));
            }
            ExprKind::Symbol(_) => {
                return Err(error(
                    expression.span,
                    "symbol runtime values are not implemented yet",
                ));
            }
        }
        Ok(())
    }

    fn number(&mut self, spelling: &str, span: Span) -> Result<(), Diagnostic> {
        let value = if spelling.contains('.') {
            EndorNumber::decimal(spelling)
        } else {
            EndorNumber::integer(spelling)
        }
        .map_err(|failure| error(span, failure.to_string()))?;
        self.constant(Value::Number(value))
    }

    fn finish(self) -> Result<CompiledCode, Diagnostic> {
        let line = self
            .scope
            .span
            .location(self.source.text())
            .map_or(1, |location| location.line);
        let code = EndorCode::from_instructions(
            &self.instructions,
            self.constants,
            self.scope.layout.clone(),
        )
        .map_err(|failure| error(self.scope.span, failure.to_string()))?
        .with_location(&self.scope.name, self.source.origin().to_string(), line);
        Ok(CompiledCode {
            code,
            children: self.children,
        })
    }
}

fn signed_literal(expression: &Expr, negative: bool) -> Option<String> {
    match &expression.kind {
        ExprKind::Number(spelling) => Some(if negative {
            format!("-{spelling}")
        } else {
            spelling.clone()
        }),
        ExprKind::Unary {
            operation: UnaryOp::Negate,
            operand,
        } => signed_literal(operand, !negative),
        _ => None,
    }
}