package basic

import (
	"context"
	"fmt"
	"math"
	"strconv"
	"strings"
)

// maxExprLen bounds an expression length so a request cannot force a parse/scan
// beyond a trivial cost and cannot smuggle a larger payload through the tool.
const maxExprLen = 256

// allowedChars lists the bytes permitted in an expression besides ASCII letters
// (letters are only admitted to spell the sqrt function; the lexer rejects any
// other identifier).
const allowedChars = "0123456789.+-*/^() \t\n"

// calculator evaluates a safe arithmetic expression. It recognizes signed numeric
// literals, + - * /, parentheses, the right-associative ^ power operator and the
// sqrt(x) function. Every other character or identifier is rejected, so the
// expression can never run arbitrary code. Division by zero and non-real results
// are errors.
type calculator struct{}

// Name returns the tool identifier.
func (calculator) Name() string { return "calculator" }

// Description explains what the tool does for the model.
func (calculator) Description() string {
	return "Safely evaluate an arithmetic expression using + - * /, parentheses, ^ (power) and sqrt(x)."
}

// Parameters returns the tool input JSON Schema.
func (calculator) Parameters() (map[string]any, error) {
	return map[string]any{
		"type": "object",
		"properties": map[string]any{
			"expression": map[string]any{"type": "string", "description": "the arithmetic expression to evaluate"},
		},
		"required": []any{"expression"},
	}, nil
}

// Execute evaluates expression and returns {"expression":..., "result":...}.
func (calculator) Execute(ctx context.Context, args any) (any, error) {
	if err := ctx.Err(); err != nil {
		return nil, err
	}
	m, err := asMap(args)
	if err != nil {
		return nil, err
	}
	expr, err := reqString(m, "expression")
	if err != nil {
		return nil, err
	}
	val, err := eval(expr)
	if err != nil {
		log().Error("calculator failed", "expression", expr, "reason", err.Error())
		return nil, err
	}
	return map[string]any{"expression": expr, "result": val}, nil
}

// eval validates, tokenizes and parses expr, returning its numeric value.
func eval(expr string) (float64, error) {
	if expr == "" {
		return 0, fmt.Errorf("%w: expression is empty", ErrEvaluation)
	}
	if len(expr) > maxExprLen {
		return 0, fmt.Errorf("%w: expression exceeds %d characters", ErrEvaluation, maxExprLen)
	}
	for _, c := range expr {
		if strings.ContainsRune(allowedChars, c) {
			continue
		}
		if c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' {
			continue // letter: the lexer admits only "sqrt"
		}
		return 0, fmt.Errorf("%w: disallowed character %q", ErrEvaluation, c)
	}
	toks, err := lex(expr)
	if err != nil {
		return 0, err
	}
	p := &parser{toks: toks}
	return p.parse()
}

func isAsciiLetter(c byte) bool {
	return c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z'
}

// tokKind describes a token category produced by the lexer.
type tokKind int

const (
	tokEOF tokKind = iota
	tokNum
	tokOp
	tokLParen
	tokRParen
	tokIdent
)

type token struct {
	kind tokKind
	text string
	val  float64
}

// lex scans s into a token stream terminated by tokEOF. Any ASCII letter run that
// is not exactly "sqrt" is rejected.
func lex(s string) ([]token, error) {
	toks := make([]token, 0, len(s)/2+1)
	for i := 0; i < len(s); {
		c := s[i]
		switch {
		case c == ' ' || c == '\t' || c == '\n':
			i++
		case c == '(':
			toks = append(toks, token{kind: tokLParen, text: "("})
			i++
		case c == ')':
			toks = append(toks, token{kind: tokRParen, text: ")"})
			i++
		case c == '+' || c == '-' || c == '*' || c == '/' || c == '^':
			toks = append(toks, token{kind: tokOp, text: string(c)})
			i++
		case isAsciiLetter(c):
			start := i
			for i < len(s) && isAsciiLetter(s[i]) {
				i++
			}
			name := s[start:i]
			if name != "sqrt" {
				return nil, fmt.Errorf("%w: unexpected identifier %q", ErrEvaluation, name)
			}
			toks = append(toks, token{kind: tokIdent, text: name})
		case c >= '0' && c <= '9' || c == '.':
			start := i
			dot := false
			for i < len(s) {
				d := s[i]
				if d >= '0' && d <= '9' {
					i++
					continue
				}
				if d == '.' && !dot {
					dot = true
					i++
					continue
				}
				break
			}
			raw := s[start:i]
			val, err := strconv.ParseFloat(raw, 64)
			if err != nil {
				return nil, fmt.Errorf("%w: invalid number %q", ErrEvaluation, raw)
			}
			toks = append(toks, token{kind: tokNum, text: raw, val: val})
		default:
			return nil, fmt.Errorf("%w: disallowed character %q", ErrEvaluation, c)
		}
	}
	toks = append(toks, token{kind: tokEOF, text: ""})
	return toks, nil
}

