module transformer

import os
import v2.ast
import v2.parser
import v2.pref as vpref
import v2.token
import v2.types

fn auto_str_parse_files_for_test(label string, source string) ([]ast.File, ast.FlatAst, &types.Environment, &vpref.Preferences) {
    tmp_file := os.join_path(os.vtmp_dir(), 'v2_auto_str_generation_${label}_${os.getpid()}.v')
    os.write_file(tmp_file, source) or { panic('failed to write temp file') }
    defer {
        os.rm(tmp_file) or {}
    }
    mut prefs := &vpref.Preferences{
        backend:     .x64
        no_parallel: true
    }
    mut file_set := token.FileSet.new()
    mut par := parser.Parser.new(prefs)
    files := par.parse_files([tmp_file], mut file_set)
    mut env := types.Environment.new()
    mut checker := types.Checker.new(prefs, file_set, env)
    checker.check_files(files)
    flat := ast.flatten_files(files)
    return files, flat, env, prefs
}

fn auto_str_legacy_generated_files(label string, source string) []ast.File {
    files, _, env, prefs := auto_str_parse_files_for_test(label, source)
    mut trans := Transformer.new_with_pref(env, prefs)
    return trans.transform_files(files)
}

fn auto_str_flat_generated_files(label string, source string) []ast.File {
    files, flat, env, prefs := auto_str_parse_files_for_test(label, source)
    mut trans := Transformer.new_with_pref(env, prefs)
    transformed_flat, _ := trans.transform_files_to_flat_via_driver(&flat, files)
    return transformed_flat.to_files()
}

fn auto_str_generated_fn_names(files []ast.File) []string {
    mut names := []string{}
    for file in files {
        for stmt in file.stmts {
            if stmt is ast.FnDecl && stmt.name.ends_with('__str') {
                names << stmt.name
            }
        }
    }
    names.sort()
    return names
}

fn assert_auto_str_names(files []ast.File, expected []string) {
    names := auto_str_generated_fn_names(files)
    for name in expected {
        assert name in names, '${name} missing from generated names ${names}'
    }
}

fn auto_str_generated_fn(files []ast.File, name string) ?ast.FnDecl {
    for file in files {
        for stmt in file.stmts {
            if stmt is ast.FnDecl && stmt.name == name {
                return stmt
            }
        }
    }
    return none
}

fn auto_str_collect_stmt_string_literals(stmt ast.Stmt, mut values []string) {
    match stmt {
        ast.AssignStmt {
            for expr in stmt.lhs {
                auto_str_collect_expr_string_literals(expr, mut values)
            }
            for expr in stmt.rhs {
                auto_str_collect_expr_string_literals(expr, mut values)
            }
        }
        ast.ExprStmt {
            auto_str_collect_expr_string_literals(stmt.expr, mut values)
        }
        ast.ReturnStmt {
            for expr in stmt.exprs {
                auto_str_collect_expr_string_literals(expr, mut values)
            }
        }
        else {}
    }
}

fn auto_str_collect_expr_string_literals(expr ast.Expr, mut values []string) {
    match expr {
        ast.AsCastExpr {
            auto_str_collect_expr_string_literals(expr.expr, mut values)
            auto_str_collect_expr_string_literals(expr.typ, mut values)
        }
        ast.CallExpr {
            for arg in expr.args {
                auto_str_collect_expr_string_literals(arg, mut values)
            }
        }
        ast.CastExpr {
            auto_str_collect_expr_string_literals(expr.expr, mut values)
            auto_str_collect_expr_string_literals(expr.typ, mut values)
        }
        ast.IfExpr {
            auto_str_collect_expr_string_literals(expr.cond, mut values)
            for nested in expr.stmts {
                auto_str_collect_stmt_string_literals(nested, mut values)
            }
            auto_str_collect_expr_string_literals(expr.else_expr, mut values)
        }
        ast.InfixExpr {
            auto_str_collect_expr_string_literals(expr.lhs, mut values)
            auto_str_collect_expr_string_literals(expr.rhs, mut values)
        }
        ast.ModifierExpr {
            auto_str_collect_expr_string_literals(expr.expr, mut values)
        }
        ast.ParenExpr {
            auto_str_collect_expr_string_literals(expr.expr, mut values)
        }
        ast.PrefixExpr {
            auto_str_collect_expr_string_literals(expr.expr, mut values)
        }
        ast.SelectorExpr {
            auto_str_collect_expr_string_literals(expr.lhs, mut values)
        }
        ast.StringLiteral {
            values << expr.value
        }
        else {}
    }
}

