package scope

import (
	"fmt"
	"jinx/pkg/lang/vm/code"
)

type ScopeChain struct {
	nameToSymbol map[string]SymbolID

	symbolScopes   []SymbolScope // All other scopes are bound to this by ID.
	functionScopes []FunctionScope
}

func NewScopeChain() ScopeChain {
	symbolScopes := make([]SymbolScope, 1)
	symbolScopes[0] = NewSymbolScope()

	functionScopes := make([]FunctionScope, 1)
	functionScopes[0] = NewFunctionScope(0, "") // Root function to house top-scope sub units

	return ScopeChain{
		nameToSymbol: make(map[string]SymbolID),

		symbolScopes:   symbolScopes,
		functionScopes: functionScopes,
	}
}

func (sc *ScopeChain) CurrentScopeID() ScopeID {
	return ScopeID(len(sc.symbolScopes) - 1)
}

func (sc *ScopeChain) IsRootScope() bool {
	return sc.CurrentScopeID() == ScopeID(0)
}

func (sc *ScopeChain) Current() *SymbolScope {
	if len(sc.functionScopes) == 0 {
		panic("root scope is missing")
	}

	return &sc.symbolScopes[sc.CurrentScopeID()]
}

func (sc *ScopeChain) CurrentFunction() *FunctionScope {
	if len(sc.functionScopes) == 0 {
		panic("root function scope is missing")
	}

	return &sc.functionScopes[len(sc.functionScopes)-1]
}

func (sc *ScopeChain) Enter() {
	sc.symbolScopes = append(sc.symbolScopes, NewSymbolScope())
}

func (sc *ScopeChain) EnterFunction(unit code.Marker) {
	sc.Enter()
	sc.functionScopes = append(sc.functionScopes, NewFunctionScope(sc.CurrentScopeID(), unit))
}

func (sc *ScopeChain) Exit() {
	if sc.CurrentScopeID() == 0 {
		return
	}

	sc.symbolScopes[sc.CurrentScopeID()] = SymbolScope{}
	sc.symbolScopes = sc.symbolScopes[:sc.CurrentScopeID()]
}

func (sc *ScopeChain) Declare(name string) (int, bool) {
	// Check whether the symbol is already declared in any of the scopes.
	if _, ok := sc.nameToSymbol[name]; ok {
		return 0, false
	}

	current := sc.Current()
	indexInScope := len(current.variableSymbols)

	symbolID := SymbolID{
		symbolKind:   SymbolKindVariable,
		scopeID:      sc.CurrentScopeID(),
		indexInScope: indexInScope,
	}

	// Declare the symbol in the current scope.
	current.variableSymbols = append(current.variableSymbols, Symbol[SymbolVariable]{
		name: name,
		data: SymbolVariable{
			localIndex: indexInScope,
		},
	})

	sc.nameToSymbol[name] = symbolID

	return indexInScope, true
}

func (sc *ScopeChain) DeclareFunction(name string, args uint) (code.Marker, bool) {
	if _, ok := sc.nameToSymbol[name]; ok {
		return "", false
	}

	current := sc.Current()
	index := len(current.functionSymbols)

	symbolID := SymbolID{
		symbolKind:   SymbolKindFunction,
		scopeID:      sc.CurrentScopeID(),
		indexInScope: index,
	}

	unitName := sc.CreateFunctionSubUnit(name)

	current.functionSymbols = append(current.functionSymbols, Symbol[SymbolFunction]{
		name: name,
		data: SymbolFunction{
			marker: unitName,
			args:   args,
		},
	})

	sc.nameToSymbol[name] = symbolID

	return unitName, true
}

func (sc *ScopeChain) DeclareAnonymous() int {
	current := sc.Current()
	index := len(current.variableSymbols)

	current.variableSymbols = append(current.variableSymbols, Symbol[SymbolVariable]{
		name: "", // An anonymous symbol has no name.
		data: SymbolVariable{
			localIndex: index,
		},
	})

	return index
}

func (sc *ScopeChain) DeclareTemporary() int {
	return len(sc.Current().variableSymbols) // :)
}

func (sc *ScopeChain) CreateAnonymousFunctionSubUnit() code.Marker {
	fnScope := sc.CurrentFunction()

	index := fnScope.subUnitCount
	fnScope.subUnitCount++

	return sc.CreateFunctionSubUnit(fmt.Sprintf("anon_%d", index))
}

func (sc *ScopeChain) CreateFunctionSubUnit(subUnitName string) code.Marker {
	fnScope := sc.CurrentFunction()

	name := fnScope.unit
	if name == "" {
		name = code.Marker(subUnitName)
	} else {
		name = name.SubUnit(subUnitName)
	}

	return name
}

func (sc *ScopeChain) Lookup(name string) (SymbolID, bool) {
	if id, ok := sc.nameToSymbol[name]; ok {
		return id, true
	}

	return SymbolID{}, false
}

func (sc *ScopeChain) GetVariable(id SymbolID) Symbol[SymbolVariable] {
	if id.symbolKind != SymbolKindVariable {
		panic("incorrect symbol id kind given")
	}

	return sc.symbolScopes[id.scopeID].variableSymbols[id.indexInScope]
}

func (sc *ScopeChain) GetFunction(id SymbolID) Symbol[SymbolFunction] {
	if id.symbolKind != SymbolKindFunction {
		panic("incorrect symbol id kind given")
	}

	return sc.symbolScopes[id.scopeID].functionSymbols[id.indexInScope]
}