// parser is a recursive-descent parser over the token stream. Grammar:
//
//	expr   := term (('+'|'-') term)*
//	term   := unary (('*'|'/') unary)*
//	unary  := ('+'|'-') unary | power
//	power  := primary ('^' unary)?
//
// ^ binds tighter than unary minus (so -2^2 == -(2^2)) and is right-associative.
type parser struct {
	toks []token
	pos  int
}

func (p *parser) cur() token { return p.toks[p.pos] }

func (p *parser) parse() (float64, error) {
	v, err := p.expr()
	if err != nil {
		return 0, err
	}
	if p.cur().kind != tokEOF {
		return 0, fmt.Errorf("%w: unexpected token %q", ErrEvaluation, p.cur().text)
	}
	return v, nil
}

func (p *parser) expr() (float64, error) {
	v, err := p.term()
	if err != nil {
		return 0, err
	}
	for p.cur().kind == tokOp && (p.cur().text == "+" || p.cur().text == "-") {
		op := p.cur().text
		p.pos++
		r, err := p.term()
		if err != nil {
			return 0, err
		}
		if op == "+" {
			v += r
		} else {
			v -= r
		}
	}
	return v, nil
}

func (p *parser) term() (float64, error) {
	v, err := p.unary()
	if err != nil {
		return 0, err
	}
	for p.cur().kind == tokOp && (p.cur().text == "*" || p.cur().text == "/") {
		op := p.cur().text
		p.pos++
		r, err := p.unary()
		if err != nil {
			return 0, err
		}
		if op == "*" {
			v *= r
		} else {
			if r == 0 {
				return 0, fmt.Errorf("%w: division by zero", ErrEvaluation)
			}
			v /= r
		}
	}
	return v, nil
}

func (p *parser) unary() (float64, error) {
	if p.cur().kind == tokOp && (p.cur().text == "+" || p.cur().text == "-") {
		neg := p.cur().text == "-"
		p.pos++
		v, err := p.unary()
		if err != nil {
			return 0, err
		}
		if neg {
			return -v, nil
		}
		return v, nil
	}
	return p.power()
}

func (p *parser) power() (float64, error) {
	base, err := p.primary()
	if err != nil {
		return 0, err
	}
	if p.cur().kind == tokOp && p.cur().text == "^" {
		p.pos++
		exp, err := p.unary()
		if err != nil {
			return 0, err
		}
		r := math.Pow(base, exp)
		if math.IsNaN(r) {
			return 0, fmt.Errorf("%w: non-real result", ErrEvaluation)
		}
		return r, nil
	}
	return base, nil
}

func (p *parser) primary() (float64, error) {
	t := p.cur()
	switch t.kind {
	case tokNum:
		p.pos++
		return t.val, nil
	case tokLParen:
		p.pos++
		v, err := p.expr()
		if err != nil {
			return 0, err
		}
		if p.cur().kind != tokRParen {
			return 0, fmt.Errorf("%w: missing closing parenthesis", ErrEvaluation)
		}
		p.pos++
		return v, nil
	case tokIdent:
		p.pos++
		if p.cur().kind != tokLParen {
			return 0, fmt.Errorf("%w: sqrt requires parentheses", ErrEvaluation)
		}
		p.pos++
		v, err := p.expr()
		if err != nil {
			return 0, err
		}
		if p.cur().kind != tokRParen {
			return 0, fmt.Errorf("%w: missing closing parenthesis in sqrt", ErrEvaluation)
		}
		p.pos++
		if v < 0 {
			return 0, fmt.Errorf("%w: sqrt of a negative number", ErrEvaluation)
		}
		return math.Sqrt(v), nil
	default:
		return 0, fmt.Errorf("%w: unexpected token %q", ErrEvaluation, t.text)
	}
}