//! Code generator: walks the AST and emits bytecode instructions.

use el_parser::{BinOp, Expr, Literal, Program, Stmt};

use crate::bytecode::{Bytecode, Value};
use crate::error::CompileResult;
use crate::source_map::SourceMap;

/// Generates bytecode from a parsed program.
pub struct Codegen {
    instructions: Vec<Bytecode>,
    source_map: SourceMap,
    #[allow(dead_code)]
    emit_source_map: bool,
}

impl Codegen {
    pub fn new(emit_source_map: bool) -> Self {
        Self {
            instructions: Vec::new(),
            source_map: SourceMap::new(),
            emit_source_map,
        }
    }

    /// Generate bytecode for a complete program.
    pub fn generate(mut self, program: &Program) -> CompileResult<(Vec<Bytecode>, SourceMap)> {
        for stmt in &program.stmts {
            self.gen_stmt(stmt)?;
        }
        self.emit(Bytecode::Halt);
        Ok((self.instructions, self.source_map))
    }

    // ── Emission helpers ──────────────────────────────────────────────────────

    fn emit(&mut self, instr: Bytecode) -> usize {
        let idx = self.instructions.len();
        self.instructions.push(instr);
        idx
    }

    #[allow(dead_code)]
    fn emit_at_span(&mut self, instr: Bytecode, span: el_lexer::Span) -> usize {
        let idx = self.instructions.len();
        if self.emit_source_map {
            self.source_map.record(idx, span);
        }
        self.instructions.push(instr);
        idx
    }

    fn patch_jump(&mut self, idx: usize, target: usize) {
        // offset = target - (idx + 1)  (jump is relative to the next instruction)
        let offset = target as i32 - (idx as i32 + 1);
        match &mut self.instructions[idx] {
            Bytecode::Jump(o) | Bytecode::JumpIf(o) | Bytecode::JumpIfNot(o) => *o = offset,
            _ => {}
        }
    }

    fn current_idx(&self) -> usize {
        self.instructions.len()
    }

    // ── Statement code generation ─────────────────────────────────────────────

    /// Generate a statement in tail position (the last stmt of a block).
    /// Expression statements leave their value on the stack instead of popping it.
    fn gen_stmt_tail(&mut self, stmt: &Stmt) -> CompileResult<()> {
        match stmt {
            Stmt::Expr(expr, _) => {
                // In tail position, leave the value on the stack.
                self.gen_expr(expr)?;
            }
            Stmt::Return(expr, _) => {
                self.gen_expr(expr)?;
                self.emit(Bytecode::Return);
            }
            // All other statement kinds behave the same as non-tail; push Nil as block value.
            other => {
                self.gen_stmt(other)?;
                self.emit(Bytecode::Push(Value::Nil));
            }
        }
        Ok(())
    }