fn auto_str_fn_string_literals(files []ast.File, fn_name string) []string {
    fn_decl := auto_str_generated_fn(files, fn_name) or {
        assert false, 'missing generated function ${fn_name}'
        return []string{}
    }
    mut values := []string{}
    for stmt in fn_decl.stmts {
        auto_str_collect_stmt_string_literals(stmt, mut values)
    }
    return values
}

fn auto_str_has_literal_prefix(values []string, prefix string) bool {
    for value in values {
        if value.starts_with(prefix) {
            return true
        }
    }
    return false
}

fn auto_str_collect_stmt_call_names(stmt ast.Stmt, mut names []string) {
    match stmt {
        ast.AssignStmt {
            for expr in stmt.lhs {
                auto_str_collect_expr_call_names(expr, mut names)
            }
            for expr in stmt.rhs {
                auto_str_collect_expr_call_names(expr, mut names)
            }
        }
        ast.ExprStmt {
            auto_str_collect_expr_call_names(stmt.expr, mut names)
        }
        ast.ForStmt {
            auto_str_collect_stmt_call_names(stmt.init, mut names)
            auto_str_collect_expr_call_names(stmt.cond, mut names)
            auto_str_collect_stmt_call_names(stmt.post, mut names)
            for nested in stmt.stmts {
                auto_str_collect_stmt_call_names(nested, mut names)
            }
        }
        ast.ReturnStmt {
            for expr in stmt.exprs {
                auto_str_collect_expr_call_names(expr, mut names)
            }
        }
        else {}
    }
}

fn auto_str_collect_expr_call_names(expr ast.Expr, mut names []string) {
    match expr {
        ast.AsCastExpr {
            auto_str_collect_expr_call_names(expr.expr, mut names)
            auto_str_collect_expr_call_names(expr.typ, mut names)
        }
        ast.CallExpr {
            if expr.lhs is ast.Ident {
                lhs := expr.lhs as ast.Ident
                names << lhs.name
            }
            for arg in expr.args {
                auto_str_collect_expr_call_names(arg, mut names)
            }
        }
        ast.CastExpr {
            auto_str_collect_expr_call_names(expr.expr, mut names)
            auto_str_collect_expr_call_names(expr.typ, mut names)
        }
        ast.IfExpr {
            auto_str_collect_expr_call_names(expr.cond, mut names)
            for nested in expr.stmts {
                auto_str_collect_stmt_call_names(nested, mut names)
            }
            auto_str_collect_expr_call_names(expr.else_expr, mut names)
        }
        ast.InfixExpr {
            auto_str_collect_expr_call_names(expr.lhs, mut names)
            auto_str_collect_expr_call_names(expr.rhs, mut names)
        }
        ast.ModifierExpr {
            auto_str_collect_expr_call_names(expr.expr, mut names)
        }
        ast.ParenExpr {
            auto_str_collect_expr_call_names(expr.expr, mut names)
        }
        ast.PrefixExpr {
            auto_str_collect_expr_call_names(expr.expr, mut names)
        }
        ast.SelectorExpr {
            auto_str_collect_expr_call_names(expr.lhs, mut names)
        }
        ast.StringInterLiteral {
            for inter in expr.inters {
                auto_str_collect_expr_call_names(inter.expr, mut names)
                auto_str_collect_expr_call_names(inter.format_expr, mut names)
            }
        }
        else {}
    }
}

fn auto_str_fn_call_names(files []ast.File, fn_name string) []string {
    fn_decl := auto_str_generated_fn(files, fn_name) or {
        assert false, 'missing generated function ${fn_name}'
        return []string{}
    }
    mut names := []string{}
    for stmt in fn_decl.stmts {
        auto_str_collect_stmt_call_names(stmt, mut names)
    }
    return names
}

fn test_array_str_generation_uses_helper_name_for_elem_type_when_metadata_is_stale() {
    env := types.Environment.new()
    mut trans := Transformer.new_with_pref(env, &vpref.Preferences{})
    stmt := trans.generate_array_str_fn('Array_int_str', 'ast__KeywordOperator__str')
    if stmt is ast.FnDecl {
        mut calls := []string{}
        for body_stmt in stmt.stmts {
            auto_str_collect_stmt_call_names(body_stmt, mut calls)
        }
        assert 'int__str' in calls, calls.str()
        assert 'ast__KeywordOperator__str__str' !in calls, calls.str()
        return
    }
    assert false, 'expected generated Array_int_str function'
}

