package semantic

import (
	"fmt"
	"strconv"

	"git.sr.ht/~mna/snow/pkg/ast"
	"git.sr.ht/~mna/snow/pkg/token"
)

// The translate pass translates the AST produced by the parser to the semantic
// tree, which has less noise and more dense and precise information about the
// program.
//
// After this pass, the semantic graph is still high-level (close to the source
// AST) but has no comment and scopes are resolved, but not types nor
// identifier references.
func translate(files []*ast.File, errh func(token.Pos, string)) *Unit {
	unit := &Unit{
		Files: make([]*File, 0, len(files)),
	}
	universe := NewUniverse(unit)
	unit.scope = universe

	tv := &translateVisitor{
		errh:      errh,
		scope:     universe,
		generated: unit,
	}
	for _, astf := range files {
		ast.Walk(tv, astf)
	}
	return unit
}

type translateVisitor struct {
	errh              func(token.Pos, string)
	preventBlockScope bool   // transitive flag to prevent a block from opening a new scope (because already opened by parent)
	scope             *Scope // current scope

	generated interface{}
}

func (t *translateVisitor) cloneNewScope(owner Node) *translateVisitor {
	return &translateVisitor{
		errh:  t.errh,
		scope: t.scope.New(owner),
	}
}

func (t *translateVisitor) registerSymbol(ident *ast.Ident, decl Decl, lbl string) {
	if !t.scope.Register(ident.Name, decl) {
		t.errh(ident.Pos(), fmt.Sprintf("%s %s: symbol already declared in this scope", lbl, ident.Name))
		return
	}
}

func (t *translateVisitor) registerGenericType(elem *GenericElem) {
	if !t.scope.Register(elem.Ident(), elem) {
		t.errh(elem.Pos(), fmt.Sprintf("generic type %s: symbol already declared in this scope", elem.Ident()))
		return
	}
}

