// Copyright (c) 2020-2024 Joe Conigliaro. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.
module ast

import v2.token

pub const empty_expr = Expr(EmptyExpr(0))
pub const empty_stmt = Stmt(EmptyStmt(0))

type EmptyExpr = u8
type EmptyStmt = u8

pub type Expr = ArrayInitExpr
    | AsCastExpr
    | AssocExpr
    | BasicLiteral
    | CallExpr
    | CallOrCastExpr
    | CastExpr
    | ComptimeExpr
    | EmptyExpr
    | FieldInit
    | FnLiteral
    | GenericArgOrIndexExpr
    | GenericArgs
    | Ident
    | IfExpr
    | IfGuardExpr
    | IndexExpr
    | InfixExpr
    | InitExpr
    | Keyword
    | KeywordOperator
    | LambdaExpr
    | LifetimeExpr
    | LockExpr
    | MapInitExpr
    | MatchExpr
    | ModifierExpr
    | OrExpr
    | ParenExpr
    | PostfixExpr
    | PrefixExpr
    | RangeExpr
    | SelectExpr
    | SelectorExpr
    | SqlExpr
    | StringInterLiteral
    | StringLiteral
    | Tuple
    | Type
    | UnsafeExpr

// TODO: decide if this going to be done like this (FieldInit)

pub type Stmt = AsmStmt
    | AssertStmt
    | AssignStmt
    | BlockStmt
    | ComptimeStmt
    | ConstDecl
    | DeferStmt
    | Directive
    | EmptyStmt
    | EnumDecl
    | ExprStmt
    | FlowControlStmt
    | FnDecl
    | ForInStmt
    | ForStmt
    | GlobalDecl
    | ImportStmt
    | InterfaceDecl
    | LabelStmt
    | ModuleStmt
    | ReturnStmt
    | StructDecl
    | TypeDecl
    | []Attribute

// pub type Decl = ConstDecl | EnumDecl | FnDecl | GlobalDecl
//     | InterfaceDecl | StructDecl | TypeDecl

// TODO: (re)implement nested sumtype like TS (was removed from v)
// currently need to cast to type in parser.type. Should I leave like
// this or add these directly to Expr until nesting is implemented?
pub type Type = AnonStructType
    | ArrayFixedType
    | ArrayType
    | ChannelType
    | FnType
    | GenericType
    | MapType
    | NilType
    | NoneType
    | OptionType
    | PointerType
    | ResultType
    | ThreadType
    | TupleType

pub fn (exprs []Expr) name_list() string {
    mut out := ''
    for i := 0; i < exprs.len; i++ {
        out = out + exprs[i].name()
        if i < exprs.len - 1 {
            out = out + ','
        }
    }
    return out
}

pub fn (t Type) name() string {
    return match t {
        GenericType {
            '${t.name.name()}[${t.params.name_list()}]'
        }
        else {
            'Type'
        }
    }
}

// TODO: fix this, should be only what is needed
pub fn (expr Expr) name() string {
    return match expr {
        AsCastExpr {
            '${expr.expr.name()} as ${expr.typ.name()}'
        }
        BasicLiteral {
            expr.value
        }
        CallExpr {
            // TODO:
            '${expr.lhs.name()}()'
        }
        CallOrCastExpr {
            '${expr.lhs.name()}(${expr.expr.name()})'
        }
        EmptyExpr {
            // TODO:
            'EmptyExpr'
        }
        GenericArgs {
            '${expr.lhs.name()}[${expr.args.name_list()}]'
        }
        Ident {
            expr.name
        }
        IndexExpr {
            '${expr.lhs.name()}[${expr.expr.name()}]'
        }
        InfixExpr {
            '${expr.lhs.name()} ${expr.op} ${expr.rhs.name()}'
        }
        Keyword {
            expr.tok.str()
        }
        LifetimeExpr {
            '^' + expr.name
        }
        ModifierExpr {
            '${expr.kind} ${expr.expr.name()}'
        }
        ParenExpr {
            '(${expr.expr.name()})'
        }
        PrefixExpr {
            '${expr.op}${expr.expr.name()}'
        }
        SelectorExpr {
            expr.lhs.name() + '.' + expr.rhs.name
        }
        StringLiteral {
            "'${expr.value}'"
        }
        Type {
            'Type'
        }
        UnsafeExpr {
            // TODO:
            'UnsafeExpr'
        }
        else {
            'Expr'
        }
    }
}

