module transform

import v3.flat
import v3.types

// ArrayIndexInfo stores array index info metadata used by transform.
struct ArrayIndexInfo {
    base_id    flat.NodeId
    index_id   flat.NodeId
    base_type  string
    value_type string
}

// EnumFromStringInfo stores enum from string info metadata used by transform.
struct EnumFromStringInfo {
    enum_type string
    fields    []string
    arg_id    flat.NodeId
}

// optional_base_type supports optional base type handling for Transformer.
fn (t &Transformer) optional_base_type(typ string) string {
    if typ.len > 1 && (typ[0] == `?` || typ[0] == `!`) {
        return typ[1..]
    }
    return typ
}

// array_index_info supports array index info handling for Transformer.
fn (mut t Transformer) array_index_info(index_id flat.NodeId) ?ArrayIndexInfo {
    if int(index_id) < 0 {
        return none
    }
    expr := t.a.nodes[int(index_id)]
    if expr.kind != .index || expr.children_count < 2 || expr.value == 'range' {
        return none
    }
    base_id := t.a.child(&expr, 0)
    index_expr_id := t.a.child(&expr, 1)
    base_type := t.resolve_expr_type(base_id)
    if !base_type.starts_with('[]') {
        return none
    }
    value_type := t.normalize_type_alias(base_type[2..])
    if value_type.len == 0 {
        return none
    }
    return ArrayIndexInfo{
        base_id:    base_id
        index_id:   index_expr_id
        base_type:  base_type
        value_type: value_type
    }
}

// is_array_index_or_expr reports whether is array index or expr applies in transform.
fn (mut t Transformer) is_array_index_or_expr(node flat.Node) bool {
    if node.kind != .or_expr || node.children_count < 2 {
        return false
    }
    _ := t.array_index_info(t.a.child(&node, 0)) or { return false }
    return true
}

// transform_array_index_or_expr transforms transform array index or expr data for transform.
fn (mut t Transformer) transform_array_index_or_expr(id flat.NodeId, node flat.Node) flat.NodeId {
    if node.children_count < 2 {
        return id
    }
    expr_id := t.a.child(&node, 0)
    body_id := t.a.child(&node, 1)
    info := t.array_index_info(expr_id) or { return id }
    base_expr := t.stable_expr_for_reuse(info.base_id)
    index_name := t.new_temp('arr_idx')
    val_name := t.new_temp('arr_val')
    outer_pending := t.pending_stmts.clone()
    t.pending_stmts.clear()
    index_expr := t.transform_expr(info.index_id)
    mut prelude := []flat.NodeId{}
    t.drain_pending(mut prelude)
    prelude << t.make_decl_assign_typed(index_name, index_expr, 'int')
    prelude << t.make_decl_assign_typed(val_name, t.zero_value_for_type(info.value_type),
        info.value_type)

    idx_ident := t.make_ident(index_name)
    lower_ok := t.make_infix(.ge, idx_ident, t.make_int_literal(0))
    upper_ok := t.make_infix(.lt, t.make_ident(index_name),
        t.make_selector(base_expr, 'len', 'int'))
    found_cond := t.make_infix(.logical_and, lower_ok, upper_ok)
    index_value := t.make_index(base_expr, t.make_ident(index_name), info.value_type)
    then_block := t.make_block(arr1(t.make_assign(t.make_ident(val_name), index_value)))
    else_block := t.make_block(t.lower_or_body_to_stmts(body_id, val_name, info.value_type,
        node.value, ''))
    t.pending_stmts = outer_pending
    for stmt in prelude {
        t.pending_stmts << stmt
    }
    t.pending_stmts << t.make_if(found_cond, then_block, else_block)
    return t.make_ident(val_name)
}

// or_body_is_nil supports or body is nil handling for Transformer.
fn (t &Transformer) or_body_is_nil(body_id flat.NodeId) bool {
    if int(body_id) < 0 {
        return false
    }
    body := t.a.nodes[int(body_id)]
    if body.kind == .nil_literal {
        return true
    }
    if body.kind != .block || body.children_count != 1 {
        return false
    }
    stmt_id := t.a.child(&body, 0)
    stmt := t.a.nodes[int(stmt_id)]
    if stmt.kind == .expr_stmt && stmt.children_count == 1 {
        return t.a.nodes[int(t.a.child(&stmt, 0))].kind == .nil_literal
    }
    return stmt.kind == .nil_literal
}

