// 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 parser

import v2.ast

@[inline]
fn (mut p Parser) expect_type() ast.Expr {
    // return p.try_type() or {
    //     p.error(err.msg())
    // }
    typ := p.try_type()
    if typ is ast.EmptyExpr {
        p.error('expecting type, got `${p.tok}`')
    }
    return typ
}

// TODO: use result or stick with empty expr?
// pub fn (mut p Parser) try_type() !ast.Expr {
fn (mut p Parser) try_type() ast.Expr {
    match p.tok {
        // pointer: `&type`
        .amp {
            p.next()
            mut lifetime := ''
            if p.tok == .xor {
                p.next()
                lifetime = p.expect_name()
            }
            return ast.Type(ast.PointerType{
                base_type: p.expect_type()
                lifetime:  lifetime
            })
        }
        // lifetime: `^a`
        .xor {
            pos := p.pos
            p.next()
            return ast.LifetimeExpr{
                name: p.expect_name()
                pos:  pos
            }
        }
        // comptime type: `$enum` | `$struct` |  etc
        .dollar {
            p.next()
            mut comptime_inner := ast.Expr(ast.empty_expr)
            if p.tok == .name {
                comptime_inner = ast.Expr(p.ident())
            } else {
                // TODO: match only allowed tokens otherwise error
                comptime_inner = ast.Expr(ast.Ident{
                    pos:  p.pos
                    name: p.tok().str()
                })
            }
            return ast.ComptimeExpr{
                expr: comptime_inner
            }
        }
        // variadic: `...type`
        .ellipsis {
            p.next()
            // TODO: what will we use here?
            return ast.PrefixExpr{
                op:   .ellipsis
                expr: p.expect_type()
            }
        }
        // atomic | shared
        // eg. typespec in struct field with modifier. other cases handled in expr()
        .key_atomic, .key_shared {
            kind := p.tok
            pos := p.pos
            p.next()
            return modifier_expr_with_expr(kind, p.expect_type(), pos)
        }
        // function: `fn(type) type`
        .key_fn {
            p.next()
            return ast.Type(p.fn_type())
        }
        // nil
        .key_nil {
            p.next()
            return ast.Type(ast.NilType{})
        }
        // none
        .key_none {
            p.next()
            return ast.Type(ast.NoneType{})
        }
        // inline / anonymous struct
        .key_struct {
            p.next()
            generic_params := if p.tok == .lsbr { p.generic_list() } else { []ast.Expr{} }
            embedded, fields := p.struct_decl_fields(.v, false, false)
            // TODO: should we use this or just StructDecl
            // even though it technically is not one? hrmm
            return ast.Type(ast.AnonStructType{
                generic_params: generic_params
                embedded:       embedded
                fields:         fields
            })
        }
        // tuple (multi return): `(type, type)`
        .lpar {
            p.next()
            // expect at least two (so we otherwise error)
            mut types := [p.expect_type()]
            p.expect(.comma)
            types << p.expect_type()
            // more than two
            for p.tok == .comma {
                p.next()
                types << p.expect_type()
            }
            p.expect(.rpar)
            return ast.Type(ast.TupleType{
                types: types
            })
        }
        // array: `[]type` | `[len]type`
        .lsbr {
            p.next()
            // dynamic array
            if p.tok == .rsbr {
                p.next()
                return ast.Type(ast.ArrayType{
                    elem_type: p.expect_type()
                })
            }
            // fixed array
            len_expr := p.expr(.lowest)
            p.expect(.rsbr)
            return ast.Type(ast.ArrayFixedType{
                len:       len_expr
                elem_type: p.expect_type()
            })
        }
        // name | chan | map
        .name {
            pos := p.pos
            // TODO: cleam this up
            name := p.ident_or_named_type()
            if p.tok == .lsbr && name !is ast.Type && p.tok != .semicolon {
                // TODO: is there a better solution than this. maybe it should be the
                // concern of p.fn_parameters() & p.struct_decl() rather than this?
                // `fn(param_a []type)` | `struct { field_a []type }`
                if name is ast.Ident && name.name.len + pos.offset < p.pos.offset {
                    return name
                }
                // TODO: using ast.GenericArgs here may not be correct,
                // perhaps we should rename it to ast.GenericTypes
                // if name is ast.Ident && p.pos == pos + name.name.len { return name }
                return ast.Type(ast.GenericType{
                    name:   name
                    params: p.generic_list()
                })
            }
            return name
        }
        // result: `!` | `!type`
        .not {
            p.next()
            return ast.Type(ast.ResultType{
                base_type: if p.tok != .semicolon { p.try_type() } else { ast.empty_expr }
                // base_type: p.tok != .semicolon { p.try_type() or { ast.empty_expr } } else { ast.empty_expr }
            })
        }
        // option: `?` | `?type`
        .question {
            p.next()
            return ast.Type(ast.OptionType{
                base_type: if p.tok != .semicolon { p.try_type() } else { ast.empty_expr }
                // base_type: if p.tok != .semicolon { p.try_type() or { ast.empty_expr } } else { ast.empty_expr }
            })
        }
        else {
            // return error('expecting type, got `${p.tok}`')
            return ast.empty_expr
        }
    }
}

// function type / signature
fn (mut p Parser) fn_type() ast.FnType {
    generic_params := if p.tok == .lsbr { p.generic_list() } else { []ast.Expr{} }
    params := p.fn_parameters()
    return_type := if p.tok != .semicolon { p.try_type() } else { ast.empty_expr }
    return fn_type_with_return_type(generic_params, params, return_type)
}

fn fn_type_with_return_type(generic_params []ast.Expr, params []ast.Parameter, return_type ast.Expr) ast.FnType {
    mut fn_type := ast.FnType{
        generic_params: generic_params
        params:         params
        return_type:    ast.empty_expr
    }
    fn_type.return_type = return_type
    return fn_type
}

// `ident` | `map[type]type | `(`chan`|`chan type`) | (`thread`|`thread type`)
fn (mut p Parser) ident_or_named_type() ast.Expr {
    pos := p.pos
    lit := p.lit
    // `map[type]type`
    if lit == 'map' {
        p.next()
        if p.tok == .lsbr {
            p.next()
            key_type := p.expect_type()
            p.expect(.rsbr)
            return ast.Type(ast.MapType{
                key_type:   key_type
                value_type: p.expect_type()
            })
        }
        // struct called `map` in builtin
        return ast.Ident{
            name: 'map'
            pos:  pos
        }
    }
    // `chan` | `chan type`
    if lit == 'chan' {
        p.next()
        elem_type := if p.tok != .semicolon { p.try_type() } else { ast.empty_expr }
        if elem_type !is ast.EmptyExpr {
            return ast.Type(ast.ChannelType{
                elem_type: elem_type
            })
        }
        // struct called `chan` in builtin
        return ast.Ident{
            name: 'chan'
            pos:  pos
        }
    }
    // `thread` | `thread type`
    else if lit == 'thread' {
        p.next()
        return ast.Type(ast.ThreadType{
            elem_type: if p.tok != .semicolon { p.try_type() } else { ast.empty_expr }
        })
    }
    // `ident` | `selector` - ident or type name
    return p.ident_or_selector_expr()
}

// `ident` | (`Type`|`Type[T]`) | (`mod.Type`|`mod.Type[T]`)
fn (mut p Parser) ident_or_type() ast.Expr {
    ident_or_selector := p.ident_or_selector_expr()
    if p.tok == .lsbr {
        return ast.Type(ast.GenericType{
            name:   ident_or_selector
            params: p.generic_list()
        })
    }
    return ident_or_selector
}