pub fn (expr Expr) pos() token.Pos {
    return match expr {
        AsCastExpr {
            expr.pos
        }
        ArrayInitExpr {
            expr.pos
        }
        AssocExpr {
            expr.pos
        }
        BasicLiteral {
            expr.pos
        }
        CallExpr {
            expr.pos
        }
        CallOrCastExpr {
            expr.pos
        }
        CastExpr {
            expr.pos
        }
        ComptimeExpr {
            expr.pos
        }
        FnLiteral {
            expr.pos
        }
        GenericArgOrIndexExpr {
            expr.pos
        }
        GenericArgs {
            expr.pos
        }
        Ident {
            expr.pos
        }
        IfExpr {
            expr.pos
        }
        IfGuardExpr {
            expr.pos
        }
        IndexExpr {
            expr.pos
        }
        InfixExpr {
            expr.pos
        }
        InitExpr {
            expr.pos
        }
        KeywordOperator {
            expr.pos
        }
        LifetimeExpr {
            expr.pos
        }
        LambdaExpr {
            expr.pos
        }
        LockExpr {
            expr.pos
        }
        MapInitExpr {
            expr.pos
        }
        MatchExpr {
            expr.pos
        }
        ModifierExpr {
            expr.pos
        }
        OrExpr {
            expr.pos
        }
        ParenExpr {
            expr.pos
        }
        PostfixExpr {
            expr.pos
        }
        PrefixExpr {
            expr.pos
        }
        RangeExpr {
            expr.pos
        }
        SelectExpr {
            expr.pos
        }
        SelectorExpr {
            expr.pos
        }
        SqlExpr {
            expr.pos
        }
        StringInterLiteral {
            expr.pos
        }
        StringLiteral {
            expr.pos
        }
        Tuple {
            expr.pos
        }
        UnsafeExpr {
            expr.pos
        }
        else {
            // Default for expressions without pos field (EmptyExpr, FieldInit, Keyword, Type)
            token.Pos{}
        }
    }
}

// File (AST container)
pub struct File {
pub:
    attributes     []Attribute
    mod            string
    name           string
    stmts          []Stmt
    imports        []ImportStmt
    selector_names map[int]string
}

pub enum Language {
    v
    c
    js
}

pub fn (lang Language) str() string {
    return match lang {
        .v { 'V' }
        .c { 'C' }
        .js { 'JS' }
    }
}

pub enum DeferMode {
    scoped   // default: defer to end of current scope
    function // defer(fn): defer to end of function
}

// Expressions
pub struct ArrayInitExpr {
pub mut:
    typ         Expr = empty_expr
    exprs       []Expr
    init        Expr = empty_expr
    cap         Expr = empty_expr
    len         Expr = empty_expr
    update_expr Expr = empty_expr // `a` in `[...a, 3, 4]`
    pos         token.Pos
}

pub struct AsCastExpr {
pub:
    expr Expr
    typ  Expr
    pos  token.Pos
}

pub struct AssocExpr {
pub:
    typ    Expr
    expr   Expr
    fields []FieldInit
    pos    token.Pos
}

pub struct BasicLiteral {
pub:
    kind  token.Token
    value string
    pos   token.Pos
}

pub struct CallExpr {
pub mut:
    lhs Expr
pub:
    args []Expr
    pos  token.Pos
}

pub struct CallOrCastExpr {
pub mut:
    lhs  Expr
    expr Expr
pub:
    pos token.Pos
}

pub struct CastExpr {
pub mut:
    typ  Expr
    expr Expr
pub:
    pos token.Pos
}

pub struct ComptimeExpr {
pub:
    expr Expr
    pos  token.Pos
}

pub struct FieldDecl {
pub:
    name                string
    typ                 Expr = empty_expr // can be empty as used for const (unless we use something else)
    value               Expr = empty_expr
    attributes          []Attribute
    is_public           bool
    is_mut              bool
    is_module_mut       bool
    is_interface_method bool
}

pub struct FieldInit {
pub:
    name string
pub mut:
    value Expr
}

// anon fn / closure
pub struct FnLiteral {
pub:
    typ           FnType
    captured_vars []Expr
    stmts         []Stmt
    pos           token.Pos
}

