// Copyright (c) 2026 Alexander Medvednikov. 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

// Cursor is a lightweight, by-value handle pointing at one FlatNode inside a
// FlatAst. Consumers walk the AST via `match c.kind() {}` and descend via
// `c.edge(i)` / `c.list_at(i)` without rehydrating to ast.Stmt / ast.Expr.
//
// A Cursor is 16 bytes: an &FlatAst pointer plus a FlatNodeId. It is safe to
// copy, store in arrays, and pass by value. An "invalid" Cursor (id < 0) is
// used as a sentinel where the legacy AST would use empty_stmt / empty_expr.
pub struct Cursor {
pub:
	flat &FlatAst = unsafe { nil }
	id   FlatNodeId
}

@[inline]
pub fn (c Cursor) is_valid() bool {
	return c.flat != unsafe { nil } && c.id >= 0 && c.id < c.flat.nodes.len
}

@[inline]
pub fn (c Cursor) kind() FlatNodeKind {
	return c.flat.nodes[c.id].kind
}

@[inline]
pub fn (c Cursor) pos() token.Pos {
	return c.flat.nodes[c.id].pos
}

@[inline]
pub fn (c Cursor) flags() u8 {
	return c.flat.nodes[c.id].flags
}

@[inline]
pub fn (c Cursor) flag(bit u8) bool {
	return (c.flat.nodes[c.id].flags & bit) != 0
}

@[inline]
pub fn (c Cursor) aux() u16 {
	return c.flat.nodes[c.id].aux
}

@[inline]
pub fn (c Cursor) extra_int() int {
	return c.flat.nodes[c.id].extra
}

// extra_str interprets the node's `extra` field as an interned string id.
// Use only on kinds where the schema documents `extra` as a string slot
// (e.g. stmt_directive's value, stmt_import's alias).
@[inline]
pub fn (c Cursor) extra_str() string {
	return c.flat.string_at(c.flat.nodes[c.id].extra)
}

// name returns the interned primary string (FlatNode.name_id) for this node.
// Most kinds use this for identifiers, type names, fn names, etc.
@[inline]
pub fn (c Cursor) name() string {
	return c.flat.string_at(c.flat.nodes[c.id].name_id)
}

@[inline]
pub fn (c Cursor) name_id() int {
	return c.flat.nodes[c.id].name_id
}

// ident reads an expr_ident cursor directly into an Ident.
pub fn (c Cursor) ident() Ident {
	if !c.is_valid() || c.kind() != .expr_ident {
		return Ident{}
	}
	return Ident{
		pos:  c.pos()
		name: c.name()
	}
}

// import_stmt reads a stmt_import cursor directly into an ImportStmt.
pub fn (c Cursor) import_stmt() ImportStmt {
	if !c.is_valid() || c.kind() != .stmt_import {
		return ImportStmt{}
	}
	mut symbols := []Expr{cap: c.edge_count()}
	for i in 0 .. c.edge_count() {
		sym := c.edge(i)
		if sym.kind() == .expr_ident {
			symbols << Expr(sym.ident())
		}
	}
	return ImportStmt{
		name:       c.name()
		alias:      c.extra_str()
		is_aliased: c.flag(flag_is_aliased)
		symbols:    symbols
	}
}

// fn_decl_signature reads a stmt_fn_decl cursor into a body-less FnDecl. This
// mirrors FlatAst.decode_fn_decl_signature without going through FlatReader.
pub fn (c Cursor) fn_decl_signature() FnDecl {
	if !c.is_valid() || c.kind() != .stmt_fn_decl {
		return FnDecl{}
	}
	return FnDecl{
		attributes: attrs_from_cursor(c.list_at(2))
		is_public:  c.flag(flag_is_public)
		is_method:  c.flag(flag_is_method)
		is_static:  c.flag(flag_is_static)
		receiver:   parameter_from_cursor(c.edge(0))
		language:   unsafe { Language(int(c.aux())) }
		name:       c.name()
		typ:        fn_type_from_cursor(c.edge(1))
		pos:        c.pos()
	}
}