// enum_from_string_info converts enum from string info data for transform.
fn (t &Transformer) enum_from_string_info(expr_id flat.NodeId) ?EnumFromStringInfo {
    if int(expr_id) < 0 {
        return none
    }
    expr := t.a.nodes[int(expr_id)]
    if expr.kind != .call || expr.children_count < 2 {
        return none
    }
    fn_id := t.a.child(&expr, 0)
    fn_node := t.a.nodes[int(fn_id)]
    if fn_node.kind != .selector || fn_node.value != 'from_string' || fn_node.children_count == 0 {
        return none
    }
    enum_id := t.a.child(&fn_node, 0)
    enum_type := t.enum_type_from_node(enum_id) or { return none }
    fields := t.enum_types[enum_type] or { return none }
    if fields.len == 0 {
        return none
    }
    return EnumFromStringInfo{
        enum_type: enum_type
        fields:    fields.clone()
        arg_id:    t.a.child(&expr, 1)
    }
}

// enum_type_from_node converts enum type from node data for transform.
fn (t &Transformer) enum_type_from_node(id flat.NodeId) ?string {
    if int(id) < 0 {
        return none
    }
    node := t.a.nodes[int(id)]
    if node.kind == .ident {
        if node.value in t.enum_types {
            return node.value
        }
        if t.cur_module.len > 0 && t.cur_module != 'main' && t.cur_module != 'builtin' {
            qname := '${t.cur_module}.${node.value}'
            if qname in t.enum_types {
                return qname
            }
        }
        for key, _ in t.enum_types {
            if key.all_after_last('.') == node.value {
                return key
            }
        }
    }
    if node.kind == .selector && node.children_count > 0 {
        base := t.a.child_node(&node, 0)
        if base.kind == .ident {
            qname := '${base.value}.${node.value}'
            if qname in t.enum_types {
                return qname
            }
        }
    }
    return none
}

// is_enum_from_string_or_expr reports whether is enum from string or expr applies in transform.
fn (mut t Transformer) is_enum_from_string_or_expr(node flat.Node) bool {
    if node.kind != .or_expr || node.children_count < 2 {
        return false
    }
    _ := t.enum_from_string_info(t.a.child(&node, 0)) or { return false }
    return true
}

// transform_enum_from_string_or_expr supports transform_enum_from_string_or_expr handling.
fn (mut t Transformer) transform_enum_from_string_or_expr(id flat.NodeId, node flat.Node) flat.NodeId {
    if node.children_count < 2 {
        return id
    }
    expr_id := t.a.child(&node, 0)
    body_id := t.a.child(&node, 1)
    info := t.enum_from_string_info(expr_id) or { return id }
    str_name := t.new_temp('enum_str')
    val_name := t.new_temp('enum_val')
    outer_pending := t.pending_stmts.clone()
    t.pending_stmts.clear()
    arg_expr := t.transform_expr(info.arg_id)
    mut prelude := []flat.NodeId{}
    t.drain_pending(mut prelude)
    prelude << t.make_decl_assign_typed(str_name, arg_expr, 'string')
    prelude << t.make_decl_assign_typed(val_name, t.zero_value_for_type(info.enum_type),
        info.enum_type)

    mut else_block := t.make_block(t.lower_or_body_to_stmts(body_id, val_name, info.enum_type,
        node.value, ''))
    for i := info.fields.len - 1; i >= 0; i-- {
        cond := t.make_call_typed('string__eq', arr2(t.make_ident(str_name),
            t.make_string_literal(info.fields[i])), 'bool')
        assign := t.make_assign(t.make_ident(val_name), t.make_int_literal(i))
        then_block := t.make_block(arr1(assign))
        else_block = t.make_if(cond, then_block, else_block)
    }
    t.pending_stmts = outer_pending
    for stmt in prelude {
        t.pending_stmts << stmt
    }
    t.pending_stmts << else_block
    return t.make_ident(val_name)
}

