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