// fn_decl reads a stmt_fn_decl cursor into a legacy FnDecl. Unlike
// fn_decl_signature, this materializes the body statements for consumers that
// still use legacy statement walkers internally.
pub fn (c Cursor) fn_decl() FnDecl {
	if !c.is_valid() || c.kind() != .stmt_fn_decl {
		return FnDecl{}
	}
	signature := c.fn_decl_signature()
	return FnDecl{
		attributes: signature.attributes
		is_public:  signature.is_public
		is_method:  signature.is_method
		is_static:  signature.is_static
		receiver:   signature.receiver
		language:   signature.language
		name:       signature.name
		typ:        signature.typ
		stmts:      c.list_at(3).stmts()
		pos:        signature.pos
	}
}

// stmt reads a cursor through FlatAst.decode_stmt. This is the escape hatch
// for legacy statement walkers; prefer cursor-specific readers when possible.
pub fn (c Cursor) stmt() Stmt {
	if !c.is_valid() {
		return empty_stmt
	}
	return c.flat.decode_stmt(c.id)
}

// expr reads a cursor through FlatAst.decode_expr. This is the escape hatch
// for legacy expression walkers; prefer cursor-specific readers when possible.
pub fn (c Cursor) expr() Expr {
	if !c.is_valid() {
		return empty_expr
	}
	return c.flat.decode_expr(c.id)
}

// type_expr reads a type-expression cursor into the legacy Expr shape used by
// signature consumers. It is intentionally narrower than FlatReader.read_expr:
// non-type payloads such as field defaults or statement bodies stay omitted.
pub fn (c Cursor) type_expr() Expr {
	if !c.is_valid() {
		return empty_expr
	}
	match c.kind() {
		.expr_empty {
			return empty_expr
		}
		.expr_basic_literal {
			return Expr(BasicLiteral{
				kind:  unsafe { token.Token(int(c.aux())) }
				value: c.name()
				pos:   c.pos()
			})
		}
		.expr_ident {
			return Expr(c.ident())
		}
		.expr_lifetime {
			return Expr(LifetimeExpr{
				name: c.name()
				pos:  c.pos()
			})
		}
		.expr_modifier {
			return Expr(ModifierExpr{
				kind: unsafe { token.Token(int(c.aux())) }
				expr: c.edge(0).type_expr()
				pos:  c.pos()
			})
		}
		.expr_prefix {
			return Expr(PrefixExpr{
				op:   unsafe { token.Token(int(c.aux())) }
				expr: c.edge(0).type_expr()
				pos:  c.pos()
			})
		}
		.expr_selector {
			rhs := c.edge(1)
			return Expr(SelectorExpr{
				lhs: c.edge(0).type_expr()
				rhs: Ident{
					name: rhs.name()
					pos:  rhs.pos()
				}
				pos: c.pos()
			})
		}
		.expr_generic_args {
			return Expr(GenericArgs{
				lhs:  c.edge(0).type_expr()
				args: type_exprs_from_edges(c, 1)
				pos:  c.pos()
			})
		}
		.expr_generic_arg_or_index {
			return Expr(GenericArgOrIndexExpr{
				lhs:  c.edge(0).type_expr()
				expr: c.edge(1).type_expr()
				pos:  c.pos()
			})
		}
		.typ_anon_struct {
			return Expr(Type(AnonStructType{
				generic_params: c.list_at(0).type_exprs()
				embedded:       c.list_at(1).type_exprs()
				fields:         field_decl_type_list(c.list_at(2))
			}))
		}
		.typ_array_fixed {
			return Expr(Type(ArrayFixedType{
				len:       c.edge(0).type_expr()
				elem_type: c.edge(1).type_expr()
			}))
		}
		.typ_array {
			return Expr(Type(ArrayType{
				elem_type: c.edge(0).type_expr()
			}))
		}
		.typ_channel {
			return Expr(Type(ChannelType{
				cap:       c.edge(0).type_expr()
				elem_type: c.edge(1).type_expr()
			}))
		}
		.typ_fn {
			return Expr(Type(FnType{
				generic_params: c.list_at(0).type_exprs()
				params:         parameter_list_from_cursor(c.list_at(1))
				return_type:    c.edge(2).type_expr()
			}))
		}
		.typ_generic {
			return Expr(Type(GenericType{
				name:   c.edge(0).type_expr()
				params: type_exprs_from_edges(c, 1)
			}))
		}
		.typ_map {
			return Expr(Type(MapType{
				key_type:   c.edge(0).type_expr()
				value_type: c.edge(1).type_expr()
			}))
		}
		.typ_nil {
			return Expr(Type(NilType{}))
		}
		.typ_none {
			return Expr(Type(NoneType{}))
		}
		.typ_option {
			return Expr(Type(OptionType{
				base_type: c.edge(0).type_expr()
			}))
		}
		.typ_pointer {
			return Expr(Type(PointerType{
				base_type: c.edge(0).type_expr()
				lifetime:  c.name()
			}))
		}
		.typ_result {
			return Expr(Type(ResultType{
				base_type: c.edge(0).type_expr()
			}))
		}
		.typ_thread {
			return Expr(Type(ThreadType{
				elem_type: c.edge(0).type_expr()
			}))
		}
		.typ_tuple {
			return Expr(Type(TupleType{
				types: type_exprs_from_edges(c, 0)
			}))
		}
		else {
			return empty_expr
		}
	}
}