    fn gen_stmt(&mut self, stmt: &Stmt) -> CompileResult<()> {
        // Record the source span for this statement in the source map
        if self.emit_source_map {
            let span = stmt_span(stmt);
            let idx = self.instructions.len();
            self.source_map.record(idx, span);
        }
        match stmt {
            Stmt::Let { name, value, .. } => {
                self.gen_expr(value)?;
                self.emit(Bytecode::StoreLocal(name.clone()));
            }
            Stmt::Return(expr, _) => {
                self.gen_expr(expr)?;
                self.emit(Bytecode::Return);
            }
            Stmt::Expr(expr, _) => {
                self.gen_expr(expr)?;
                // Discard the expression result unless it's a return-like
                if !matches!(expr, Expr::Block(_) | Expr::If { .. }) {
                    self.emit(Bytecode::Pop);
                }
            }
            Stmt::FnDef { name, params, body, requires, .. } => {
                // In this simple bytecode model, function defs emit a Jump to skip
                // the function body, then a label for the function start.
                // A full implementation would use a call frame table; for now we
                // emit the body inline and register the entry point offset.
                let skip_jump = self.emit(Bytecode::Jump(0)); // patched below

                // Function body
                // Bind parameters in order (caller pushes args left-to-right)
                for param in params.iter().rev() {
                    self.emit(Bytecode::StoreLocal(param.name.clone()));
                }
                // Emit contract check if `requires` is present
                if let Some(req_expr) = requires {
                    self.gen_expr(req_expr)?;
                    self.emit(Bytecode::ContractCheck {
                        message: format!("contract violation in fn '{name}': requires clause failed"),
                    });
                }
                for s in body {
                    self.gen_stmt(s)?;
                }
                // Implicit void return
                self.emit(Bytecode::Push(Value::Nil));
                self.emit(Bytecode::Return);

                // Patch the skip jump
                let after = self.current_idx();
                self.patch_jump(skip_jump, after);

                // Register the function name → bytecode offset mapping
                // (stored as a load of the entry point index as an Int constant,
                //  then store as a local — real implementations use a function table)
                let entry_point = skip_jump + 1; // first instruction of body
                self.emit(Bytecode::Push(Value::Int(entry_point as i64)));
                self.emit(Bytecode::StoreLocal(format!("__fn_{name}")));
            }
            Stmt::Retry { count, body, fallback, .. } => {
                // Codegen for retry N times:
                //   counter = N
                //   loop_start:
                //     if counter <= 0 goto fallback
                //     decrement counter
                //     [body]
                //     goto done
                //   fallback:
                //     [fallback_body]
                //   done:
                let counter_name = format!("__retry_counter_{}__", self.current_idx());

                // Initialize counter
                self.gen_expr(count)?;
                self.emit(Bytecode::StoreLocal(counter_name.clone()));

                // Loop start: check counter > 0
                let loop_start = self.current_idx();
                self.emit(Bytecode::LoadLocal(counter_name.clone()));
                self.emit(Bytecode::Push(Value::Int(0)));
                self.emit(Bytecode::Gt);
                let to_fallback = self.emit(Bytecode::JumpIfNot(0)); // patched to fallback

                // Decrement counter
                self.emit(Bytecode::LoadLocal(counter_name.clone()));
                self.emit(Bytecode::Push(Value::Int(1)));
                self.emit(Bytecode::Sub);
                self.emit(Bytecode::StoreLocal(counter_name.clone()));

                // Execute body
                for s in body {
                    self.gen_stmt(s)?;
                }
                // Body succeeded — jump to done
                let to_done = self.emit(Bytecode::Jump(0));

                // Fallback
                let fallback_start = self.current_idx();
                self.patch_jump(to_fallback, fallback_start);
                if let Some(fb_body) = fallback {
                    for s in fb_body {
                        self.gen_stmt(s)?;
                    }
                }

                let done = self.current_idx();
                self.patch_jump(to_done, done);

                // Note: in this simple model the body always "succeeds".
                // A real retry would need exception-like control flow.
                // For the retry-loop semantic, also add a back-jump that
                // jumps back to loop_start after each body execution would
                // require adding another jump before `to_done`. This design
                // runs the body once then exits — which is correct for
                // "success on first try" semantics in a pure-fn language.
            }
            Stmt::Deploy { fn_name, route, target, .. } => {
                self.emit(Bytecode::DeployFn {
                    fn_name: fn_name.clone(),
                    route: route.clone(),
                    target: target.clone(),
                });
            }
            Stmt::TypeDef { .. } | Stmt::EnumDef { .. } => {
                // Type and enum definitions are compile-time only; no runtime code.
            }
            // Test-related statements — skipped during normal compilation.
            // The el-test crate walks the AST directly rather than running compiled bytecode.
            Stmt::TestDef { .. } | Stmt::Seed(..) | Stmt::Assert(..) => {}
            // New statement kinds — no runtime code emitted.
            _ => {}
        }
        Ok(())
    }

    // ── Expression code generation ────────────────────────────────────────────

