use full_moon::ast::{Ast, Stmt, LastStmt, Expression};
use full_moon::node::Node;
use crate::common::diagnostics::{DiagnosticBundle, Diagnostic, DiagnosticLevel};
use crate::common::spans::Span;

pub fn validate_ast(ast: &Ast, file_id: usize) -> Result<(), DiagnosticBundle> {
    let mut bundle = DiagnosticBundle::new();
    
    let block = ast.nodes();
    for stmt in block.stmts() {
        validate_stmt(stmt, &mut bundle, file_id);
    }
    
    if let Some(last_stmt) = block.last_stmt() {
        validate_last_stmt(last_stmt, &mut bundle, file_id);
    }

    if bundle.diagnostics.is_empty() {
        Ok(())
    } else {
        Err(bundle)
    }
}

fn get_span(node: impl Node, file_id: usize) -> Option<Span> {
    node.tokens().next().and_then(|t| {
        let start = t.token().start_position();
        node.tokens().last().map(|last_t| {
            let end = last_t.token().end_position();
            Span::new(file_id, start.bytes() as u32, end.bytes() as u32)
        })
    })
}

fn validate_stmt(stmt: &Stmt, bundle: &mut DiagnosticBundle, file_id: usize) {
    match stmt {
        Stmt::LocalAssignment(assignment) => {
            if assignment.expressions().len() == 0 {
                bundle.push(Diagnostic {
                    message: "Local declaration must have an initializer in this subset.".into(),
                    level: DiagnosticLevel::Error,
                    span: get_span(assignment, file_id),
                });
            }
            // Check for multiple assignments (not supported in subset)
            if assignment.names().len() > 1 || assignment.expressions().len() > 1 {
                 bundle.push(Diagnostic {
                    message: "Multiple assignments are not supported in this subset.".into(),
                    level: DiagnosticLevel::Error,
                    span: get_span(assignment, file_id),
                });
            }
            for expr in assignment.expressions() {
                validate_expr(expr, bundle, file_id);
            }
        }
        Stmt::Assignment(assignment) => {
            if assignment.variables().len() > 1 || assignment.expressions().len() > 1 {
                bundle.push(Diagnostic {
                    message: "Multiple assignments are not supported in this subset.".into(),
                    level: DiagnosticLevel::Error,
                    span: get_span(assignment, file_id),
                });
            }
            for expr in assignment.expressions() {
                validate_expr(expr, bundle, file_id);
            }
        }
        Stmt::FunctionCall(call) => {
             validate_expr(&Expression::FunctionCall(call.clone()), bundle, file_id);
        }
        Stmt::If(if_stmt) => {
            validate_expr(if_stmt.condition(), bundle, file_id);
            for s in if_stmt.block().stmts() {
                validate_stmt(s, bundle, file_id);
            }
            if let Some(else_ifs) = if_stmt.else_if() {
                for else_if in else_ifs {
                    validate_expr(else_if.condition(), bundle, file_id);
                    for s in else_if.block().stmts() {
                        validate_stmt(s, bundle, file_id);
                    }
                }
            }
            if let Some(else_block) = if_stmt.else_block() {
                 for s in else_block.stmts() {
                    validate_stmt(s, bundle, file_id);
                }
            }
        }
        Stmt::While(while_stmt) => {
            validate_expr(while_stmt.condition(), bundle, file_id);
            for s in while_stmt.block().stmts() {
                validate_stmt(s, bundle, file_id);
            }
        }
        Stmt::Do(do_stmt) => {
            for s in do_stmt.block().stmts() {
                validate_stmt(s, bundle, file_id);
            }
            if let Some(last) = do_stmt.block().last_stmt() {
                validate_last_stmt(last, bundle, file_id);
            }
        }
        Stmt::FunctionDeclaration(f) => {
             // Check for varargs
             if f.body().parameters().iter().any(|p| matches!(p, full_moon::ast::Parameter::Ellipsis(_))) {
                  bundle.push(Diagnostic {
                    message: "Varargs (...) are not supported in this subset.".into(),
                    level: DiagnosticLevel::Error,
                    span: get_span(f, file_id),
                });
             }
             for s in f.body().block().stmts() {
                 validate_stmt(s, bundle, file_id);
             }
             if let Some(last) = f.body().block().last_stmt() {
                 validate_last_stmt(last, bundle, file_id);
             }
        }
        _ => {
            // For color-square-lua, we might need to support more statement types if used.
            // Let's see what's failing.
            bundle.push(Diagnostic {
                message: format!("Statement type not supported in this subset: {:?}", stmt),
                level: DiagnosticLevel::Error,
                span: get_span(stmt, file_id),
            });
        }
    }
}

fn validate_expr(expr: &Expression, bundle: &mut DiagnosticBundle, file_id: usize) {
    match expr {
        Expression::BinaryOperator { lhs, rhs, .. } => {
            validate_expr(lhs, bundle, file_id);
            validate_expr(rhs, bundle, file_id);
        }
        Expression::UnaryOperator { expression, .. } => {
            validate_expr(expression, bundle, file_id);
        }
        Expression::Number(_) | Expression::String(_) | Expression::Symbol(_) => {}
        Expression::Var(_) => {
        }
        Expression::FunctionCall(_) => {
        }
        Expression::TableConstructor(_) => {
        }
        Expression::Function(_) => {
             bundle.push(Diagnostic {
                message: "Anonymous functions/closures are not supported in this subset.".into(),
                level: DiagnosticLevel::Error,
                span: get_span(expr, file_id),
            });
        }
        _ => {
            bundle.push(Diagnostic {
                message: "Expression type not supported in this subset.".into(),
                level: DiagnosticLevel::Error,
                span: get_span(expr, file_id),
            });
        }
    }
}

fn validate_last_stmt(last_stmt: &LastStmt, bundle: &mut DiagnosticBundle, file_id: usize) {
    match last_stmt {
        LastStmt::Return(ret) => {
            if ret.returns().len() > 1 {
                 bundle.push(Diagnostic {
                    message: "Multiple return values are not supported in this subset.".into(),
                    level: DiagnosticLevel::Error,
                    span: get_span(ret, file_id),
                });
            }
            for expr in ret.returns() {
                validate_expr(expr, bundle, file_id);
            }
        }
        LastStmt::Break(_) => {
        }
        _ => {}
    }
}