package semantic

import (
	"bytes"
	"fmt"
	"io"
	"sort"
	"strings"

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

// Scope registers declared identifiers in their lexical scope. Scopes form
// a tree rooted in the Universe scope (the scope of the Unit node).
type Scope struct {
	// Parent is the parent scope of this scope, it is nil for the Universe
	// scope.
	Parent *Scope
	// Children is the list of immediate children of this scope.
	Children []*Scope
	// Owner is the node that created this scope. It can be *Unit, *File,
	// *Fn, *If, *Guard, *Struct or *Block.
	Owner Node
	// Symbols is the registry of symbols-to-Object mapping in this scope.
	Symbols map[string]Decl
	// ID is a unique identifier for the scope. A scope's parents are guaranteed
	// to have a smaller ID, and conversely a scope's children are guaranteed
	// to have a greater ID.
	ID int
	// OrderIndependent indicates if the lookup of symbols is order-
	// dependent (i.e. if a symbol needs to be declared before use).
	OrderIndependent bool

	nextID func() int
}

// NewUniverse creates a new universe scope.
func NewUniverse(owner *Unit) *Scope {
	s := newScope(owner, nil)
	// temporarily allow registering universe objects
	s.Parent = s
	defer func() { s.Parent = nil }()

	// TODO: not super happy with how the pre-declared universe identifiers
	// are handled, refactor this...

	for id, decl := range universeIdents {
		vv := decl.(*Var)
		vv.scope = s
		s.Register(id, decl)
	}

	// add the @extern attribute as a struct - added here as it needs
	// to be able to create its child scope.
	// TODO: nothing prevents using this struct in non-attribute situations.
	// Is that ok? If so it should get codegen'd.
	varNms := []string{"import", "pkg", "symbol"}
	str := &Struct{Vars: make([]*Var, len(varNms))}
	str.ident = ExternAttrName
	str.typ = &StructType{Decl: str}
	str.scope = s
	strScope := s.New(str)
	for i, varNm := range varNms {
		var v Var
		v.ident = varNm
		v.typ = &BasicType{Kind: String}
		v.scope = strScope
		v.Ctx = Immutable
		if varNm == "pkg" {
			// NOTE: set pkg default value to empty string, just so it doesn't have to be specified
			// Eventually, this should be a custom init where only import and symbol are required.
			v.Value = &LitString{Repr: `""`}
		}
		str.Vars[i] = &v
		strScope.Register(varNm, &v)
	}
	str.BodyScope = strScope
	s.Register(ExternAttrName, str)

	return s
}

// New creates a new child scope for s, opened by node.
func (s *Scope) New(owner Node) *Scope {
	return newScope(owner, s)
}

func newScope(owner Node, parent *Scope) *Scope {
	var orderInd bool

	switch owner.(type) {
	case *Unit:
		if parent != nil {
			panic(fmt.Sprintf("non-universe scope created with invalid owner %T", owner))
		}
		orderInd = true

	case *File, *Struct:
		if parent == nil {
			panic(fmt.Sprintf("universe scope created with invalid owner %T", owner))
		}
		orderInd = true

	case *Fn, *If, *Guard, *Block:
		if parent == nil {
			panic(fmt.Sprintf("universe scope created with invalid owner %T", owner))
		}
	}

	s := &Scope{
		Parent:           parent,
		Owner:            owner,
		OrderIndependent: orderInd,
		Symbols:          make(map[string]Decl),
	}
	if parent == nil {
		var id int
		s.nextID = func() int {
			id++
			return id
		}
	} else {
		s.nextID = parent.nextID
		parent.Children = append(parent.Children, s)
	}
	s.ID = s.nextID()
	return s
}

// IsUniverse returns true if s is the universe scope.
func (s *Scope) IsUniverse() bool {
	return s.Parent == nil
}

// IsTopLevel returns true if s is a top-level (file) scope.
func (s *Scope) IsTopLevel() bool {
	return s.Parent != nil && s.Parent.IsUniverse()
}

// Register adds the identifier to the scope. It returns true if the symbol was
// added, false if it already existed in this scope.
func (s *Scope) Register(ident string, decl Decl) bool {
	if s.IsUniverse() {
		panic(fmt.Sprintf("attempt to register %s into universe scope", ident))
	}

	if _, ok := s.Symbols[ident]; ok {
		return false
	}
	s.Symbols[ident] = decl
	return true
}

// Lookup looks up the symbol in that scope, without looking at the parent
// chain. It returns the declaration if the symbol exists, or nil otherwise.
// It doesn't care about position, because when looking up a symbol in a
// specific scope, only one such symbol can exist.
func (s *Scope) Lookup(symbol string) Decl {
	return s.Symbols[symbol]
}

// LookupChain looks up the symbol in that scope or any parent scopes. It
// returns the declaration if found, or nil.  If a valid pos is provided, the
// existing object will be returned only if it was declared at or before that
// position, except if the scope is order-independent.
func (s *Scope) LookupChain(symbol string, pos token.Pos) Decl {
	for ; s != nil; s = s.Parent {
		if decl := s.Symbols[symbol]; decl != nil && (s.OrderIndependent || !pos.IsValid() || decl.Pos() <= pos) {
			return decl
		}
	}
	return nil
}

// WriteTo prints the scope and its contents to w, indented by n copies of a
// spacing prefix. If recurse is true, children scopes are printed too, with
// increasing indentation.
func (s *Scope) WriteTo(w io.Writer, n int, recurse bool) {
	const ind = ".  "
	indn := strings.Repeat(ind, n)

	fmt.Fprintf(w, "%s%d %T {\n", indn, s.ID, s.Owner)

	syms := make([]string, 0, len(s.Symbols))
	for sym := range s.Symbols {
		syms = append(syms, sym)
	}
	sort.Strings(syms)

	indn1 := indn + ind
	for _, sym := range syms {
		fmt.Fprintf(w, "%s%s\n", indn1, sym)
	}

	if recurse {
		for _, child := range s.Children {
			child.WriteTo(w, n+1, recurse)
		}
	}

	fmt.Fprintf(w, "%s}\n", indn)
}

// String returns a string representation of the scope and its symbols, without
// children scopes.
func (s *Scope) String() string {
	var buf bytes.Buffer
	s.WriteTo(&buf, 0, false)
	return buf.String()
}

// the list of identifiers to add in the universe scope.
var universeIdents = func() map[string]Decl {
	m := make(map[string]Decl)

	for b := kindStart + 1; b < kindEnd; b++ {
		v := &Var{Ctx: Typ}
		v.typ = &BasicType{Kind: b}
		v.ident = b.String()
		m[v.Ident()] = v
	}

	// add true and false constants
	for _, id := range []string{TrueLitName, FalseLitName} {
		v := &Var{Ctx: Immutable}
		v.typ = &BasicType{Kind: Bool}
		v.ident = id
		m[v.Ident()] = v
	}

	return m
}()