pub struct GenericArgs {
pub:
    lhs  Expr
    args []Expr // concrete types and lifetimes
    pos  token.Pos
}

pub struct GenericArgOrIndexExpr {
pub:
    lhs  Expr
    expr Expr
    pos  token.Pos
}

pub struct Ident {
pub mut:
    pos  token.Pos
    name string
}

pub struct IfExpr {
pub mut:
    cond      Expr = empty_expr
    else_expr Expr = empty_expr
    stmts     []Stmt
    pos       token.Pos
}

pub struct IfGuardExpr {
pub:
    stmt AssignStmt
    pos  token.Pos
}

pub struct InfixExpr {
pub mut:
    op  token.Token
    lhs Expr
    rhs Expr
    pos token.Pos
}

pub struct IndexExpr {
pub mut:
    lhs      Expr
    expr     Expr
    is_gated bool
    pos      token.Pos
}

pub struct InitExpr {
pub mut:
    typ    Expr
    fields []FieldInit
    pos    token.Pos
}

pub struct Keyword {
pub:
    tok token.Token
}

pub struct KeywordOperator {
pub:
    op    token.Token
    exprs []Expr
    pos   token.Pos
}

pub struct Tuple {
pub:
    exprs []Expr
    pos   token.Pos
}

pub struct LambdaExpr {
pub:
    args []Ident
    expr Expr
    pos  token.Pos
}

pub struct LockExpr {
pub:
    lock_exprs  []Expr
    rlock_exprs []Expr
    stmts       []Stmt
    pos         token.Pos
}

pub struct MapInitExpr {
pub:
    typ  Expr = empty_expr
    keys []Expr
    vals []Expr
    pos  token.Pos
}

pub struct MatchBranch {
pub:
    cond  []Expr
    stmts []Stmt
    pos   token.Pos
}

pub struct MatchExpr {
pub:
    expr     Expr
    branches []MatchBranch
    pos      token.Pos
}

// TODO: possibly merge modifiers into a bitfield where
// multiple modifiers are used. this could also be used
// for struct decl `mut:`, `pub mut:` etc. consider this
// [flag]
// pub enum Modifier {
//     .atomic
//     .global
//     .mutable
//     .public
//     .shared
//     .static
//     .volatile
// }

pub struct ModifierExpr {
pub mut:
    kind token.Token
    expr Expr
    pos  token.Pos
}

// pub fn (expr ModifierExpr) unwrap() Expr {
//     return expr.expr
// }

pub struct OrExpr {
pub:
    expr  Expr
    stmts []Stmt
    pos   token.Pos
}

pub struct Parameter {
pub mut:
    name   string
    typ    Expr
    is_mut bool
    pos    token.Pos
}

pub struct LifetimeExpr {
pub:
    name string
    pos  token.Pos
}

// name_str returns the parameter name.
pub fn (p Parameter) name_str() string {
    return p.name
}

pub struct ParenExpr {
pub:
    expr Expr
    pos  token.Pos
}

pub struct PostfixExpr {
pub:
    op   token.Token
    expr Expr
    pos  token.Pos
}

pub struct PrefixExpr {
pub mut:
    op   token.Token
    expr Expr
    pos  token.Pos
}

pub struct RangeExpr {
pub:
    op    token.Token // `..` exclusive | `...` inclusive
    start Expr
    end   Expr
    pos   token.Pos
}

pub struct SelectExpr {
pub:
    pos   token.Pos
    stmt  Stmt
    stmts []Stmt
    next  Expr = empty_expr
}

pub struct SelectorExpr {
pub mut:
    lhs Expr
    rhs Ident
    pos token.Pos
}

pub fn (se SelectorExpr) leftmost() Expr {
    if se.lhs is SelectorExpr {
        return se.lhs.leftmost()
    }
    return se.lhs
}

// pub fn (expr Expr) str() string {
//     return 'Expr.str() - ${expr.type_name()}'
// }

pub fn (se SelectorExpr) name() string {
    return se.lhs.name() + '.' + se.rhs.name
}

pub enum StringLiteralKind {
    c
    js
    raw
    v
}

