package ast

import (
	"encoding/json"
	"fmt"
	"strings"

	"git.usebox.net/micro-lang/tokens"
)

// for JSON
type node struct {
	Type   string `json:"type"`
	Object any    `json:"object"`
}

var (
	TypeNumber = Type{Value: tokens.Token{Id: tokens.TNumber}, Const: true}
	TypeBool   = Type{Value: tokens.Token{Id: tokens.TBool}, Const: true}
	TypeString = Type{Value: tokens.Token{Id: tokens.TString}, Const: true}
	TypeArray  = Type{Value: tokens.Token{Id: tokens.TArray}, Const: true}
	TypeNone   = Type{Value: tokens.Token{Id: tokens.None}, Const: true}
)

type Type struct {
	Value  tokens.Token `json:"value"`
	Params []Type       `json:"params,omitempty"`
	Ret    *Type        `json:"type,omitempty"`
	Len    *int         `json:"len,omitempty"`
	Const  bool         `json:"const"`
	Inst   bool         `json:"-"`
	Ref    bool         `json:"-"`
}

func NewStructType(name tokens.Token) *Type {
	return &Type{Value: tokens.Token{Id: tokens.TStruct, Value: name.Value, Loc: name.Loc}}
}

func NewFuncType(name tokens.Token, params []Type, ret *Type) *Type {
	return &Type{Value: tokens.Token{Id: tokens.TFunc, Loc: name.Loc}, Params: params, Ret: ret}
}

func (t Type) MarshalJSON() ([]byte, error) {
	type Object Type
	return json.MarshalIndent(node{"type", Object(t)}, "", "  ")
}

func (t Type) String() string {
	if t.Value.Id == tokens.TFunc {
		var b = strings.Builder{}
		b.WriteString(t.Value.Id.String())
		if t.Params == nil {
			// not specified (e.g. println)
			b.WriteString(" (...)")
		} else if len(t.Params) > 0 {
			b.WriteString(" (")
			for n, param := range t.Params {
				b.WriteString(param.String())
				if n < len(t.Params)-1 {
					b.WriteString(", ")
				}
			}
			b.WriteString(")")
		} else {
			b.WriteString(" ()")
		}

		if t.Ret != nil {
			b.WriteString(" ")
			b.WriteString(t.Ret.String())
		}

		return b.String()
	}

	if t.Value.Id == tokens.TStruct {
		return fmt.Sprintf("struct %s", t.Value.Value)
	}

	if t.Value.Id == tokens.TArray {
		if t.Ret != nil {
			return fmt.Sprintf("array [%d]%s", *t.Len, t.Ret.String())
		}
		return "array"
	}

	return t.Value.Id.String()
}

func (t Type) Compare(other *Type) bool {
	if other == nil {
		return false
	}
	if t.Value.Id != other.Value.Id {
		return false
	}

	if t.Value.Id == tokens.TArray {
		// this is len accepting any length of array
		if t.Len == nil || other.Len == nil {
			return true
		}
		return *t.Len == *other.Len
	}

	if t.Value.Id == tokens.TStruct {
		return t.Value.Value == other.Value.Value
	}

	if t.Value.Id == tokens.TFunc {
		// func checks
		if len(t.Params) != len(other.Params) {
			return false
		}
		for i := range t.Params {
			if !t.Params[i].Compare(&other.Params[i]) {
				return false
			}
		}
		if t.Ret != nil {
			return t.Ret.Compare(other.Ret)
		}
		return t.Ret == other.Ret
	}

	return true
}

func (t Type) ResolveConst() bool {
	if t.Value.Id == tokens.TArray {
		return t.Ret.Const
	}
	return t.Const
}

type Expr interface {
	Resolve() *Type
}

type ExprList struct {
	Exprs []Expr `json:"exprs"`
}

func (n ExprList) MarshalJSON() ([]byte, error) {
	type Object ExprList
	return json.MarshalIndent(node{"exprlist", Object(n)}, "", "  ")
}

