package scope

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

type ScopeChain struct {
	symbolScopes   []SymbolScope // All other scopes are bound to this by ID.
	functionScopes []FunctionScope
	loopScopes     []LoopScope
}

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

	loopScopes := make([]LoopScope, 0)

	return ScopeChain{
		symbolScopes:   symbolScopes,
		functionScopes: functionScopes,
		loopScopes:     loopScopes,
	}
}

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) CurrentLoop() *LoopScope {
	if len(sc.loopScopes) == 0 {
		return nil
	}

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

func (sc *ScopeChain) CalculateTopLocalIndex() int {
	topFunctionScope := sc.CurrentFunction()

	sumLocals := 0
	for scopeIndex := len(sc.symbolScopes) - 1; scopeIndex >= 0; scopeIndex-- {
		scope := sc.symbolScopes[scopeIndex]
		sumLocals += len(scope.variableSymbols)

		// Function scope signifies start of stack boundry.
		if ScopeID(scopeIndex) == topFunctionScope.id {
			break
		}
	}

	return sumLocals
}

func (sc *ScopeChain) CanSymbolBeDeclared(name string) bool {
	topFunctionScope := sc.CurrentFunction()

	for scopeIndex := len(sc.symbolScopes) - 1; scopeIndex >= 0; scopeIndex-- {
		scope := sc.symbolScopes[scopeIndex]

		if _, ok := scope.nameToSymbol[name]; ok {
			return false
		}

		// Shadowing only allowed with at least one function scope in the chain.
		if ScopeID(scopeIndex) == topFunctionScope.id {
			return true
		}
	}

	return true
}

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) EnterConstructor(unit code.Marker, thisLocal int) {
	sc.Enter()
	sc.functionScopes = append(sc.functionScopes, NewMethodFunctionScope(sc.CurrentScopeID(), unit, thisLocal))
}

func (sc *ScopeChain) EnterMethod(unit code.Marker) {
	sc.Enter()
	sc.functionScopes = append(sc.functionScopes, NewMethodFunctionScope(sc.CurrentScopeID(), unit, -1))
}

func (sc *ScopeChain) EnterLoop() (code.Marker, code.Marker) {
	parentMarker := sc.CreateAnonymousFunctionSubUnit()

	breakMarker := parentMarker.SubMarker("end")
	continueMarker := parentMarker.SubMarker("start")

	sc.Enter()
	sc.loopScopes = append(sc.loopScopes, NewLoopScope(sc.CurrentScopeID(), breakMarker, continueMarker))

	return breakMarker, continueMarker
}

func (sc *ScopeChain) Exit() int {
	if sc.CurrentScopeID() == 0 {
		return 0
	}
	id := sc.CurrentScopeID()
	stackSpace := len(sc.symbolScopes[id].variableSymbols)

	sc.symbolScopes = sc.symbolScopes[:id]

	if sc.CurrentLoop() != nil && sc.CurrentLoop().id == id {
		sc.loopScopes = sc.loopScopes[:len(sc.loopScopes)-1]
	}

	if sc.CurrentFunction().id == id {
		sc.functionScopes = sc.functionScopes[:len(sc.functionScopes)-1]
	}

	return stackSpace
}

func (sc *ScopeChain) Declare(name string) (int, bool) {
	// Check whether the symbol is already declared in any of the scopes.
	if !sc.CanSymbolBeDeclared(name) {
		return 0, false
	}

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

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

	localIndex := sc.CalculateTopLocalIndex()

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

	current.nameToSymbol[name] = symbolID

	return localIndex, true
}

func (sc *ScopeChain) DeclareAnonymous() int {
	current := sc.Current()
	localIndex := sc.CalculateTopLocalIndex()

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

	return localIndex
}

func (sc *ScopeChain) DeclareTemporary() int {
	return sc.CalculateTopLocalIndex() // :)
}

func (sc *ScopeChain) DeclareGlobal(name, module, author string) bool {
	if !sc.CanSymbolBeDeclared(name) {
		return false
	}

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

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

	globalID := fmt.Sprintf("%s:%s:%s", author, module, name)

	// Declare the symbol in the current scope.
	current.globalSymbols = append(current.globalSymbols, Symbol[SymbolGlobal]{
		name: name,
		data: SymbolGlobal{
			id: globalID,
		},
	})

	current.nameToSymbol[name] = symbolID

	return true
}

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) {
	var id SymbolID
	for scopeIndex := len(sc.symbolScopes) - 1; scopeIndex >= 0; scopeIndex-- {
		scope := sc.symbolScopes[scopeIndex]

		if foundID, ok := scope.nameToSymbol[name]; ok {
			id = foundID
			break
		}

		if scopeIndex == 0 {
			return SymbolID{}, false
		}
	}

	// Check whether the symbol is outside the current function scope.
	fnScope := sc.CurrentFunction()
	if id.scopeID < fnScope.id {
		// Return env symbol instead of a local symbol.
		return SymbolID{
			symbolKind:   SymbolKindEnv,
			scopeID:      id.scopeID,
			indexInScope: id.indexInScope,
		}, true
	}

	return id, true
}

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) GetEnv(id SymbolID) Symbol[SymbolEnv] {
	if id.symbolKind != SymbolKindEnv {
		panic("incorrect symbol id kind given")
	}

	symbol := sc.symbolScopes[id.scopeID].variableSymbols[id.indexInScope]

	// Add the local to the function scope, if it is not already there.
	fnScope := sc.CurrentFunction()

	var indexInEnv int
	alreadyUsed := false
	for i, env := range fnScope.outsideSymbolsInEnv {
		if env == id {
			alreadyUsed = true
			indexInEnv = i
			break
		}
	}

	if !alreadyUsed {
		indexInEnv = len(fnScope.outsideSymbolsInEnv)
		fnScope.outsideSymbolsInEnv = append(fnScope.outsideSymbolsInEnv, id)
	}

	return Symbol[SymbolEnv]{
		name: symbol.name,
		data: SymbolEnv{
			indexInEnv: indexInEnv,
		},
	}
}

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

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