package go_llm

import (
	"context"
	"encoding/json"
	"fmt"
	"log/slog"
	"reflect"
	"time"

	"gitea.stevedudenhoeffer.com/steve/go-llm/schema"
)

type Function struct {
	Name        string      `json:"name"`
	Description string      `json:"description,omitempty"`
	Strict      bool        `json:"strict,omitempty"`
	Parameters  schema.Type `json:"parameters"`

	Forced bool `json:"forced,omitempty"`

	// Timeout is the maximum time to wait for the function to complete
	Timeout time.Duration `json:"-"`

	// fn is the function to call, only set if this is constructed with NewFunction
	fn reflect.Value

	paramType reflect.Type
}

func (f Function) WithSyntheticField(name string, description string) Function {
	if obj, o := f.Parameters.(schema.Object); o {
		f.Parameters = obj.WithSyntheticField(name, description)
	}

	return f
}

func (f Function) WithSyntheticFields(fieldsAndDescriptions map[string]string) Function {
	if obj, o := f.Parameters.(schema.Object); o {
		for k, v := range fieldsAndDescriptions {
			obj = obj.WithSyntheticField(k, v)
		}
		f.Parameters = obj
	}

	return f
}

func (f Function) Execute(ctx *Context, input string) (any, error) {
	if !f.fn.IsValid() {
		return "", fmt.Errorf("function %s is not implemented", f.Name)
	}

	slog.Info("Function.Execute", "name", f.Name, "input", input, "f", f.paramType)
	// first, we need to parse the input into the struct
	p := reflect.New(f.paramType)
	fmt.Println("Function.Execute", f.Name, "input:", input)

	var vals map[string]any
	err := json.Unmarshal([]byte(input), &vals)

	var syntheticFields map[string]string

	// first eat up any synthetic fields
	if obj, o := f.Parameters.(schema.Object); o {
		for k := range obj.SyntheticFields() {
			key := schema.SyntheticFieldPrefix + k
			if val, ok := vals[key]; ok {
				if syntheticFields == nil {
					syntheticFields = map[string]string{}
				}

				syntheticFields[k] = fmt.Sprint(val)
				delete(vals, key)
			}
		}
	}

	// now for any remaining fields, re-marshal them into json and then unmarshal into the struct
	b, err := json.Marshal(vals)
	if err != nil {
		return "", fmt.Errorf("failed to marshal input: %w (input: %s)", err, input)
	}

	// now we can unmarshal the input into the struct
	err = json.Unmarshal(b, p.Interface())
	if err != nil {
		return "", fmt.Errorf("failed to unmarshal input: %w (input: %s)", err, input)
	}

	// now we can call the function
	exec := func(ctx *Context) (any, error) {
		out := f.fn.Call([]reflect.Value{reflect.ValueOf(ctx), p.Elem()})

		if len(out) != 2 {
			return "", fmt.Errorf("function %s must return two values, got %d", f.Name, len(out))
		}

		if out[1].IsNil() {
			return out[0].Interface(), nil
		}

		return "", out[1].Interface().(error)
	}

	var cancel context.CancelFunc
	if f.Timeout > 0 {
		ctx, cancel = ctx.WithTimeout(f.Timeout)
		defer cancel()
	}

	return exec(ctx)
}

type FunctionCall struct {
	Name      string `json:"name,omitempty"`
	Arguments string `json:"arguments,omitempty"`
}

func (fc *FunctionCall) toRaw() map[string]any {
	res := map[string]interface{}{
		"name": fc.Name,
	}

	if fc.Arguments != "" {
		res["arguments"] = fc.Arguments
	}

	return res
}