func (n ExprList) Resolve() *Type {
	return n.Exprs[0].Resolve()
}

type Var struct {
	Name  tokens.Token `json:"ident"`
	Type  Type         `json:"type"`
	Initv Expr         `json:"initval,omitempty"`
}

func (n Var) MarshalJSON() ([]byte, error) {
	type Object Var
	return json.MarshalIndent(node{"var", Object(n)}, "", "  ")
}

func (n Var) Resolve() *Type {
	return &n.Type
}

type Variable struct {
	Name  tokens.Token `json:"ident"`
	Index Expr         `json:"index,omitempty"`
	Type  Type         `json:"-"`
}

func (n Variable) MarshalJSON() ([]byte, error) {
	type Object Variable
	return json.MarshalIndent(node{"variable", Object(n)}, "", "  ")
}

func (n Variable) Resolve() *Type {
	if n.Type.Value.Id == tokens.TArray && n.Index != nil {
		return n.Type.Ret
	}
	return &n.Type
}

func (n Variable) String() string {
	if n.Index != nil {
		return fmt.Sprintf("%s[%s]", n.Name.Value, n.Index.Resolve().Value)
	}
	if n.Type.Value.Id == tokens.TStruct {
		return fmt.Sprintf("<struct %s>.%s", n.Type.Value.Value, n.Name.Value)
	}
	return n.Name.String()
}

type Assign struct {
	Left  Expr            `json:"left"`
	Right Expr            `json:"right"`
	Loc   tokens.Location `json:"-"`
}

func (n Assign) MarshalJSON() ([]byte, error) {
	type Object Assign
	return json.MarshalIndent(node{"variable", Object(n)}, "", "  ")
}

func (n Assign) Resolve() *Type {
	return n.Left.Resolve()
}

type Binary struct {
	Left  Expr         `json:"left"`
	Op    tokens.Token `json:"op"`
	Right Expr         `json:"right"`
	Type  *Type
}

func (n Binary) MarshalJSON() ([]byte, error) {
	type Object Binary
	return json.MarshalIndent(node{"binary", Object(n)}, "", "  ")
}

func (n Binary) Resolve() *Type {
	if n.Type != nil {
		return n.Type
	} else {
		return n.Left.Resolve()
	}
}

type Unary struct {
	Op    tokens.Token `json:"op"`
	Right Expr         `json:"expr"`
}

func (n Unary) MarshalJSON() ([]byte, error) {
	type Object Unary
	return json.MarshalIndent(node{"unary", Object(n)}, "", "  ")
}

func (n Unary) Resolve() *Type {
	if n.Op.Id == tokens.TestE {
		return &TypeBool
	}
	return n.Right.Resolve()
}

type Number int64

func (n Number) String() string {
	return fmt.Sprintf("%d", n)
}

func (n Number) Resolve() *Type {
	return &TypeNumber
}

type Literal struct {
	Value  tokens.Token `json:"value"`
	Numval Number       `json:"numval,omitempty"`
	Type   Type         `json:"-"`
}

func (n Literal) MarshalJSON() ([]byte, error) {
	type Object Literal
	return json.MarshalIndent(node{"literal", Object(n)}, "", "  ")
}

func (n Literal) Resolve() *Type {
	return &n.Type
}

type Group struct {
	Expr Expr `json:"expr"`
}

func (n Group) MarshalJSON() ([]byte, error) {
	type Object Group
	return json.MarshalIndent(node{"group", Object(n)}, "", "  ")
}

func (n Group) Resolve() *Type {
	return n.Expr.Resolve()
}

type Block struct {
	Stmts []Expr `json:"stmts"`
}

func (n Block) MarshalJSON() ([]byte, error) {
	type Object Block
	return json.MarshalIndent(node{"block", Object(n)}, "", "  ")
}

func (n Block) Resolve() *Type {
	// not an expression
	return &TypeNone
}