pub fn (s StringLiteralKind) str() string {
    return match s {
        .c { 'c' }
        .js { 'js' }
        .raw { 'r' }
        .v { 'v' }
    }
}

// TODO: allow overriding this method in main v compiler
// that is why this method was renamed from `from_string`
@[direct_array_access]
pub fn StringLiteralKind.from_string_tinyv(s string) StringLiteralKind {
    match s[0] {
        `c` {
            return .c
        }
        `j` {
            if s[1] == `s` {
                return .js
            }
        }
        `r` {
            return .raw
        }
        else {}
    }

    return .v
}

// NOTE: I'm using two nodes StringLiteral & StringInterLiteral
// to avoid the extra array allocations when not needed.
pub struct StringLiteral {
pub:
    kind  StringLiteralKind
    value string
    pos   token.Pos
}

pub struct StringInterLiteral {
pub:
    kind   StringLiteralKind
    values []string
    inters []StringInter
    pos    token.Pos
}

pub struct StringInter {
pub:
    format    StringInterFormat
    width     int
    precision int
    expr      Expr
    // TEMP: prob removed once individual
    // fields are set, precision etc
    format_expr  Expr = empty_expr
    resolved_fmt string // resolved sprintf format specifier (e.g. '%d', '%s', '%lld'), set by transformer
}

pub enum StringInterFormat {
    unformatted
    binary
    character
    decimal
    exponent
    exponent_short
    float
    hex
    octal
    pointer_address
    string
}

pub fn StringInterFormat.from_u8(c u8) StringInterFormat {
    return match c {
        `b` { .binary }
        `c` { .character }
        `d` { .decimal }
        `e`, `E` { .exponent }
        `g`, `G` { .exponent_short }
        `f`, `F` { .float }
        `x`, `X` { .hex }
        `o` { .octal }
        `p` { .pointer_address }
        `s` { .string }
        else { .unformatted }
    }
}

pub fn (sif StringInterFormat) str() string {
    return match sif {
        .unformatted { '' }
        .binary { 'b' }
        .character { 'c' }
        .decimal { 'd' }
        .exponent { 'e' }
        .exponent_short { 'g' }
        .float { 'f' }
        .hex { 'x' }
        .octal { 'o' }
        .pointer_address { 'p' }
        .string { 's' }
    }
}

pub struct SqlExpr {
pub:
    expr       Expr
    table_name string
    is_count   bool
    is_create  bool
    pos        token.Pos
}

pub struct UnsafeExpr {
pub:
    stmts []Stmt
    pos   token.Pos
}

// Statements
pub struct AsmStmt {
pub:
    arch string
}

pub struct AssertStmt {
pub:
    expr  Expr
    extra Expr = empty_expr
}

pub struct AssignStmt {
pub:
    op  token.Token
    lhs []Expr
    rhs []Expr
    pos token.Pos
}

pub fn (attributes []Attribute) has(name string) bool {
    for attribute in attributes {
        if attribute.name == name {
            return true
        }
        // Also check value when it's a simple identifier (e.g., @[flag])
        if attribute.name == '' {
            if attribute.value is Ident && attribute.value.name == name {
                return true
            }
        }
    }
    return false
}

pub struct Attribute {
pub:
    name          string
    value         Expr
    comptime_cond Expr
    pos           token.Pos
}

pub struct BlockStmt {
pub:
    stmts []Stmt
}

pub struct ComptimeStmt {
pub:
    stmt Stmt
}

pub struct ConstDecl {
pub:
    is_public bool
    fields    []FieldInit
}

pub struct DeferStmt {
pub:
    mode  DeferMode
    stmts []Stmt
}

// #flag / #include
pub struct Directive {
pub:
    name    string
    value   string
    ct_cond string // optional comptime condition e.g. 'linux', 'darwin' for `#include linux <pty.h>`
}

pub struct EnumDecl {
pub:
    attributes []Attribute
    is_public  bool
    name       string
    as_type    Expr = empty_expr
    fields     []FieldDecl
}

pub struct ExprStmt {
pub mut:
    expr Expr
}

pub struct FlowControlStmt {
pub:
    op    token.Token
    label string
}

// enum FnArributes {
//     method
//     static
// }

pub struct FnDecl {
pub:
    attributes []Attribute
    is_public  bool
    is_method  bool
    is_static  bool
    receiver   Parameter
    language   Language = .v
    name       string
    typ        FnType
    stmts      []Stmt
    pos        token.Pos
}

