glyph-runtime 0.0.1

//! AST to IR transformation
//!
//! This module converts the Glyph AST into the intermediate representation (IR).

use crate::ir::*;
use glyph_parser::ast::*;
use glyph_types::Value;
use std::collections::HashMap;

/// Loop information for break/continue
struct LoopInfo {
    break_label: String,
    continue_label: String,
}

pub struct AstToIr {
    builder: IRBuilder,
    /// Variable scope tracking
    scopes: Vec<HashMap<String, ()>>,
    /// Loop stack for break/continue
    loop_stack: Vec<LoopInfo>,
}

impl AstToIr {
    pub fn new(program: &ProgramDecorator) -> Self {
        let builder = IRBuilder::new(
            program.name.clone(),
            program.version.clone(),
            program.requires.clone(),
        );

        Self {
            builder,
            scopes: vec![HashMap::new()], // Global scope
            loop_stack: Vec::new(),
        }
    }

    pub fn transform(mut self, module: &GlyphModule) -> IRModule {
        // Process all functions
        for stmt in &module.statements {
            if let Statement::FunctionDef(func) = stmt {
                self.transform_function(func);
            }
        }

        // Set entry point if main function exists
        // Note: entry_point is handled by the builder's finish method

        self.builder.finish()
    }

    fn transform_function(&mut self, func: &Function) {
        // Start new function
        let params: Vec<String> = func.params.iter().map(|p| p.name.clone()).collect();
        self.builder
            .start_function(func.name.clone(), params.clone(), func.is_async);

        // New scope for function
        self.push_scope();

        // Add parameters to scope
        for param in &params {
            self.scopes.last_mut().unwrap().insert(param.clone(), ());
        }

        // Transform function body
        for stmt in &func.body {
            self.transform_statement(stmt);
        }

        // Ensure function ends with return
        if let Some(last_stmt) = func.body.last() {
            if !matches!(last_stmt, Statement::Return(_)) {
                self.builder.emit(IRInstruction::LoadConst(Value::None));
                self.builder.terminate(IRTerminator::Return);
            }
        } else {
            self.builder.emit(IRInstruction::LoadConst(Value::None));
            self.builder.terminate(IRTerminator::Return);
        }

        self.pop_scope();
    }

    fn transform_statement(&mut self, stmt: &Statement) {
        match stmt {
            Statement::Expression(expr) => {
                self.transform_expression(expr);
                self.builder.emit(IRInstruction::Pop);
            }
            Statement::Return(expr) => {
                if let Some(expr) = expr {
                    self.transform_expression(expr);
                } else {
                    self.builder.emit(IRInstruction::LoadConst(Value::None));
                }
                self.builder.terminate(IRTerminator::Return);
            }
            Statement::Let { name, value, .. } => {
                self.transform_expression(value);
                self.builder.emit(IRInstruction::StoreVar(name.clone()));
                self.add_to_scope(name.clone());
            }
            Statement::Assignment { target, value } => {
                self.transform_expression(value);
                self.builder.emit(IRInstruction::StoreVar(target.clone()));
            }
            Statement::If(if_stmt) => {
                self.transform_if_statement(if_stmt);
            }
            Statement::Match(match_stmt) => {
                self.transform_match_statement(match_stmt);
            }
            Statement::While { condition, body } => {
                self.transform_while_loop(condition, body);
            }
            Statement::For {
                variable,
                iterable,
                body,
            } => {
                self.transform_for_loop(variable, iterable, body);
            }
            Statement::Break => {
                if let Some(loop_info) = self.loop_stack.last() {
                    self.builder
                        .terminate(IRTerminator::Jump(loop_info.break_label.clone()));
                    // Create a new unreachable block for any dead code after break
                    let dead_label = self.builder.fresh_label();
                    self.builder.new_block(dead_label);
                } else {
                    panic!("Break statement outside of loop");
                }
            }
            Statement::Continue => {
                if let Some(loop_info) = self.loop_stack.last() {
                    self.builder
                        .terminate(IRTerminator::Jump(loop_info.continue_label.clone()));
                    // Create a new unreachable block for any dead code after continue
                    let dead_label = self.builder.fresh_label();
                    self.builder.new_block(dead_label);
                } else {
                    panic!("Continue statement outside of loop");
                }
            }
            Statement::FunctionDef(_) => {
                // Nested functions not supported yet
                panic!("Nested functions not supported");
            }
        }
    }