func (t *translateVisitor) Visit(n ast.Node) ast.Visitor {
	switch n := n.(type) {

	case *ast.File:
		// ************** FILE *****************

		var f File
		f.pos = n.Pos()
		f.scope = t.scope // the file is in the Unit's scope

		tt := t.cloneNewScope(&f)
		for _, decl := range n.Decls {
			ast.Walk(tt, decl)
			switch gn := tt.generated.(type) {
			case []*Var:
				f.Vars = append(f.Vars, gn...)
			case *Fn:
				f.Fns = append(f.Fns, gn)
			case *Struct:
				f.Structs = append(f.Structs, gn)
			default:
				panic(fmt.Sprintf("invalid top-level declaration in file: %T", gn))
			}
		}
		unit := t.generated.(*Unit)
		unit.Files = append(unit.Files, &f)

		// ************** DECLARATIONS *****************

	case *ast.FnDecl:
		var fn Fn
		fn.pos = n.Name.Pos()
		fn.scope = t.scope
		fn.ident = n.Name.Name
		fn.IsRef = n.Ref.IsValid()
		fn.Attrs = make([]*Call, len(n.Attrs))
		fn.Params = make([]*Var, len(n.Signature.Params))
		t.registerSymbol(n.Name, &fn, token.Fn.String())

		// attributes are in the same scope as the fn
		for i, attr := range n.Attrs {
			ast.Walk(t, attr)
			call := t.generated.(*Call)
			call.AttrOf = &fn
			fn.Attrs[i] = call
		}

		// Return type is in a new scope, below the fn's, but above the body's and
		// params, that the generic clause is also part of. That's because the
		// return type could be $T, and if so it must refer to its generic clause
		// (unless it is a method, in which case it should resolve to the struct's
		// generic clause).
		if str, ok := fn.scope.Owner.(*Struct); ok {
			fn.MethodOf = str
		}

		tgen := t
		if gp := n.GenericParams; gp != nil {
			if fn.MethodOf != nil {
				// TODO: should that really be prevented?
				t.errh(n.GenericParams.Pos(), "invalid generic clause on struct method")
			} else {
				tgen = t.cloneNewScope(&fn)
				fn.GenericParams = tgen.buildExplicitGenericClause(gp)
			}
		}

		if n.Signature.RetType != nil {
			ast.Walk(tgen, n.Signature.RetType)
			fn.ReturnExpr = tgen.generated.(Expr)
		}

		// parameters and body (possibly missing for external functions) are in the body's scope
		tbody := tgen.cloneNewScope(&fn)
		tbody.preventBlockScope = true
		if fn.MethodOf != nil {
			generateSelfVar(&fn, n.Body, tbody.scope)
		}
		for i, param := range n.Signature.Params {
			ast.Walk(tbody, param)
			v := tbody.generated.(*Var)
			v.ParamOf = &fn
			fn.Params[i] = v
		}
		if n.Body != nil {
			ast.Walk(tbody, n.Body)
			fn.Body = tbody.generated.(*Block)
		}
		t.generated = &fn

	case *ast.VarDecl:
		ctx := Immutable
		if n.Kw == token.Var {
			ctx = Mutable
		}

		list := make([]*Var, len(n.Vars))
		for i, vd := range n.Vars {
			v := t.buildVarOrExpr(n.Kw, vd.Name, vd.Type).(*Var)
			v.Ctx = ctx
			if vd.Value != nil {
				ast.Walk(t, vd.Value)
				v.Value = t.generated.(Expr)
			}
			if str, ok := v.scope.Owner.(*Struct); ok {
				v.PropOf = str
			}
			list[i] = v
		}
		t.generated = list

	case *ast.StructDecl:
		var str Struct
		str.pos = n.Name.Pos()
		str.scope = t.scope
		str.ident = n.Name.Name
		t.registerSymbol(n.Name, &str, token.Struct.String())

		tt := t.cloneNewScope(&str)
		str.BodyScope = tt.scope

		if gp := n.GenericParams; gp != nil {
			str.GenericParams = tt.buildExplicitGenericClause(gp)
		}

		for _, decl := range n.Decls {
			ast.Walk(tt, decl)
			switch gn := tt.generated.(type) {
			case []*Var:
				str.Vars = append(str.Vars, gn...)
			case *Fn:
				str.Fns = append(str.Fns, gn)
			case *Struct:
				str.Structs = append(str.Structs, gn)
			default:
				panic(fmt.Sprintf("invalid declaration in struct %s: %T", str.ident, gn))
			}
		}
		t.generated = &str

		// ************** STATEMENTS *****************

	case *ast.IfStmt:
		var ifs If
		ifs.pos = n.Pos()
		ifs.scope = t.scope // the scope containing the if, not the scope of the if body

		// declarations in conditions and if body are in a new scope
		tt := t.cloneNewScope(&ifs)
		for _, cond := range n.Conds {
			ast.Walk(tt, cond)
			ifs.Conds = append(ifs.Conds, tt.generated.(Expr))
		}
		if n.Body != nil {
			tt.preventBlockScope = true
			ast.Walk(tt, n.Body)
			ifs.Body = tt.generated.(*Block)
		}

		// the else starts in the same scope as the if, and the new scope will be open
		// by walking it - it's an If or a Block.
		if n.Else != nil {
			ast.Walk(t, n.Else)
			ifs.Else = t.generated.(Stmt)
		}
		t.generated = &ifs

	case *ast.GuardStmt:
		var g Guard
		g.pos = n.Pos()
		g.scope = t.scope

		// declarations in conditions are in the same scope as the guard
		for _, cond := range n.Conds {
			ast.Walk(t, cond)
			g.Conds = append(g.Conds, t.generated.(Expr))
		}

		// the Else is a Block opened by the Guard, but has no other symbols than those
		// in the block.
		if n.Else != nil {
			tt := t.cloneNewScope(&g)
			tt.preventBlockScope = true
			ast.Walk(tt, n.Else)
			g.Else = tt.generated.(*Block)
		}
		t.generated = &g

	case *ast.Block:
		preventScope := t.preventBlockScope
		t.preventBlockScope = false

		var b Block
		b.pos = n.Pos()
		b.scope = t.scope

		tt := t
		if !preventScope {
			tt = t.cloneNewScope(&b)
		}
		for _, stmt := range n.Stmts {
			ast.Walk(tt, stmt)
			switch nn := tt.generated.(type) {
			case Stmt:
				b.Stmts = append(b.Stmts, nn)
			case []*Var:
				for _, stmt := range nn {
					b.Stmts = append(b.Stmts, stmt)
				}
			default:
				panic(fmt.Sprintf("invalid generated block node: %T", nn))
			}
		}
		t.generated = &b

	case *ast.ReturnStmt:
		var r Return
		r.pos = n.Pos()
		r.scope = t.scope

		if n.Value != nil {
			ast.Walk(t, n.Value)
			r.Value = t.generated.(Expr)
		}
		t.generated = &r

	case *ast.AssignStmt:
		var a Assign
		a.pos = n.Pos()
		a.scope = t.scope
		a.Op = token.Assign

		if n.Left != nil {
			ast.Walk(t, n.Left)
			a.Left = t.generated.(Expr)
		}
		if n.Right != nil {
			ast.Walk(t, n.Right)
			a.Right = t.generated.(Expr)
		}
		t.generated = &a

	case *ast.ExprStmt:
		var e ExprStmt
		e.pos = n.Pos()
		e.scope = t.scope

		if n.Value != nil {
			ast.Walk(t, n.Value)
			e.Value = t.generated.(Expr)
		}
		t.generated = &e

		// ************** EXPRESSIONS *****************

	case *ast.FnType:
		var fn FnTypeExpr
		fn.pos = n.Pos()
		fn.scope = t.scope
		fn.Params = make([]Expr, len(n.Params))

		for i, p := range n.Params {
			ast.Walk(t, p)
			fn.Params[i] = t.generated.(Expr)
		}
		if n.RetType != nil {
			ast.Walk(t, n.RetType)
			fn.Return = t.generated.(Expr)
		}
		t.generated = &fn

	case *ast.TupleType:
		var tup TupleTypeExpr
		tup.pos = n.Pos()
		tup.scope = t.scope
		tup.Fields = make([]Expr, len(n.Types))

		for i, field := range n.Types {
			ast.Walk(t, field)
			tup.Fields[i] = t.generated.(Expr)
		}
		t.generated = &tup

	case *ast.TupleExpr:
		var tup TupleVal
		tup.pos = n.Pos()
		tup.scope = t.scope
		tup.Values = make([]Expr, len(n.Values))

		for i, val := range n.Values {
			ast.Walk(t, val)
			tup.Values[i] = t.generated.(Expr)
		}
		t.generated = &tup

	case *ast.BinaryExpr:
		var bin Binary
		bin.pos = n.Pos()
		bin.scope = t.scope
		bin.Op = n.Op

		if n.Left != nil {
			ast.Walk(t, n.Left)
			bin.Left = t.generated.(Expr)
		}
		if n.Right != nil {
			ast.Walk(t, n.Right)
			bin.Right = t.generated.(Expr)
		}
		t.generated = &bin

	case *ast.UnaryExpr:
		var un Unary
		un.pos = n.Pos()
		un.scope = t.scope
		un.Op = n.Op

		if n.Right != nil {
			ast.Walk(t, n.Right)
			un.Right = t.generated.(Expr)
		}
		t.generated = &un

	case *ast.ParenExpr:
		var par Paren
		par.pos = n.Pos()
		par.scope = t.scope

		if n.Value != nil {
			ast.Walk(t, n.Value)
			par.Value = t.generated.(Expr)
		}
		t.generated = &par

	case *ast.CallExpr:
		var call Call
		call.pos = n.Pos()
		call.scope = t.scope
		call.Labels = make([]string, len(n.Args))
		call.Args = make([]Expr, len(n.Args))

		if n.Left != nil {
			ast.Walk(t, n.Left)
			call.Fun = t.generated.(Expr)
		}
		for i, arg := range n.Args {
			ast.Walk(t, arg)
			if arg.Label != nil {
				call.Labels[i] = arg.Label.Name
			}
			call.Args[i] = t.generated.(Expr)
		}
		t.generated = &call

	case *ast.SelectorExpr:
		var sel Selector
		sel.pos = n.Pos()
		sel.scope = t.scope

		if n.Left != nil {
			ast.Walk(t, n.Left)
			sel.Left = t.generated.(Expr)
		}
		if n.Sel != nil {
			ast.Walk(t, n.Sel)
			switch expr := t.generated.(type) {
			case *GenericInst:
				sel.Sel = expr
			case *Ident:
				sel.Sel = expr
			default:
				panic(fmt.Sprintf("invalid node type for right-hand side of selector: %T", expr))
			}
		}
		t.generated = &sel

	case *ast.GenericIdent:
		var gi GenericInst
		gi.pos = n.Pos()
		gi.scope = t.scope
		if n.Ident != nil {
			ast.Walk(t, n.Ident)
			gi.GenericDecl = t.generated.(*Ident)
		}
		if n.GenericArgs != nil {
			gi.TypeExprs = make([]Expr, len(n.GenericArgs.Types))
			for i, typ := range n.GenericArgs.Types {
				ast.Walk(t, typ.Type)
				gi.TypeExprs[i] = t.generated.(Expr)
			}
		}
		t.generated = &gi

	case *ast.Ident:
		var id Ident
		id.pos = n.Pos()
		id.scope = t.scope
		id.Name = n.Name
		id.Index = -1
		if ix, err := strconv.Atoi(n.Name); err == nil && ix >= 0 {
			id.Index = ix
		}
		t.generated = &id

	case *ast.BasicLit:
		var ce commonExpr
		ce.pos = n.Pos()
		ce.scope = t.scope

		var expr Expr
		switch n.Literal {
		case token.String:
			// TODO: currently, the string lit supports the same escapes as Go's
			val, err := strconv.Unquote(n.Value)
			if err != nil {
				panic(fmt.Sprintf("invalid string literal value for %s: %s", n.Value, err))
			}
			expr = &LitString{
				commonExpr: ce,
				Repr:       n.Value,
				Value:      val,
			}

		case token.Int:
			// TODO: currently, the integer lit only supports digits, but could be underscore and prefixes in the future
			val, err := strconv.Atoi(n.Value)
			if err != nil {
				panic(fmt.Sprintf("invalid integer literal value for %s: %s", n.Value, err))
			}
			expr = &LitInt{
				commonExpr: ce,
				Repr:       n.Value,
				Value:      val,
			}

		default:
			panic(fmt.Sprintf("invalid literal token: %s", n.Literal))
		}
		t.generated = expr

		// ************** SUPPORTING NODES *****************

	case *ast.ParamDef:
		// ParamDef are treated just like VarDecl if there is a Name, otherwise
		// just an Expr (the type).
		ve := t.buildVarOrExpr(token.Let, n.Name, n.Type)
		if v, ok := ve.(*Var); ok {
			v.Ctx = Immutable
		}
		t.generated = ve

	case *ast.ExprItem:
		// ExprItem ignores the label (it is extracted directly in CallExpr, the
		// only place where this is allowed) and returns just the Expr.
		ast.Walk(t, n.Expr)

	case *ast.Attr:
		// Attributes are translated to Call expressions.
		var call Call
		call.pos = n.Pos()
		call.scope = t.scope
		call.Labels = make([]string, len(n.Fields))
		call.Args = make([]Expr, len(n.Fields))

		if n.Name != nil {
			ast.Walk(t, n.Name)
			call.Fun = t.generated.(Expr)
		}
		for i, field := range n.Fields {
			ast.Walk(t, field.Value)
			if field.Key != nil {
				call.Labels[i] = field.Key.Name
			}
			call.Args[i] = t.generated.(Expr)
		}
		t.generated = &call

	case *ast.ElseClause:
		return t

	default:
		if n != nil {
			panic(fmt.Sprintf("invalid AST node type: %T", n))
		}
	}
	return nil
}