@[inline]
pub fn (c Cursor) edge_count() int {
	return c.flat.nodes[c.id].edge_count
}

// edge returns a Cursor over the i-th direct child edge of this node. Out-of-
// range indices return an invalid Cursor (id == invalid_flat_node_id), which
// callers can detect via `c.is_valid()`. This matches FlatReader's behaviour
// where `r.edge(n, i)` may return invalid_flat_node_id for missing slots.
@[inline]
pub fn (c Cursor) edge(i int) Cursor {
	return Cursor{
		flat: c.flat
		id:   c.flat.child_at(c.id, i)
	}
}

// list_at treats edge `i` as a reference to an `aux_list` node and returns a
// CursorList over its children. The flat schema represents list-typed fields
// (e.g. StructDecl.fields, FnDecl.stmts, EnumDecl.attributes) as a single
// edge to an aux_list whose children are the actual list items. Use `edge()`
// for fields stored as direct child edges of the parent (e.g. AssignStmt's
// LHS/RHS, AssertStmt's expr/extra).
@[inline]
pub fn (c Cursor) list_at(edge_i int) CursorList {
	return CursorList{
		flat:      c.flat
		parent_id: c.flat.child_at(c.id, edge_i)
	}
}

// for_body_list views a stmt_for node's body as a CursorList. The stmt_for flat
// layout is edges [init, cond, post, body_stmt_0, body_stmt_1, ...] — the body
// is trailing edges of the for-node itself, NOT a separate aux_list — so the
// body list is this node's children from edge index 3 onward.
@[inline]
pub fn (c Cursor) for_body_list() CursorList {
	return CursorList{
		flat:      c.flat
		parent_id: c.id
		offset:    3
	}
}

// CursorList is a view over the children of an `aux_list` node. It is the
// flat equivalent of `[]ast.Stmt` / `[]ast.Expr` / `[]FieldDecl` etc., except
// no slice is materialised — `at(i)` decodes one child at a time.
pub struct CursorList {
pub:
	flat      &FlatAst = unsafe { nil }
	parent_id FlatNodeId
	// offset skips the first `offset` child edges of `parent_id`. Default 0 (the
	// whole child list). Lets a trailing edge range be viewed as a list when the
	// items are direct child edges of a node rather than an aux_list — e.g. a
	// ForStmt body lives in edges [3..] of the stmt_for node itself (see
	// `Cursor.for_body_list`).
	offset int
}

@[inline]
pub fn (l CursorList) len() int {
	if l.flat == unsafe { nil } || l.parent_id < 0 || l.parent_id >= l.flat.nodes.len {
		return 0
	}
	return l.flat.nodes[l.parent_id].edge_count - l.offset
}

@[inline]
pub fn (l CursorList) at(i int) Cursor {
	return Cursor{
		flat: l.flat
		id:   l.flat.child_at(l.parent_id, l.offset + i)
	}
}