    fn transform_expression(&mut self, expr: &Expression) {
        match expr {
            Expression::Literal(lit) => {
                let value = match lit {
                    Literal::Int(n) => Value::Int(*n),
                    Literal::Float(f) => Value::Float(*f),
                    Literal::String(s) => Value::Str(s.clone()),
                    Literal::Bool(b) => Value::Bool(*b),
                    Literal::Unit => Value::None,
                };
                self.builder.emit(IRInstruction::LoadConst(value));
            }
            Expression::Identifier(name) => {
                self.builder.emit(IRInstruction::LoadVar(name.clone()));
            }
            Expression::BinaryOp { left, op, right } => {
                self.transform_expression(left);
                self.transform_expression(right);
                let ir_op = match op {
                    BinaryOperator::Add => BinaryOp::Add,
                    BinaryOperator::Subtract => BinaryOp::Sub,
                    BinaryOperator::Multiply => BinaryOp::Mul,
                    BinaryOperator::Divide => BinaryOp::Div,
                    BinaryOperator::Modulo => BinaryOp::Mod,
                    BinaryOperator::Power => BinaryOp::Pow,
                    BinaryOperator::Equal => BinaryOp::Eq,
                    BinaryOperator::NotEqual => BinaryOp::Ne,
                    BinaryOperator::Less => BinaryOp::Lt,
                    BinaryOperator::Greater => BinaryOp::Gt,
                    BinaryOperator::LessEqual => BinaryOp::Le,
                    BinaryOperator::GreaterEqual => BinaryOp::Ge,
                    BinaryOperator::And => BinaryOp::And,
                    BinaryOperator::Or => BinaryOp::Or,
                };
                self.builder.emit(IRInstruction::BinaryOp(ir_op));
            }
            Expression::UnaryOp { op, operand } => {
                self.transform_expression(operand);
                let ir_op = match op {
                    UnaryOperator::Negate => UnaryOp::Neg,
                    UnaryOperator::Not => UnaryOp::Not,
                };
                self.builder.emit(IRInstruction::UnaryOp(ir_op));
            }
            Expression::Call { func, args, kwargs } => {
                // Handle intrinsic calls
                if let Expression::Attribute { value, attr } = func.as_ref() {
                    if let Expression::Identifier(module) = value.as_ref() {
                        let intrinsic_name = format!("{module}.{attr}");

                        // Push positional args
                        for arg in args {
                            self.transform_expression(arg);
                        }

                        // Push keyword args as dict
                        if !kwargs.is_empty() {
                            for (key, value) in kwargs {
                                self.builder
                                    .emit(IRInstruction::LoadConst(Value::Str(key.clone())));
                                self.transform_expression(value);
                            }
                            self.builder.emit(IRInstruction::MakeDict(kwargs.len()));
                        }

                        let total_args = args.len() + if kwargs.is_empty() { 0 } else { 1 };
                        self.builder.emit(IRInstruction::CallIntrinsic {
                            name: intrinsic_name,
                            args_count: total_args,
                        });
                        return;
                    }
                }

                // Regular function call
                if let Expression::Identifier(name) = func.as_ref() {
                    // Push arguments
                    for arg in args {
                        self.transform_expression(arg);
                    }
                    if !kwargs.is_empty() {
                        panic!("Keyword arguments in regular function calls not yet supported");
                    }

                    self.builder.emit(IRInstruction::Call {
                        func: name.clone(),
                        args_count: args.len(),
                    });
                } else {
                    panic!("Complex function expressions not yet supported");
                }
            }
            Expression::List(elements) => {
                for elem in elements {
                    self.transform_expression(elem);
                }
                self.builder.emit(IRInstruction::MakeList(elements.len()));
            }
            Expression::Dict(pairs) => {
                for (key, value) in pairs {
                    self.transform_expression(key);
                    self.transform_expression(value);
                }
                self.builder.emit(IRInstruction::MakeDict(pairs.len()));
            }
            Expression::Attribute { value, attr } => {
                self.transform_expression(value);
                self.builder.emit(IRInstruction::GetAttr(attr.clone()));
            }
            Expression::Await(expr) => {
                self.transform_expression(expr);
                self.builder.emit(IRInstruction::Await);
            }
            Expression::FString(_parts) => {
                // F-string support - for now just emit empty string
                self.builder
                    .emit(IRInstruction::LoadConst(Value::Str("".to_string())));
            }
            Expression::IfExpr {
                test,
                if_true,
                if_false,
            } => {
                let then_label = self.builder.fresh_label();
                let else_label = self.builder.fresh_label();
                let end_label = self.builder.fresh_label();

                self.transform_expression(test);
                self.builder.terminate(IRTerminator::JumpIf {
                    then_block: then_label.clone(),
                    else_block: else_label.clone(),
                });

                self.builder.new_block(then_label);
                self.transform_expression(if_true);
                self.builder
                    .terminate(IRTerminator::Jump(end_label.clone()));

                self.builder.new_block(else_label);
                self.transform_expression(if_false);
                self.builder
                    .terminate(IRTerminator::Jump(end_label.clone()));

                self.builder.new_block(end_label);
            }
        }
    }