fn auto_str_collect_stmt_selector_names(stmt ast.Stmt, mut names []string) {
    match stmt {
        ast.AssignStmt {
            for expr in stmt.lhs {
                auto_str_collect_expr_selector_names(expr, mut names)
            }
            for expr in stmt.rhs {
                auto_str_collect_expr_selector_names(expr, mut names)
            }
        }
        ast.ExprStmt {
            auto_str_collect_expr_selector_names(stmt.expr, mut names)
        }
        ast.ReturnStmt {
            for expr in stmt.exprs {
                auto_str_collect_expr_selector_names(expr, mut names)
            }
        }
        else {}
    }
}

fn auto_str_collect_expr_selector_names(expr ast.Expr, mut names []string) {
    match expr {
        ast.AsCastExpr {
            auto_str_collect_expr_selector_names(expr.expr, mut names)
            auto_str_collect_expr_selector_names(expr.typ, mut names)
        }
        ast.CallExpr {
            for arg in expr.args {
                auto_str_collect_expr_selector_names(arg, mut names)
            }
        }
        ast.CastExpr {
            auto_str_collect_expr_selector_names(expr.expr, mut names)
            auto_str_collect_expr_selector_names(expr.typ, mut names)
        }
        ast.IfExpr {
            auto_str_collect_expr_selector_names(expr.cond, mut names)
            for nested in expr.stmts {
                auto_str_collect_stmt_selector_names(nested, mut names)
            }
            auto_str_collect_expr_selector_names(expr.else_expr, mut names)
        }
        ast.InfixExpr {
            auto_str_collect_expr_selector_names(expr.lhs, mut names)
            auto_str_collect_expr_selector_names(expr.rhs, mut names)
        }
        ast.ModifierExpr {
            auto_str_collect_expr_selector_names(expr.expr, mut names)
        }
        ast.ParenExpr {
            auto_str_collect_expr_selector_names(expr.expr, mut names)
        }
        ast.PrefixExpr {
            auto_str_collect_expr_selector_names(expr.expr, mut names)
        }
        ast.SelectorExpr {
            names << expr.rhs.name
            auto_str_collect_expr_selector_names(expr.lhs, mut names)
        }
        ast.StringInterLiteral {
            for inter in expr.inters {
                auto_str_collect_expr_selector_names(inter.expr, mut names)
                auto_str_collect_expr_selector_names(inter.format_expr, mut names)
            }
        }
        else {}
    }
}

fn auto_str_fn_selector_names(files []ast.File, fn_name string) []string {
    fn_decl := auto_str_generated_fn(files, fn_name) or {
        assert false, 'missing generated function ${fn_name}'
        return []string{}
    }
    mut names := []string{}
    for stmt in fn_decl.stmts {
        auto_str_collect_stmt_selector_names(stmt, mut names)
    }
    return names
}

fn assert_auto_str_simple_point_body(files []ast.File) {
    literals := auto_str_fn_string_literals(files, 'Point__str')
    calls := auto_str_fn_call_names(files, 'Point__str')
    assert 'Point{}' !in literals
    assert auto_str_has_literal_prefix(literals, 'Point{')
    assert '    x: ' in literals
    assert '    y: ' in literals
    assert 'strings__Builder__write_string' in calls
    assert 'strings__Builder__str' in calls
}

fn assert_auto_str_recursive_node_body(files []ast.File) {
    literals := auto_str_fn_string_literals(files, 'Node__str')
    calls := auto_str_fn_call_names(files, 'Node__str')
    assert 'Node{}' !in literals
    assert auto_str_has_literal_prefix(literals, 'Node{')
    assert '    value: ' in literals
    assert '    left: ' in literals
    assert '    right: ' in literals
    assert 'Tree__str' in calls
    assert 'string__replace' in calls
}

fn assert_auto_str_tree_sumtype_body(files []ast.File) {
    literals := auto_str_fn_string_literals(files, 'Tree__str')
    calls := auto_str_fn_call_names(files, 'Tree__str')
    selectors := auto_str_fn_selector_names(files, 'Tree__str')
    assert 'Tree(' in literals
    assert ')' in literals
    assert 'Tree{}' in literals
    assert 'Empty__str' in calls
    assert 'Node__str' in calls
    assert '_tag' in selectors
}

fn assert_auto_str_main_uses_node_str(files []ast.File) {
    calls := auto_str_fn_call_names(files, 'main')
    assert 'Node__str' in calls
}