// type_exprs reads every item in a cursor list through Cursor.type_expr.
pub fn (l CursorList) type_exprs() []Expr {
	mut out := []Expr{cap: l.len()}
	for i in 0 .. l.len() {
		out << l.at(i).type_expr()
	}
	return out
}

// stmts reads every item in a cursor list through FlatAst.decode_stmt. This is
// the escape hatch for legacy statement walkers; prefer cursor-specific
// readers when the caller only needs a declaration signature or metadata.
pub fn (l CursorList) stmts() []Stmt {
	mut out := []Stmt{cap: l.len()}
	for i in 0 .. l.len() {
		c := l.at(i)
		if !c.is_valid() {
			continue
		}
		out << c.stmt()
	}
	return out
}

// attribute_expr reads the small expression subset used inside attributes.
// Attribute payloads are normally identifiers or strings; type_expr handles
// the type-like fallback cases without opening the full FlatReader.
pub fn (c Cursor) attribute_expr() Expr {
	if !c.is_valid() {
		return empty_expr
	}
	match c.kind() {
		.expr_empty {
			return empty_expr
		}
		.expr_basic_literal {
			return Expr(BasicLiteral{
				kind:  unsafe { token.Token(int(c.aux())) }
				value: c.name()
				pos:   c.pos()
			})
		}
		.expr_ident {
			return Expr(c.ident())
		}
		.expr_string {
			return Expr(StringLiteral{
				kind:  unsafe { StringLiteralKind(int(c.aux())) }
				value: c.name()
				pos:   c.pos()
			})
		}
		.expr_selector {
			rhs := c.edge(1)
			return Expr(SelectorExpr{
				lhs: c.edge(0).attribute_expr()
				rhs: Ident{
					name: rhs.name()
					pos:  rhs.pos()
				}
				pos: c.pos()
			})
		}
		else {
			return c.type_expr()
		}
	}
}

// attribute reads an aux_attribute cursor into the legacy Attribute shape.
pub fn (c Cursor) attribute() Attribute {
	if !c.is_valid() || c.kind() != .aux_attribute {
		return Attribute{}
	}
	return Attribute{
		name:  c.name()
		value: c.edge(0).attribute_expr()
		// comptime_cond (`@[if cond ?]`) can be an arbitrary expression, not the
		// ident/string subset attribute_expr handles — decode it fully (mirrors
		// FlatReader.read_attribute), else complex conditions silently become
		// empty_expr and `@[if ...]` functions are never elided.
		comptime_cond: c.edge(1).expr()
		pos:           c.pos()
	}
}

// attributes reads every aux_attribute in a cursor list.
pub fn (l CursorList) attributes() []Attribute {
	mut out := []Attribute{cap: l.len()}
	for i in 0 .. l.len() {
		attr := l.at(i)
		if !attr.is_valid() {
			continue
		}
		out << attr.attribute()
	}
	return out
}

// import_stmts reads every stmt_import in a cursor list.
pub fn (l CursorList) import_stmts() []ImportStmt {
	mut out := []ImportStmt{cap: l.len()}
	for i in 0 .. l.len() {
		imp := l.at(i)
		if !imp.is_valid() || imp.kind() != .stmt_import {
			continue
		}
		out << imp.import_stmt()
	}
	return out
}

// field_init reads an aux_field_init cursor. Field values are still legacy
// expression consumers today, so this materializes only the value expression,
// not the parent declaration.
pub fn (c Cursor) field_init() FieldInit {
	if !c.is_valid() || c.kind() != .aux_field_init {
		return FieldInit{}
	}
	value_c := c.edge(0)
	return FieldInit{
		name:  c.name()
		value: value_c.expr()
	}
}

// field_inits reads every aux_field_init in a cursor list.
pub fn (l CursorList) field_inits() []FieldInit {
	mut out := []FieldInit{cap: l.len()}
	for i in 0 .. l.len() {
		out << l.at(i).field_init()
	}
	return out
}