pub struct ForStmt {
pub mut:
    init  Stmt = empty_stmt // initialization
    cond  Expr = empty_expr // condition
    post  Stmt = empty_stmt // post iteration (afterthought)
    stmts []Stmt
}

// NOTE: used as the initializer for ForStmt
pub struct ForInStmt {
pub mut:
    // key            string
    // value          string
    // value_is_mut bool
    // expr         Expr
    // TODO:
    key   Expr = empty_expr
    value Expr
    expr  Expr
}

pub struct GlobalDecl {
pub:
    attributes []Attribute
    fields     []FieldDecl
    is_public  bool
}

pub struct ImportStmt {
pub:
    name       string
    alias      string
    is_aliased bool
    symbols    []Expr
}

pub struct InterfaceDecl {
pub:
    is_public      bool
    attributes     []Attribute
    name           string
    generic_params []Expr
    embedded       []Expr
    fields         []FieldDecl
}

pub struct LabelStmt {
pub:
    name string
    stmt Stmt = empty_stmt
}

pub struct ModuleStmt {
pub:
    name string
}

pub struct ReturnStmt {
pub:
    exprs []Expr
}

pub struct StructDecl {
pub:
    attributes     []Attribute
    is_public      bool
    is_union       bool
    implements     []Expr
    embedded       []Expr
    language       Language = .v
    name           string
    generic_params []Expr
    fields         []FieldDecl
    pos            token.Pos
}

pub struct TypeDecl {
pub:
    is_public      bool
    language       Language
    name           string
    generic_params []Expr
    base_type      Expr = empty_expr
    variants       []Expr
}

// Type Nodes
pub struct ArrayType {
pub:
    elem_type Expr = empty_expr
}

pub struct ArrayFixedType {
pub:
    len       Expr = empty_expr
    elem_type Expr = empty_expr
}

pub struct ChannelType {
pub:
    // s255: cap defaults to empty_expr (like ThreadType/OptionType/etc.). The
    // parser omits cap for a bare `chan T` (type.v: `ChannelType{elem_type: ...}`),
    // which previously left it a zero-valued Expr (invalid sum-type tag, null
    // payload). Encoding that via FlatBuilder.add_expr crashes the arm64 self-host
    // — an exhaustive match on the unmatched tag falls into the first arm
    // (ArrayInitExpr) and derefs the null payload (same class as s251).
    cap       Expr = empty_expr
    elem_type Expr = empty_expr
}

pub struct ThreadType {
pub:
    elem_type Expr = empty_expr
}

pub struct FnType {
pub:
    generic_params []Expr
    params         []Parameter
pub mut:
    return_type Expr = empty_expr
}

pub fn (ft &FnType) str() string {
    mut s := 'fn('
    for i := 0; i < ft.params.len; i++ {
        param := ft.params[i]
        s = s + param.name + param.typ.name()
        if i < ft.params.len - 1 {
            s = s + ', '
        }
    }
    if ft.return_type is EmptyExpr {
        s = s + ')'
    } else {
        s = s + ') ' + ft.return_type.name()
    }
    return s
}

pub fn (ident &Ident) str() string {
    return ident.name.clone()
}

// pub fn (expr Expr) str() string {
//     return match expr {
//         Ident {
//             expr.name
//         }
//         SelectorExpr {
//             expr.lhs.str() + '.' + expr.rhs.name
//         }
//         else {
//             'missing Expr.str() for ${expr.type_name()}'
//             // expr.str()
//         }
//     }
// }

pub struct AnonStructType {
pub:
    generic_params []Expr
    embedded       []Expr
    fields         []FieldDecl
}

pub struct GenericType {
pub:
    name   Expr = empty_expr
    params []Expr
}

pub struct MapType {
pub:
    key_type   Expr = empty_expr
    value_type Expr = empty_expr
}

pub struct NilType {}

pub struct NoneType {}

pub struct OptionType {
pub:
    base_type Expr = empty_expr
}

pub struct PointerType {
pub:
    base_type Expr = empty_expr
    lifetime  string
}

pub struct ResultType {
pub:
    base_type Expr = empty_expr
}

pub struct TupleType {
pub:
    types []Expr
}