package compiler

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

type ScopeID int

type ScopeChain struct {
	nameToSymbol map[string]SymbolID
	scopes       []Scope
}

func NewScopeChain() ScopeChain {
	scopes := make([]Scope, 1)
	scopes[0] = NewFunctionScope("") // Top-most scope is a function scope, so it can have sub-units

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

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

func (sc *ScopeChain) Current() *Scope {
	return &sc.scopes[sc.CurrentScopeID()]
}

func (sc *ScopeChain) CurrentFunction() *Scope {
	// TODO: Probably should make this lookup constant by making a seperate array of function scopes
	for i := len(sc.scopes) - 1; i <= 0; i++ {
		if sc.scopes[i].kind == ScopeKindFunction {
			return &sc.scopes[i]
		}
	}

	panic("top scope should always be a function scope")
}

func (sc *ScopeChain) Enter() {
	sc.scopes = append(sc.scopes, NewNormalScope())
}

func (sc *ScopeChain) EnterFunction(unitName string) {
	sc.scopes = append(sc.scopes, NewFunctionScope(unitName))
}

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

	sc.scopes[sc.CurrentScopeID()] = Scope{}
	sc.scopes = sc.scopes[: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) (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,
		},
	})

	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()
	data := fnScope.data.(ScopeFunction)

	index := data.subUnitCount
	data.subUnitCount++

	fnScope.data = data

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

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

	name := data.unitName
	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.scopes[id.scopeID].variableSymbols[id.indexInScope]
}

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

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

type SymbolID struct {
	symbolKind   SymbolKind
	scopeID      ScopeID
	indexInScope int
}

type ScopeKind int

const (
	ScopeKindNormal ScopeKind = iota
	ScopeKindFunction
	ScopeKindLoop
)

type Scope struct {
	variableSymbols []Symbol[SymbolVariable]
	functionSymbols []Symbol[SymbolFunction]

	kind ScopeKind
	data any
}

func NewNormalScope() Scope {
	return Scope{
		variableSymbols: make([]Symbol[SymbolVariable], 0),
		functionSymbols: make([]Symbol[SymbolFunction], 0),
		kind:            ScopeKindNormal,
		data:            ScopeNormal{},
	}
}

func NewFunctionScope(unitName string) Scope {
	return Scope{
		variableSymbols: make([]Symbol[SymbolVariable], 0),
		functionSymbols: make([]Symbol[SymbolFunction], 0),
		kind:            ScopeKindFunction,
		data: ScopeFunction{
			unitName:     code.Marker(unitName),
			subUnitCount: 0,
		},
	}
}

type ScopeNormal struct{}

type ScopeFunction struct {
	unitName     code.Marker
	subUnitCount int
}

type ScopeLoop struct {
	breakMarker    code.Marker
	continueMarker code.Marker
}