    fn transform_if_statement(&mut self, if_stmt: &IfStatement) {
        let then_label = self.builder.fresh_label();
        let else_label = self.builder.fresh_label();
        let end_label = self.builder.fresh_label();

        // Test condition
        self.transform_expression(&if_stmt.condition);
        self.builder.terminate(IRTerminator::JumpIf {
            then_block: then_label.clone(),
            else_block: else_label.clone(),
        });

        // Then block
        self.builder.new_block(then_label);
        for stmt in &if_stmt.then_body {
            self.transform_statement(stmt);
        }
        self.builder
            .terminate(IRTerminator::Jump(end_label.clone()));

        // Else block (including elif)
        self.builder.new_block(else_label);
        if !if_stmt.elif_clauses.is_empty() || if_stmt.else_body.is_some() {
            // Handle elif/else
            if let Some(else_body) = &if_stmt.else_body {
                for stmt in else_body {
                    self.transform_statement(stmt);
                }
            }
        }
        self.builder
            .terminate(IRTerminator::Jump(end_label.clone()));

        // Continue after if
        self.builder.new_block(end_label);
    }

    fn transform_match_statement(&mut self, match_stmt: &MatchStatement) {
        self.transform_expression(&match_stmt.subject);

        let mut cases = Vec::new();
        let mut case_labels = Vec::new();
        let end_label = self.builder.fresh_label();

        // First, collect all the patterns and labels
        for case in &match_stmt.cases {
            let case_label = self.builder.fresh_label();
            let pattern = self.transform_pattern(&case.pattern);
            cases.push((pattern, case_label.clone()));
            case_labels.push((case_label, case));
        }

        // Emit the match terminator which will do the pattern matching
        self.builder.terminate(IRTerminator::Match {
            cases,
            default: None, // No default for now - all patterns should be exhaustive
        });

        // Now generate code for each case
        for (case_label, case) in case_labels {
            self.builder.new_block(case_label);
            self.push_scope(); // New scope for pattern bindings

            // Handle variable bindings from pattern
            if let Pattern::Variable(name) = &case.pattern {
                // For variable patterns, we need to bind the matched value
                self.builder.emit(IRInstruction::StoreVar(name.clone()));
                self.add_to_scope(name.clone());
            }

            for stmt in &case.body {
                self.transform_statement(stmt);
            }

            // Only add jump if we haven't already terminated this block
            if !self.builder.is_terminated() {
                self.builder
                    .terminate(IRTerminator::Jump(end_label.clone()));
            }
            self.pop_scope();
        }

        self.builder.new_block(end_label);
    }

    fn transform_pattern(&mut self, pattern: &Pattern) -> IRPattern {
        match pattern {
            Pattern::Literal(lit) => {
                let value = match lit {
                    Literal::Int(n) => Value::Int(*n),
                    Literal::Float(f) => Value::Float(*f),
                    Literal::String(s) => Value::Str(s.clone()),
                    Literal::Bool(b) => Value::Bool(*b),
                    Literal::Unit => Value::None,
                };
                IRPattern::Literal(value)
            }
            Pattern::Variable(name) => {
                self.add_to_scope(name.clone());
                IRPattern::Variable(name.clone())
            }
            Pattern::Constructor { name, args } => {
                let ir_args = args.iter().map(|p| self.transform_pattern(p)).collect();
                IRPattern::Constructor {
                    name: name.clone(),
                    args: ir_args,
                }
            }
            Pattern::Wildcard => IRPattern::Wildcard,
        }
    }

    fn push_scope(&mut self) {
        self.scopes.push(HashMap::new());
    }

    fn pop_scope(&mut self) {
        self.scopes.pop();
    }