// or_expr_types supports or expr types handling for Transformer.
fn (t &Transformer) or_expr_types(expr_id flat.NodeId, fallback_type string) (string, string) {
    if !isnil(t.tc) {
        if typ := t.tc.expr_type(expr_id) {
            if typ is types.OptionType {
                base_name := t.value_type_name(typ.base_type)
                source_name := if base_name == 'int' && typ.base_type is types.Void {
                    '?void'
                } else {
                    '?${base_name}'
                }
                return source_name, base_name
            }
            if typ is types.ResultType {
                base_name := t.value_type_name(typ.base_type)
                source_name := if base_name == 'int' && typ.base_type is types.Void {
                    '!void'
                } else {
                    '!${base_name}'
                }
                return source_name, base_name
            }
        }
    }
    mut expr_type := t.node_type(expr_id)
    if expr_type.len == 0 || expr_type == 'Optional' {
        if fallback_type.len > 0 && !t.is_optional_type_name(fallback_type) {
            expr_type = '?${fallback_type}'
        } else {
            expr_type = fallback_type
        }
    }
    base_type := t.optional_base_type(expr_type)
    mut value_type := if base_type.len > 0 { base_type } else { fallback_type }
    if value_type.len == 0 || value_type == 'void' || value_type == '!' || value_type == '?' {
        value_type = 'int'
    }
    return expr_type, value_type
}

// value_type_name returns value type name data for Transformer.
fn (t &Transformer) value_type_name(typ types.Type) string {
    if typ is types.Void {
        return 'void'
    }
    return typ.name()
}

// is_optional_type_name reports whether is optional type name applies in transform.
fn (t &Transformer) is_optional_type_name(typ string) bool {
    return typ.len > 0 && (typ[0] == `?` || typ[0] == `!`)
}

// qualify_optional_type supports qualify optional type handling for Transformer.
fn (t &Transformer) qualify_optional_type(typ string) string {
    if typ.len < 2 || (typ[0] != `?` && typ[0] != `!`) {
        return typ
    }
    base := typ[1..]
    if base.contains('.') || base.len == 0 {
        return typ
    }
    qualified := t.qualify_type(base)
    if qualified != base {
        return typ[..1] + qualified
    }
    return typ
}

// qualify_type supports qualify type handling for Transformer.
fn (t &Transformer) qualify_type(name string) string {
    if name.contains('.') || name.len == 0 {
        return name
    }
    if t.cur_module.len > 0 && t.cur_module != 'main' && t.cur_module != 'builtin' {
        qname := '${t.cur_module}.${name}'
        if qname in t.sum_types || qname in t.structs {
            return qname
        }
    }
    if qualified := t.qualified_types[name] {
        return qualified
    }
    return name
}

// make_decl_assign_typed builds make decl assign typed data for transform.
fn (mut t Transformer) make_decl_assign_typed(name string, rhs flat.NodeId, typ string) flat.NodeId {
    decl := t.make_decl_assign(name, rhs)
    if typ.len > 0 {
        t.a.nodes[int(decl)].typ = typ
        t.set_var_type(name, typ)
    }
    return decl
}

// zero_value_for_type supports zero value for type handling for Transformer.
fn (mut t Transformer) zero_value_for_type(typ string) flat.NodeId {
    mut clean := typ
    if clean.starts_with('&') {
        return t.a.add(.nil_literal)
    }
    if clean.len > 1 && (clean[0] == `?` || clean[0] == `!`) {
        return t.make_optional_none(clean)
    }
    clean = t.normalize_type_alias(clean)
    if clean.starts_with('fn(') || clean.starts_with('fn (') || clean.starts_with('fn_ptr:') {
        return t.make_cast(clean, t.make_int_literal(0), clean)
    }
    if clean.starts_with('[]') {
        return t.make_array_init(clean[2..])
    }
    if clean.starts_with('map[') {
        return t.make_map_init(clean)
    }
    if clean == 'string' {
        return t.make_string_literal('')
    }
    if clean == 'bool' {
        return t.make_bool_literal(false)
    }
    if clean in ['f32', 'f64'] {
        return t.make_float_literal('0.0')
    }
    if clean in ['void', ''] {
        return t.make_int_literal(0)
    }
    if clean in ['int', 'i8', 'i16', 'i32', 'i64', 'isize', 'usize', 'u8', 'byte', 'u16', 'u32', 'u64', 'rune', 'char']
        || clean in t.enum_types {
        return t.make_int_literal(0)
    }
    return t.make_struct_init(clean)
}

