package macro

import (
	"fmt"

	"github.com/hellerve/argos/ast"
)

type Macro struct {
	params      []string
	transformer *ast.AST
}

type MacroEnv map[string]Macro

func NewMacroEnv() MacroEnv {
	return make(MacroEnv)
}

func handleMac(l []*ast.AST, m MacroEnv) (*ast.AST, error) {
	if len(l) != 4 {
		return nil, fmt.Errorf("mac takes 3 arguments")
	}

	nameThunk := l[1]

	if nameThunk.Tag != ast.Symbol {
		return nil, fmt.Errorf("First argument to mac must be symbol, got %s", nameThunk.Pretty())
	}

	name := nameThunk.Val.(string)

	paramsThunk := l[2]

	if paramsThunk.Tag != ast.List {
		return nil, fmt.Errorf("second argument to mac must be list, got %s", paramsThunk.Pretty())
	}

	var params []string

	for _, elem := range paramsThunk.Val.([]*ast.AST) {
		if elem.Tag != ast.Symbol {
			return nil, fmt.Errorf("Invalid value in parameter list to mac: %s", elem.Pretty())
		}
		params = append(params, elem.Val.(string))
	}

	m[name] = Macro{params, l[3]}

	nilVal := ast.AST{ast.Symbol, "nil"}
	return &nilVal, nil
}

func transformList(l []*ast.AST, args map[string]*ast.AST) (*ast.AST, error) {
	var res []*ast.AST

	for _, elem := range l {
		transformed, err := doTransform(elem, args)

		if err != nil {
			return nil, err
		}

		res = append(res, transformed)
	}

	lst := ast.AST{ast.List, res}

	return &lst, nil
}

func doTransform(transformer *ast.AST, args map[string]*ast.AST) (*ast.AST, error) {
	switch transformer.Tag {
	case ast.List:
		lst := transformer.Val.([]*ast.AST)
		if len(lst) > 0 && lst[0].Tag == ast.Symbol && lst[0].Val.(string) == "list" {
			return transformList(lst[1:], args)
		}
		if len(lst) > 0 && lst[0].Tag == ast.Symbol && lst[0].Val.(string) == "quote" {
			return lst[1], nil
		}
		return transformer, nil
	case ast.Quoted:
		return transformer.Val.(*ast.AST), nil
	case ast.Symbol:
		key := transformer.Val.(string)
		val, ok := args[key]
		if !ok {
			return nil, fmt.Errorf("Unbound symbol %s", key)
		}
		return val, nil
	}
	return transformer, nil
}

func transform(l *ast.AST, macro Macro, m MacroEnv) (*ast.AST, error) {
	args := make(map[string]*ast.AST)
	vals := l.Val.([]*ast.AST)[1:]

	if len(vals) != len(macro.params) {
		return nil, fmt.Errorf("Wrong number of arguments to macro")
	}

	for i, arg := range vals {
		expanded, err := Expand(arg, m)
		if err != nil {
			return nil, err
		}
		args[macro.params[i]] = expanded
	}

	return doTransform(macro.transformer, args)
}

func traverse(l *ast.AST, m MacroEnv) (*ast.AST, error) {
	vals := l.Val.([]*ast.AST)
	if len(vals) == 0 {
		return l, nil
	}

	if vals[0].Tag == ast.Symbol && vals[0].Val.(string) == "mac" {
		return handleMac(vals, m)
	}

	if vals[0].Tag == ast.Symbol {
		sym := vals[0].Val.(string)

		macro, ok := m[sym]

		if ok {
			return transform(l, macro, m)
		}
	}

	return l, nil
}

func Expand(l *ast.AST, m MacroEnv) (*ast.AST, error) {
	if l.Tag != ast.List {
		return l, nil
	}
	return traverse(l, m)
}