// field_decl reads an aux_field_decl cursor. Set decode_value to false when a
// caller only needs declaration metadata and type expressions.
pub fn (c Cursor) field_decl(decode_value bool) FieldDecl {
	if !c.is_valid() || c.kind() != .aux_field_decl {
		return FieldDecl{}
	}
	value_c := c.edge(1)
	return FieldDecl{
		name:                c.name()
		typ:                 c.edge(0).type_expr()
		value:               if decode_value { value_c.expr() } else { empty_expr }
		attributes:          c.list_at(2).attributes()
		is_public:           c.flag(flag_is_public)
		is_mut:              c.flag(flag_is_mut)
		is_module_mut:       c.flag(flag_field_is_module_mut)
		is_interface_method: c.flag(flag_field_is_interface_method)
	}
}

// field_decls reads every aux_field_decl in a cursor list.
pub fn (l CursorList) field_decls(decode_values bool) []FieldDecl {
	mut out := []FieldDecl{cap: l.len()}
	for i in 0 .. l.len() {
		out << l.at(i).field_decl(decode_values)
	}
	return out
}

// const_decl reads a stmt_const_decl cursor.
pub fn (c Cursor) const_decl() ConstDecl {
	if !c.is_valid() || c.kind() != .stmt_const_decl {
		return ConstDecl{}
	}
	return ConstDecl{
		is_public: c.flag(flag_is_public)
		fields:    c.list_at(0).field_inits()
	}
}

// enum_decl reads a stmt_enum_decl cursor. Set decode_values to false when
// only field names/attributes are needed.
pub fn (c Cursor) enum_decl(decode_values bool) EnumDecl {
	if !c.is_valid() || c.kind() != .stmt_enum_decl {
		return EnumDecl{}
	}
	return EnumDecl{
		attributes: c.list_at(1).attributes()
		is_public:  c.flag(flag_is_public)
		name:       c.name()
		as_type:    c.edge(0).type_expr()
		fields:     c.list_at(2).field_decls(decode_values)
	}
}

// global_decl reads a stmt_global_decl cursor. Set decode_values to false when
// only field metadata and declared types are needed.
pub fn (c Cursor) global_decl(decode_values bool) GlobalDecl {
	if !c.is_valid() || c.kind() != .stmt_global_decl {
		return GlobalDecl{}
	}
	return GlobalDecl{
		attributes: c.list_at(0).attributes()
		fields:     c.list_at(1).field_decls(decode_values)
		is_public:  c.flag(flag_is_public)
	}
}

// interface_decl reads a stmt_interface_decl cursor.
pub fn (c Cursor) interface_decl() InterfaceDecl {
	if !c.is_valid() || c.kind() != .stmt_interface_decl {
		return InterfaceDecl{}
	}
	return InterfaceDecl{
		is_public:      c.flag(flag_is_public)
		attributes:     c.list_at(0).attributes()
		name:           c.name()
		generic_params: c.list_at(1).type_exprs()
		embedded:       c.list_at(2).type_exprs()
		fields:         c.list_at(3).field_decls(true)
	}
}

// struct_decl reads a stmt_struct_decl cursor.
pub fn (c Cursor) struct_decl() StructDecl {
	if !c.is_valid() || c.kind() != .stmt_struct_decl {
		return StructDecl{}
	}
	return StructDecl{
		attributes:     c.list_at(0).attributes()
		is_public:      c.flag(flag_is_public)
		is_union:       c.flag(flag_is_union)
		implements:     c.list_at(1).type_exprs()
		embedded:       c.list_at(2).type_exprs()
		language:       unsafe { Language(int(c.aux())) }
		name:           c.name()
		generic_params: c.list_at(3).type_exprs()
		fields:         c.list_at(4).field_decls(true)
		pos:            c.pos()
	}
}

// type_decl reads a stmt_type_decl cursor.
pub fn (c Cursor) type_decl() TypeDecl {
	if !c.is_valid() || c.kind() != .stmt_type_decl {
		return TypeDecl{}
	}
	return TypeDecl{
		is_public:      c.flag(flag_is_public)
		language:       unsafe { Language(int(c.aux())) }
		name:           c.name()
		generic_params: c.list_at(2).type_exprs()
		base_type:      c.edge(0).type_expr()
		variants:       c.list_at(3).type_exprs()
	}
}

fn attrs_from_cursor(list CursorList) []Attribute {
	return list.attributes()
}