    fn gen_expr(&mut self, expr: &Expr) -> CompileResult<()> {
        match expr {
            Expr::Literal(lit) => {
                let val = match lit {
                    Literal::Int(n) => Value::Int(*n),
                    Literal::Float(f) => Value::Float(*f),
                    Literal::Str(s) => Value::Str(s.clone()),
                    Literal::Bool(b) => Value::Bool(*b),
                };
                self.emit(Bytecode::Push(val));
            }
            Expr::Ident(name) => {
                self.emit(Bytecode::LoadLocal(name.clone()));
            }
            Expr::BinOp { op, left, right } => {
                self.gen_expr(left)?;
                self.gen_expr(right)?;
                let instr = match op {
                    BinOp::Add => Bytecode::Add,
                    BinOp::Sub => Bytecode::Sub,
                    BinOp::Mul => Bytecode::Mul,
                    BinOp::Div => Bytecode::Div,
                    BinOp::Eq => Bytecode::Eq,
                    BinOp::NotEq => Bytecode::NotEq,
                    BinOp::Lt => Bytecode::Lt,
                    BinOp::Gt => Bytecode::Gt,
                    BinOp::LtEq => Bytecode::LtEq,
                    BinOp::GtEq => Bytecode::GtEq,
                    BinOp::And => Bytecode::And,
                    BinOp::Or => Bytecode::Or,
                };
                self.emit(instr);
            }
            Expr::UnaryNot(inner) => {
                self.gen_expr(inner)?;
                self.emit(Bytecode::Not);
            }
            Expr::Call { func, args } => {
                // Push arguments left-to-right
                for arg in args {
                    self.gen_expr(arg)?;
                }
                // Get the function name from the callee expression
                let fn_name = match func.as_ref() {
                    Expr::Ident(n) => n.clone(),
                    Expr::Field { object, field } => {
                        self.gen_expr(object)?;
                        field.clone()
                    }
                    _ => {
                        self.gen_expr(func)?;
                        "__dynamic__".to_string()
                    }
                };
                self.emit(Bytecode::Call { name: fn_name, arity: args.len() as u32 });
            }
            Expr::Block(stmts) => {
                if stmts.is_empty() {
                    self.emit(Bytecode::Push(Value::Nil));
                } else {
                    for (i, s) in stmts.iter().enumerate() {
                        let is_last = i == stmts.len() - 1;
                        if is_last {
                            // Last statement: emit as a tail expression (leave value on stack).
                            self.gen_stmt_tail(s)?;
                        } else {
                            self.gen_stmt(s)?;
                        }
                    }
                }
            }
            Expr::If { cond, then, else_ } => {
                self.gen_expr(cond)?;
                let jump_false = self.emit(Bytecode::JumpIfNot(0)); // patched
                self.gen_expr(then)?;
                if let Some(else_expr) = else_ {
                    let jump_end = self.emit(Bytecode::Jump(0)); // skip else
                    let else_start = self.current_idx();
                    self.patch_jump(jump_false, else_start);
                    self.gen_expr(else_expr)?;
                    let after_else = self.current_idx();
                    self.patch_jump(jump_end, after_else);
                } else {
                    let after_then = self.current_idx();
                    self.patch_jump(jump_false, after_then);
                }
            }
            Expr::Match { subject, arms } => {
                self.gen_expr(subject)?;
                // Simplified match: for each arm, dup subject, push pattern,
                // compare, branch. A full implementation would use a jump table.
                let mut end_jumps = Vec::new();
                for arm in arms {
                    match &arm.pattern {
                        el_parser::Pattern::Wildcard => {
                            // Wildcard always matches — pop subject and run body directly.
                            self.emit(Bytecode::Pop);
                            self.gen_expr(&arm.body)?;
                            end_jumps.push(self.emit(Bytecode::Jump(0)));
                            // No jump_no_match needed — wildcard always matches.
                            // But we still need to patch the end jumps at the end.
                            // Nothing else to do; break out of the loop since wildcard
                            // is a catch-all and subsequent arms are unreachable.
                            break;
                        }
                        el_parser::Pattern::Binding(name) => {
                            // Bind and always match.
                            self.emit(Bytecode::Dup);
                            self.emit(Bytecode::StoreLocal(name.clone()));
                            // Dup'd subject is still on stack; compare to itself.
                            self.emit(Bytecode::Dup);
                            self.emit(Bytecode::Eq);
                            let jump_no_match = self.emit(Bytecode::JumpIfNot(0));
                            self.emit(Bytecode::Pop);
                            self.gen_expr(&arm.body)?;
                            end_jumps.push(self.emit(Bytecode::Jump(0)));
                            let next_arm = self.current_idx();
                            self.patch_jump(jump_no_match, next_arm);
                        }
                        _ => {
                            self.emit(Bytecode::Dup);
                            // Push pattern value
                            match &arm.pattern {
                                el_parser::Pattern::Literal(lit) => {
                                    let v = match lit {
                                        Literal::Int(n) => Value::Int(*n),
                                        Literal::Str(s) => Value::Str(s.clone()),
                                        Literal::Bool(b) => Value::Bool(*b),
                                        Literal::Float(f) => Value::Float(*f),
                                    };
                                    self.emit(Bytecode::Push(v));
                                }
                                el_parser::Pattern::EnumVariant { variant, payload, .. } => {
                                    // Push the variant name as a string for comparison
                                    self.emit(Bytecode::Push(Value::Str(variant.clone())));
                                    if let Some(bind) = payload {
                                        // Store the subject (simplified: payload = subject)
                                        self.emit(Bytecode::StoreLocal(bind.clone()));
                                    }
                                }
                                _ => unreachable!("wildcard and binding handled above"),
                            }
                            self.emit(Bytecode::Eq);
                            let jump_no_match = self.emit(Bytecode::JumpIfNot(0));
                            // Pop subject from stack
                            self.emit(Bytecode::Pop);
                            // Generate arm body
                            self.gen_expr(&arm.body)?;
                            end_jumps.push(self.emit(Bytecode::Jump(0)));
                            let next_arm = self.current_idx();
                            self.patch_jump(jump_no_match, next_arm);
                        }
                    }
                }
                // Default fallthrough: pop subject, push nil
                self.emit(Bytecode::Pop);
                self.emit(Bytecode::Push(Value::Nil));
                let end = self.current_idx();
                for j in end_jumps {
                    self.patch_jump(j, end);
                }
            }
            Expr::Activate { type_name, query } => {
                self.emit(Bytecode::Activate {
                    type_name: type_name.clone(),
                    query: query.clone(),
                });
            }
            Expr::Sealed(stmts) => {
                self.emit(Bytecode::SealedBegin);
                for s in stmts {
                    self.gen_stmt(s)?;
                }
                self.emit(Bytecode::SealedEnd);
                self.emit(Bytecode::Push(Value::Nil));
            }
            Expr::Field { object, field } => {
                self.gen_expr(object)?;
                self.emit(Bytecode::GetField(field.clone()));
            }
            Expr::Array(elems) => {
                // Build a list by pushing all elements then collecting
                // In this simple VM we push a List value directly
                // For a stack-based VM we'd emit individual pushes + a BuildList instr;
                // here we inline the value since it's all literals at codegen time
                for e in elems {
                    self.gen_expr(e)?;
                }
                // Emit a "build list of N" — we use a Call to a builtin
                self.emit(Bytecode::Call {
                    name: "__build_list__".to_string(),
                    arity: elems.len() as u32,
                });
            }
            Expr::Path { segments } => {
                // Emit the last segment as a string value (enum variant reference)
                let variant = segments.last().cloned().unwrap_or_default();
                self.emit(Bytecode::Push(Value::Str(variant)));
            }
            Expr::Index { object, index } => {
                self.gen_expr(object)?;
                self.gen_expr(index)?;
                self.emit(Bytecode::GetIndex);
            }
            Expr::StructLit { type_name, fields, .. } => {
                // Push each field value in declaration order
                for (_, field_expr) in fields {
                    self.gen_expr(field_expr)?;
                }
                let field_names: Vec<String> = fields.iter().map(|(n, _)| n.clone()).collect();
                self.emit(Bytecode::BuildStruct {
                    type_name: type_name.clone(),
                    fields: field_names,
                });
            }
            Expr::With { base, updates } => {
                // Generate base struct clone then apply updates
                self.gen_expr(base)?;
                for (field, val_expr) in updates {
                    self.gen_expr(val_expr)?;
                    self.emit(Bytecode::SetField(field.clone()));
                }
            }
            Expr::Reason { query } => {
                self.emit(Bytecode::Reason { query: query.clone() });
            }
            Expr::Parallel { entries } => {
                // For parallel, emit each expression sequentially and collect into a Map
                // A full implementation would use threads; here we collect results into a Map
                let n = entries.len() as u32;
                for (name, expr) in entries {
                    self.emit(Bytecode::Push(Value::Str(name.clone())));
                    self.gen_expr(expr)?;
                }
                self.emit(Bytecode::BuildMap(n));
            }
            Expr::Trace { label, body } => {
                self.emit(Bytecode::TraceBegin { label: label.clone() });
                for s in body {
                    self.gen_stmt(s)?;
                }
                self.emit(Bytecode::TraceEnd { label: label.clone() });
                self.emit(Bytecode::Push(Value::Nil));
            }
            // New expression kinds — push Nil as placeholder
            _ => {
                self.emit(Bytecode::Push(Value::Nil));
            }
        }
        Ok(())
    }
}