// make_panic_stmt builds make panic stmt data for transform.
fn (mut t Transformer) make_panic_stmt(message string) flat.NodeId {
    call := t.make_call('panic', arr1(t.make_string_literal(message)))
    return t.make_expr_stmt(call)
}

// make_none_return_stmt builds make none return stmt data for transform.
fn (mut t Transformer) make_none_return_stmt() flat.NodeId {
    return t.make_return(t.a.add(.none_expr), t.cur_fn_ret_type)
}

// make_none_return_stmt_with_err builds make none return stmt with err data for transform.
fn (mut t Transformer) make_none_return_stmt_with_err(err_source string) flat.NodeId {
    if err_source.len == 0 {
        return t.make_none_return_stmt()
    }
    err_expr := t.make_selector(t.make_ident(err_source), 'err', 'IError')
    return t.make_none_return_stmt_with_err_expr(err_expr)
}

fn (mut t Transformer) make_none_return_stmt_with_err_expr(err_expr flat.NodeId) flat.NodeId {
    if int(err_expr) < 0 {
        return t.make_none_return_stmt()
    }
    return t.make_return(t.make_optional_none_with_err(t.cur_fn_ret_type, err_expr),
        t.cur_fn_ret_type)
}

// lower_or_expr_to_temp converts lower or expr to temp data for transform.
fn (mut t Transformer) lower_or_expr_to_temp(id flat.NodeId, node flat.Node) flat.NodeId {
    if node.children_count < 2 {
        return id
    }
    expr_id := t.a.child(&node, 0)
    body_id := t.a.child(&node, 1)
    expr_type, value_type := t.or_expr_types(expr_id, node.typ)
    if !t.is_optional_type_name(expr_type) {
        return t.transform_expr(expr_id)
    }
    is_void := value_type.len == 0 || value_type == 'void'

    opt_tmp := t.new_temp('or_opt')
    val_tmp := t.new_temp('or_val')

    outer_pending := t.pending_stmts.clone()
    t.pending_stmts.clear()
    new_expr := t.transform_expr(expr_id)
    mut prelude := []flat.NodeId{}
    t.drain_pending(mut prelude)

    if is_void {
        prelude << t.make_decl_assign_typed(opt_tmp, new_expr, expr_type)
        opt_ident := t.make_ident(opt_tmp)
        not_ok := t.make_prefix(.not, t.make_selector(opt_ident, 'ok', 'bool'))
        else_block := t.make_block(t.lower_or_body_to_stmts(body_id, '', '', node.value, opt_tmp))
        if_stmt := t.make_if(not_ok, else_block, t.make_empty())
        t.pending_stmts = outer_pending
        for stmt in prelude {
            t.pending_stmts << stmt
        }
        t.pending_stmts << if_stmt
        return t.make_int_literal(0)
    }

    prelude << t.make_decl_assign_typed(opt_tmp, new_expr, expr_type)
    prelude << t.make_decl_assign_typed(val_tmp, t.zero_value_for_type(value_type), value_type)

    opt_ident := t.make_ident(opt_tmp)
    ok_cond := t.make_selector(opt_ident, 'ok', 'bool')
    value_expr := t.make_selector(t.make_ident(opt_tmp), 'value', value_type)
    then_assign := t.make_assign(t.make_ident(val_tmp), value_expr)
    then_block := t.make_block(arr1(then_assign))
    else_block := t.make_block(t.lower_or_body_to_stmts(body_id, val_tmp, value_type, node.value,
        opt_tmp))
    if_stmt := t.make_if(ok_cond, then_block, else_block)
    t.pending_stmts = outer_pending
    for stmt in prelude {
        t.pending_stmts << stmt
    }
    t.pending_stmts << if_stmt
    return t.make_ident(val_tmp)
}

// lower_or_body_to_stmts converts lower or body to stmts data for transform.
fn (mut t Transformer) lower_or_body_to_stmts(body_id flat.NodeId, target_name string, target_type string, mode string, err_source string) []flat.NodeId {
    err_expr := if err_source.len > 0 {
        t.make_selector(t.make_ident(err_source), 'err', 'IError')
    } else {
        flat.empty_node
    }
    return t.lower_or_body_to_stmts_with_err_expr(body_id, target_name, target_type, mode, err_expr)
}