    fn add_to_scope(&mut self, name: String) {
        if let Some(scope) = self.scopes.last_mut() {
            scope.insert(name, ());
        }
    }

    fn transform_while_loop(&mut self, condition: &Expression, body: &[Statement]) {
        let loop_label = self.builder.fresh_label();
        let body_label = self.builder.fresh_label();
        let end_label = self.builder.fresh_label();

        // Push loop info for break/continue
        self.loop_stack.push(LoopInfo {
            break_label: end_label.clone(),
            continue_label: loop_label.clone(),
        });

        // Jump to loop start
        self.builder
            .terminate(IRTerminator::Jump(loop_label.clone()));

        // Loop start: evaluate condition
        self.builder.new_block(loop_label.clone());
        self.transform_expression(condition);
        self.builder.terminate(IRTerminator::JumpIf {
            then_block: body_label.clone(),
            else_block: end_label.clone(),
        });

        // Loop body
        self.builder.new_block(body_label);
        self.push_scope();
        for stmt in body {
            self.transform_statement(stmt);
        }
        self.pop_scope();

        // Jump back to loop start
        self.builder.terminate(IRTerminator::Jump(loop_label));

        // End of loop
        self.builder.new_block(end_label);
        self.loop_stack.pop();
    }

    fn transform_for_loop(&mut self, variable: &str, iterable: &Expression, body: &[Statement]) {
        // For loops are transformed into while loops with an iterator
        // For simplicity, we'll implement a basic version for lists

        let iter_var = format!("__iter_{variable}");
        let index_var = format!("__index_{variable}");
        let len_var = format!("__len_{variable}");

        // Store the iterable
        self.transform_expression(iterable);
        self.builder.emit(IRInstruction::StoreVar(iter_var.clone()));

        // Get length of iterable
        self.builder.emit(IRInstruction::LoadVar(iter_var.clone()));
        self.builder.emit(IRInstruction::CallMethod {
            name: "len".to_string(),
            argc: 0,
        });
        self.builder.emit(IRInstruction::StoreVar(len_var.clone()));

        // Initialize index to 0
        self.builder.emit(IRInstruction::LoadConst(Value::Int(0)));
        self.builder
            .emit(IRInstruction::StoreVar(index_var.clone()));

        // Loop labels
        let loop_label = self.builder.fresh_label();
        let body_label = self.builder.fresh_label();
        let end_label = self.builder.fresh_label();

        // Push loop info for break/continue
        self.loop_stack.push(LoopInfo {
            break_label: end_label.clone(),
            continue_label: loop_label.clone(),
        });

        // Jump to loop start
        self.builder
            .terminate(IRTerminator::Jump(loop_label.clone()));

        // Loop start: check if index < length
        self.builder.new_block(loop_label.clone());
        self.builder.emit(IRInstruction::LoadVar(index_var.clone()));
        self.builder.emit(IRInstruction::LoadVar(len_var.clone()));
        self.builder.emit(IRInstruction::BinaryOp(BinaryOp::Lt));
        self.builder.terminate(IRTerminator::JumpIf {
            then_block: body_label.clone(),
            else_block: end_label.clone(),
        });

        // Loop body
        self.builder.new_block(body_label.clone());
        self.push_scope();

        // Get current item and bind to loop variable
        self.builder.emit(IRInstruction::LoadVar(iter_var.clone()));
        self.builder.emit(IRInstruction::LoadVar(index_var.clone()));
        self.builder.emit(IRInstruction::GetItem);
        self.builder
            .emit(IRInstruction::StoreVar(variable.to_string()));
        self.add_to_scope(variable.to_string());

        // Execute loop body
        for stmt in body {
            self.transform_statement(stmt);
        }

        // Increment index
        self.builder.emit(IRInstruction::LoadVar(index_var.clone()));
        self.builder.emit(IRInstruction::LoadConst(Value::Int(1)));
        self.builder.emit(IRInstruction::BinaryOp(BinaryOp::Add));
        self.builder.emit(IRInstruction::StoreVar(index_var));

        self.pop_scope();

        // Jump back to loop start
        self.builder.terminate(IRTerminator::Jump(loop_label));

        // End of loop
        self.builder.new_block(end_label);
        self.loop_stack.pop();
    }
}

/// Transform AST to IR
pub fn ast_to_ir(module: &GlyphModule) -> IRModule {
    let transformer = AstToIr::new(&module.program);
    transformer.transform(module)
}