// ── Helper: extract a representative span from a statement ────────────────────

fn stmt_span(stmt: &Stmt) -> el_lexer::Span {
    match stmt {
        Stmt::Let { span, .. }
        | Stmt::Return(_, span)
        | Stmt::Expr(_, span)
        | Stmt::FnDef { span, .. }
        | Stmt::TypeDef { span, .. }
        | Stmt::EnumDef { span, .. }
        | Stmt::TestDef { span, .. }
        | Stmt::Seed(_, span)
        | Stmt::Assert(_, span) => *span,
        _ => el_lexer::Span::new(0, 0, 0, 0),
    }
}

#[cfg(test)]
mod tests {
    use el_lexer::tokenize;
    use el_parser::parse;
    use super::*;

    fn gen(src: &str) -> Vec<Bytecode> {
        let tokens = tokenize(src).unwrap();
        let prog = parse(tokens, src.to_string()).unwrap();
        let cg = Codegen::new(false);
        let (bc, _) = cg.generate(&prog).unwrap();
        bc
    }

    #[test]
    fn test_push_int() {
        let bc = gen("42");
        assert!(matches!(&bc[0], Bytecode::Push(Value::Int(42))));
    }

    #[test]
    fn test_let_store() {
        let bc = gen("let x = 1");
        assert!(matches!(&bc[1], Bytecode::StoreLocal(n) if n == "x"));
    }

    #[test]
    fn test_add() {
        let bc = gen("1 + 2");
        assert!(bc.iter().any(|b| matches!(b, Bytecode::Add)));
    }

    #[test]
    fn test_halt_at_end() {
        let bc = gen("42");
        assert!(matches!(bc.last(), Some(Bytecode::Halt)));
    }

    #[test]
    fn test_activate_emitted() {
        let bc = gen(r#"activate User where "query""#);
        assert!(bc.iter().any(|b| matches!(b, Bytecode::Activate { .. })));
    }

    #[test]
    fn test_sealed_markers() {
        let bc = gen("sealed { let x = 1 }");
        assert!(bc.iter().any(|b| matches!(b, Bytecode::SealedBegin)));
        assert!(bc.iter().any(|b| matches!(b, Bytecode::SealedEnd)));
    }
}