func (t *translateVisitor) buildExplicitGenericClause(astGC *ast.GenericClause) *GenericClause {
	gc := &GenericClause{
		Explicit: true,
		Elems:    make([]*GenericElem, len(astGC.Types)),
	}

	for i, param := range astGC.Types {
		ast.Walk(t, param)
		var ge GenericElem
		v := t.generated.(*Var)
		ge.pos = v.Pos()
		ge.scope = v.Scope()
		ge.ident = v.Ident()
		gc.Elems[i] = &ge
		t.registerGenericType(&ge)
	}
	return gc
}

func (t *translateVisitor) buildVarOrExpr(kw token.Token, ident *ast.Ident, typ ast.Expr) Node {
	var v Var
	v.scope = t.scope
	if typ != nil {
		ast.Walk(t, typ)
		v.TypeExpr = t.generated.(Expr)
	}

	if ident != nil {
		v.pos = ident.Pos()
		v.ident = ident.Name
		if !ident.IsGenericTypeName() {
			t.registerSymbol(ident, &v, kw.String())
		}
		return &v
	}
	return v.TypeExpr
}

func generateSelfVar(fn *Fn, body *ast.Block, bodyScope *Scope) {
	// nothing to do if there is no body
	if body == nil {
		return
	}

	// just register the "self" var in the scope
	var self Var
	self.Ctx = Immutable
	if fn.IsRef {
		self.Ctx = Mutable
	}
	self.pos = body.Pos()
	self.scope = bodyScope
	self.ident = SelfVarName
	self.typ = &StructType{Decl: fn.MethodOf}
	bodyScope.Register(SelfVarName, &self)
}