fn test_auto_struct_str_generates_field_struct_helper_from_legacy_and_flat() {
    source := '
module main

struct Point {
    x int
    y int
}

fn show(point Point) string {
    return "\${point}"
}
'
    legacy := auto_str_legacy_generated_files('simple_struct_legacy', source)
    flat := auto_str_flat_generated_files('simple_struct_flat', source)
    assert_auto_str_names(legacy, ['Point__str'])
    assert_auto_str_names(flat, ['Point__str'])
    assert_auto_str_simple_point_body(legacy)
    assert_auto_str_simple_point_body(flat)
}

fn test_auto_sumtype_str_generates_wrapper_and_variant_helpers_from_legacy_and_flat() {
    source := '
module main

type Tree = Empty | Node

struct Empty {}

struct Node {
    value int
}

fn show(tree Tree) string {
    return "\${tree}"
}
'
    legacy := auto_str_legacy_generated_files('sumtype_legacy', source)
    flat := auto_str_flat_generated_files('sumtype_flat', source)
    assert_auto_str_names(legacy, ['Empty__str', 'Node__str', 'Tree__str'])
    assert_auto_str_names(flat, ['Empty__str', 'Node__str', 'Tree__str'])
    assert_auto_str_tree_sumtype_body(legacy)
    assert_auto_str_tree_sumtype_body(flat)
}

fn test_auto_recursive_struct_sumtype_str_generates_helpers_from_legacy_and_flat() {
    source := '
module main

type Tree = Empty | Node

struct Empty {}

struct Node {
    value int
    left  Tree
    right Tree
}

fn show(node Node) string {
    return "\${node}"
}
'
    legacy := auto_str_legacy_generated_files('recursive_legacy', source)
    flat := auto_str_flat_generated_files('recursive_flat', source)
    assert_auto_str_names(legacy, ['Empty__str', 'Node__str', 'Tree__str'])
    assert_auto_str_names(flat, ['Empty__str', 'Node__str', 'Tree__str'])
    assert_auto_str_recursive_node_body(legacy)
    assert_auto_str_recursive_node_body(flat)
    assert_auto_str_tree_sumtype_body(legacy)
    assert_auto_str_tree_sumtype_body(flat)
}

fn test_auto_recursive_struct_sumtype_str_from_call_arg_interpolation_legacy_and_flat() {
    source := '
module main

type Tree = Empty | Node

struct Empty {}

struct Node {
    value int
    left  Tree
    right Tree
}

fn sink(s string) {
    _ = s
}

fn main() {
    leaf := Node{7, Empty{}, Empty{}}
    root := Node{9, leaf, Empty{}}
    sink("\${root}")
}
'
    legacy := auto_str_legacy_generated_files('recursive_call_arg_legacy', source)
    flat := auto_str_flat_generated_files('recursive_call_arg_flat', source)
    assert_auto_str_names(legacy, ['Empty__str', 'Node__str', 'Tree__str'])
    assert_auto_str_names(flat, ['Empty__str', 'Node__str', 'Tree__str'])
    assert_auto_str_main_uses_node_str(legacy)
    assert_auto_str_main_uses_node_str(flat)
    assert_auto_str_recursive_node_body(legacy)
    assert_auto_str_recursive_node_body(flat)
    assert_auto_str_tree_sumtype_body(legacy)
    assert_auto_str_tree_sumtype_body(flat)
}

fn test_auto_recursive_struct_sumtype_str_from_implicit_main_call_arg_legacy_and_flat() {
    source := '
type Tree = Empty | Node

struct Empty {}

struct Node {
    value int
    left  Tree
    right Tree
}

fn sink(s string) {
    _ = s
}

fn main() {
    node1 := Node{30, Empty{}, Empty{}}
    node2 := Node{20, Empty{}, Empty{}}
    tree := Node{10, node1, node2}
    sink("\${tree}")
}
'
    legacy := auto_str_legacy_generated_files('recursive_implicit_main_call_arg_legacy', source)
    flat := auto_str_flat_generated_files('recursive_implicit_main_call_arg_flat', source)
    assert_auto_str_names(legacy, ['Empty__str', 'Node__str', 'Tree__str'])
    assert_auto_str_names(flat, ['Empty__str', 'Node__str', 'Tree__str'])
    assert_auto_str_main_uses_node_str(legacy)
    assert_auto_str_main_uses_node_str(flat)
    assert_auto_str_recursive_node_body(legacy)
    assert_auto_str_recursive_node_body(flat)
    assert_auto_str_tree_sumtype_body(legacy)
    assert_auto_str_tree_sumtype_body(flat)
}