fn fn_type_from_cursor(c Cursor) FnType {
	if !c.is_valid() || c.kind() != .typ_fn {
		return FnType{}
	}
	return FnType{
		generic_params: c.list_at(0).type_exprs()
		params:         parameter_list_from_cursor(c.list_at(1))
		return_type:    c.edge(2).type_expr()
	}
}

fn type_exprs_from_edges(c Cursor, start int) []Expr {
	if !c.is_valid() || start >= c.edge_count() {
		return []Expr{}
	}
	mut out := []Expr{cap: c.edge_count() - start}
	for i in start .. c.edge_count() {
		out << c.edge(i).type_expr()
	}
	return out
}

fn parameter_from_cursor(c Cursor) Parameter {
	if !c.is_valid() {
		return Parameter{}
	}
	return Parameter{
		name:   c.name()
		typ:    c.edge(0).type_expr()
		is_mut: c.flag(flag_is_mut)
		pos:    c.pos()
	}
}

fn parameter_list_from_cursor(list CursorList) []Parameter {
	mut out := []Parameter{cap: list.len()}
	for i in 0 .. list.len() {
		out << parameter_from_cursor(list.at(i))
	}
	return out
}

fn field_decl_type_list(list CursorList) []FieldDecl {
	mut out := []FieldDecl{cap: list.len()}
	for i in 0 .. list.len() {
		field := list.at(i)
		if !field.is_valid() {
			continue
		}
		out << FieldDecl{
			name:                field.name()
			typ:                 field.edge(0).type_expr()
			is_public:           field.flag(flag_is_public)
			is_mut:              field.flag(flag_is_mut)
			is_module_mut:       field.flag(flag_field_is_module_mut)
			is_interface_method: field.flag(flag_field_is_interface_method)
		}
	}
	return out
}

// FileCursor is a typed wrapper over a FlatFile entry. It exposes the
// file-level metadata (name, mod, selector_names) and the three top-level
// child lists (attributes, imports, stmts) without rehydrating a full
// ast.File.
pub struct FileCursor {
pub:
	flat &FlatAst = unsafe { nil }
	idx  int // index into flat.files
}

// file_cursor returns a FileCursor over the i-th FlatFile in this FlatAst.
@[inline]
pub fn (flat &FlatAst) file_cursor(idx int) FileCursor {
	return FileCursor{
		flat: unsafe { flat }
		idx:  idx
	}
}

// file_cursors returns one FileCursor per FlatFile. Allocates an int-sized
// array; for hot loops prefer `for i in 0 .. flat.files.len { flat.file_cursor(i) }`.
pub fn (flat &FlatAst) file_cursors() []FileCursor {
	mut out := []FileCursor{cap: flat.files.len}
	for i in 0 .. flat.files.len {
		out << flat.file_cursor(i)
	}
	return out
}

@[inline]
pub fn (fc FileCursor) flat_file() FlatFile {
	return fc.flat.files[fc.idx]
}

// root returns a Cursor positioned at this file's root FlatNode (kind == .file).
@[inline]
pub fn (fc FileCursor) root() Cursor {
	return Cursor{
		flat: fc.flat
		id:   fc.flat.files[fc.idx].file_id
	}
}

@[inline]
pub fn (fc FileCursor) name() string {
	return fc.flat.string_at(fc.flat.files[fc.idx].name_idx)
}

@[inline]
pub fn (fc FileCursor) mod() string {
	return fc.flat.string_at(fc.flat.files[fc.idx].mod_idx)
}

@[inline]
pub fn (fc FileCursor) selector_names() map[int]string {
	return fc.flat.files[fc.idx].selector_names
}

// attrs returns the file's top-level attribute list (edge 0 of the file node).
@[inline]
pub fn (fc FileCursor) attrs() CursorList {
	return fc.root().list_at(0)
}

// imports returns the file's top-level import list (edge 1 of the file node).
@[inline]
pub fn (fc FileCursor) imports() CursorList {
	return fc.root().list_at(1)
}

// stmts returns the file's top-level statement list (edge 2 of the file node).
@[inline]
pub fn (fc FileCursor) stmts() CursorList {
	return fc.root().list_at(2)
}