fn (mut t Transformer) lower_or_body_to_stmts_with_err_expr(body_id flat.NodeId, target_name string, target_type string, mode string, err_expr flat.NodeId) []flat.NodeId {
    if mode == '!' || mode == '?' {
        if t.is_optional_type_name(t.cur_fn_ret_type) {
            return arr1(t.make_none_return_stmt_with_err_expr(err_expr))
        }
        return arr1(t.make_panic_stmt('option/result propagation failed'))
    }
    if int(body_id) < 0 {
        return []flat.NodeId{}
    }
    body := t.a.nodes[int(body_id)]
    if body.kind != .block {
        if body.kind == .none_expr && t.is_optional_type_name(t.cur_fn_ret_type) {
            return arr1(t.make_none_return_stmt())
        }
        if target_name.len == 0 {
            value := t.transform_expr(body_id)
            mut result := []flat.NodeId{}
            t.drain_pending(mut result)
            if t.node_type(body_id) != 'void' {
                result << t.make_expr_stmt(value)
            }
            return result
        }
        return arr1(t.make_assign(t.make_ident(target_name), t.transform_expr(body_id)))
    }
    mut result := []flat.NodeId{}
    if body.children_count == 0 {
        return result
    }
    // The or-body is its own scope: `err` (and any locals it declares) are bound only
    // for its lowering and restored afterwards, so the previous binding (e.g. an outer
    // `err := 1`) survives and a subsequent `${err}` is not mis-typed as `IError`.
    // Mirrors transform_if_guard_else_block.
    saved_var_types := t.var_types.clone()
    t.set_var_type('err', 'IError')
    err_value := if int(err_expr) >= 0 {
        err_expr
    } else {
        t.make_struct_init('IError')
    }
    result << t.make_decl_assign_typed('err', err_value, 'IError')
    for i in 0 .. body.children_count {
        child_id := t.a.child(&body, i)
        child := t.a.nodes[int(child_id)]
        is_last := i == body.children_count - 1
        if is_last && child.kind == .expr_stmt && child.children_count > 0 {
            inner_id := t.a.child(&child, 0)
            inner := t.a.nodes[int(inner_id)]
            if inner.kind == .call && t.is_noreturn_call(inner) {
                expanded := t.transform_stmt(child_id)
                t.drain_pending(mut result)
                for eid in expanded {
                    result << eid
                }
            } else if inner.kind == .call && t.is_error_call(inner)
                && t.is_optional_type_name(t.cur_fn_ret_type) {
                result << t.make_return(inner_id, t.cur_fn_ret_type)
            } else if inner.kind == .none_expr && t.is_optional_type_name(t.cur_fn_ret_type) {
                result << t.make_none_return_stmt()
            } else if t.node_type(inner_id) == 'void' {
                expanded := t.transform_stmt(child_id)
                t.drain_pending(mut result)
                for eid in expanded {
                    result << eid
                }
            } else if target_name.len == 0 {
                expanded := t.transform_stmt(child_id)
                t.drain_pending(mut result)
                for eid in expanded {
                    result << eid
                }
            } else {
                value := if target_type in t.sum_types
                    || t.resolve_sum_name(target_type) in t.sum_types {
                    t.wrap_sum_value(inner_id, target_type)
                } else {
                    t.transform_expr_for_type(inner_id, target_type)
                }
                t.drain_pending(mut result)
                result << t.make_assign(t.make_ident(target_name), value)
            }
        } else {
            expanded := t.transform_stmt(child_id)
            t.drain_pending(mut result)
            for eid in expanded {
                result << eid
            }
        }
    }
    _ = target_type
    t.var_types = saved_var_types
    return result
}

fn (t &Transformer) is_error_call(node flat.Node) bool {
    if node.kind != .call || node.children_count == 0 {
        return false
    }
    fn_node := t.a.child_node(&node, 0)
    return fn_node.value == 'error' || fn_node.value == 'error_with_code'
}

// is_noreturn_call reports whether is noreturn call applies in transform.
fn (t &Transformer) is_noreturn_call(node flat.Node) bool {
    if node.kind != .call || node.children_count == 0 {
        return false
    }
    fn_node := t.a.child_node(&node, 0)
    return fn_node.value in ['panic', 'exit']
}