package interpreter

import (
	"fmt"

	"usebox.net/micro-lang/ast"
	"usebox.net/micro-lang/errors"
	"usebox.net/micro-lang/tokens"
)

type Callable interface {
	Name() string
	String() string
	Call(interp *Interpreter, args []any, loc tokens.Location) (any, error)
	Type() *ast.Type
	Params() []ast.Type
	Ret() *ast.Type
}

var builtInFuncs = map[string]Callable{
	"len":     builtInLen{},
	"println": builtInPrintln{},
	"panic":   builtInPanic{},
}

func BuiltInTypes() map[string]ast.Type {
	types := map[string]ast.Type{}
	for k, v := range builtInFuncs {
		types[k] = *v.Type()
	}
	return types
}

type builtInLen struct{}

func (n builtInLen) Name() string {
	return "'len'"
}

func (n builtInLen) Type() *ast.Type {
	return ast.NewFuncType(tokens.Token{}, n.Params(), n.Ret())
}

func (n builtInLen) String() string {
	return n.Type().String()
}

func (n builtInLen) Params() []ast.Type {
	// won't be arity or type checked
	return []ast.Type{ast.TypeArray}
}

func (n builtInLen) Ret() *ast.Type {
	return &ast.TypeNumber
}

func (n builtInLen) Call(interp *Interpreter, args []any, loc tokens.Location) (any, error) {
	vals, ok := args[0].([]any)
	if !ok {
		// shouldn't happen
		return nil, errors.NewError(errors.Unexpected, loc, "type mismatch in call to 'len'")
	}
	// return
	panic(&PanicJump{typ: PanicJumpReturn, value: ast.Number(len(vals))})
}

type builtInPrintln struct{}

func (n builtInPrintln) Name() string {
	return "'println'"
}

func (n builtInPrintln) Type() *ast.Type {
	return ast.NewFuncType(tokens.Token{}, n.Params(), n.Ret())
}

func (n builtInPrintln) String() string {
	return n.Type().String()
}

func (n builtInPrintln) Params() []ast.Type {
	// won't be arity or type checked
	return nil
}

func (n builtInPrintln) Ret() *ast.Type {
	return &ast.TypeNumber
}

func (n builtInPrintln) Call(interp *Interpreter, args []any, loc tokens.Location) (any, error) {
	var count int

	for i := range args {
		if args[i] == nil {
			continue
		}
		written, err := fmt.Print(args[i])
		if err != nil {
			return nil, err
		}
		count += written
	}
	fmt.Println()
	count++

	// return
	panic(&PanicJump{typ: PanicJumpReturn, value: ast.Number(count)})
}

type builtInPanic struct{}

func (n builtInPanic) Name() string {
	return "'panic'"
}

func (n builtInPanic) Type() *ast.Type {
	return ast.NewFuncType(tokens.Token{}, n.Params(), n.Ret())
}

func (n builtInPanic) String() string {
	return n.Type().String()
}

func (n builtInPanic) Params() []ast.Type {
	return []ast.Type{ast.TypeString}
}

func (n builtInPanic) Ret() *ast.Type {
	// no return (returns none)
	return nil
}

func (n builtInPanic) Call(interp *Interpreter, args []any, loc tokens.Location) (any, error) {
	return nil, fmt.Errorf("[%s] panic: %s", loc, args[0])
}

type Function struct {
	fun     ast.Func
	closure *Env
}

func (f Function) Name() string {
	return f.fun.Name.String()
}

func (f Function) Type() *ast.Type {
	return ast.NewFuncType(f.fun.Name, f.Params(), f.Ret())
}

func (f Function) String() string {
	return f.Type().String()
}

func (f Function) Params() []ast.Type {
	params := make([]ast.Type, 0, 1)
	for _, p := range f.fun.Params {
		params = append(params, p.Type)
	}
	return params
}

func (f Function) Ret() *ast.Type {
	return f.fun.Ret
}

func (f Function) Call(interp *Interpreter, args []any, loc tokens.Location) (result any, err error) {
	pEnv := interp.env
	interp.env = NewEnv(f.closure)
	defer func() {
		interp.env = pEnv
	}()

	for n, v := range f.fun.Params {
		interp.env = interp.env.Set(v.Name.Value, Var{Value: args[n], Loc: v.Name.Loc})
	}

	// tail call optimization
	for {
		// wrap the evaluation in a function
		var tailCall *PanicJump
		tailCall, result, err = func() (tailCall *PanicJump, result any, err error) {
			// handle tail call
			// will call this function again without setting up a new call frame
			defer func() {
				if r := recover(); r != nil {
					if val, ok := r.(*PanicJump); ok && val.typ == PanicJumpTailCall {
						tailCall = val
						return
					}
					panic(r)
				}
			}()
			result, err = interp.evaluate(f.fun.Body)
			return nil, result, err
		}()
		if tailCall == nil {
			break
		}

		// XXX: can we optimize this?
		// if the callee can be a variable expression, we probably can't
		call := tailCall.value.(ast.Call)
		callee, err := interp.evaluate(call.Callee)
		if err != nil {
			return nil, err
		}
		if fun, ok := callee.(Callable); !ok || fun.Name() != f.Name() {
			// can't be optimized
			val, err := interp.evaluate(call)
			if err != nil {
				return nil, err
			}
			panic(&PanicJump{typ: PanicJumpReturn, value: val, loc: tailCall.loc})
		}

		args, err := interp.evalArgs(f.Name(), call.Loc, f.Params(), call.Args)
		if err != nil {
			return nil, err
		}

		for n, v := range f.fun.Params {
			interp.env.Update(v.Name.Value, args[n])
		}
	}

	return result, err
}