type For struct {
	Cond  Expr `json:"cond,omitempty"`
	Stmts Expr `json:"stmts"`
}

func (n For) MarshalJSON() ([]byte, error) {
	type Object For
	return json.MarshalIndent(node{"for", Object(n)}, "", "  ")
}

func (n For) Resolve() *Type {
	// not an expression
	return &TypeNone
}

type ForIn struct {
	Name  tokens.Token `json:"ident"`
	Expr  Expr         `json:"expr"`
	Stmts Expr         `json:"stmts"`
}

func (n ForIn) MarshalJSON() ([]byte, error) {
	type Object ForIn
	return json.MarshalIndent(node{"forin", Object(n)}, "", "  ")
}

func (n ForIn) Resolve() *Type {
	// not an expression
	return &TypeNone
}

type IfElse struct {
	Cond  Expr `json:"cond"`
	True  Expr `json:"true"`
	False Expr `json:"false"`
}

func (n IfElse) MarshalJSON() ([]byte, error) {
	type Object IfElse
	return json.MarshalIndent(node{"ifelse", Object(n)}, "", "  ")
}

func (n IfElse) Resolve() *Type {
	// not an expression
	return &TypeNone
}

type Call struct {
	Callee Expr            `json:"callee"`
	Loc    tokens.Location `json:"-"`
	Args   []Expr          `json:"args"`
}

func (n Call) MarshalJSON() ([]byte, error) {
	type Object Call
	return json.MarshalIndent(node{"call", Object(n)}, "", "  ")
}

func (n Call) Resolve() *Type {
	typ := n.Callee.Resolve().Ret
	if typ == nil {
		typ = &TypeNone
	}
	return typ
}

type Struct struct {
	Name tokens.Token `json:"ident"`
	Body Block        `json:"body"`
}

func (n Struct) MarshalJSON() ([]byte, error) {
	type Object Struct
	return json.MarshalIndent(node{"struct", Object(n)}, "", "  ")
}

func (n Struct) Resolve() *Type {
	return NewStructType(n.Name)
}

type GetExpr struct {
	Object Expr `json:"object"`
	Expr   Expr `json:"expr"`
}

func (n GetExpr) MarshalJSON() ([]byte, error) {
	type Object GetExpr
	return json.MarshalIndent(node{"get", Object(n)}, "", "  ")
}

func (n GetExpr) Resolve() *Type {
	// this is likely to be wrong
	return n.Expr.Resolve()
}

type Func struct {
	Name   tokens.Token `json:"ident"`
	Params []Var        `json:"params"`
	Ret    *Type        `json:"ret"`
	Body   Block        `json:"body"`
}

func (n Func) MarshalJSON() ([]byte, error) {
	type Object Func
	return json.MarshalIndent(node{"func", Object(n)}, "", "  ")
}

func (n Func) Resolve() *Type {
	if n.Ret == nil {
		return &TypeNone
	} else {
		return n.Ret
	}
}

type Return struct {
	Value Expr            `json:"expr"`
	Loc   tokens.Location `json:"-"`
	Error bool            `json:"error"`
}

func (n Return) MarshalJSON() ([]byte, error) {
	type Object Return
	return json.MarshalIndent(node{"return", Object(n)}, "", "  ")
}

func (n Return) Resolve() *Type {
	return n.Value.Resolve()
}

type Continue struct {
	Loc tokens.Location `json:"-"`
}

func (n Continue) MarshalJSON() ([]byte, error) {
	type Object Continue
	return json.MarshalIndent(node{"continue", Object(n)}, "", "  ")
}

func (n Continue) Resolve() *Type {
	return &TypeNone
}

type Break struct {
	Loc tokens.Location `json:"-"`
}

func (n Break) MarshalJSON() ([]byte, error) {
	type Object Break
	return json.MarshalIndent(node{"break", Object(n)}, "", "  ")
}

func (n Break) Resolve() *Type {
	return &TypeNone
}