package agent

import (
	"encoding/json"
	"fmt"
	"strings"

	"ollie/backend"
)

// parseTextToolCalls extracts tool calls from assistant text when the model
// emits them as plain text instead of using the function calling API.
// Recognizes formats:
//   - tool_name: {"key": ...}           (JSON object)
//   - tool_name: key=[...], key2=...    (shorthand: key=value pairs → JSON object)
//
// Returns nil if no valid tool calls are found.
func parseTextToolCalls(text string, tools []backend.Tool) []backend.ToolCall {
	toolNames := make(map[string]bool, len(tools))
	for _, t := range tools {
		toolNames[t.Name] = true
	}

	var calls []backend.ToolCall
	for _, name := range []string{"shell"} {
		if !toolNames[name] {
			continue
		}
		prefix := name + ":"
		idx := strings.Index(text, prefix)
		if idx < 0 {
			continue
		}
		rest := strings.TrimSpace(text[idx+len(prefix):])
		if rest == "" {
			continue
		}

		var argsJSON string
		if rest[0] == '{' {
			// Direct JSON object
			argsJSON = extractJSONObject(rest)
			if argsJSON == "" {
				argsJSON = extractJSONObject(relaxJSON(rest))
			}
		} else {
			// Shorthand format: key=[...], key2=value
			argsJSON = shorthandToJSON(rest)
		}
		if argsJSON == "" {
			continue
		}
		var check json.RawMessage
		if json.Unmarshal([]byte(argsJSON), &check) != nil {
			continue
		}
		calls = append(calls, backend.ToolCall{
			ID:        fmt.Sprintf("text-%s-%d", name, len(calls)),
			Name:      name,
			Arguments: json.RawMessage(argsJSON),
		})
	}
	return calls
}

// shorthandToJSON converts "key=[...], key2=value" format to a JSON object.
// Handles: steps=[{code: "date"}], timeout=30
func shorthandToJSON(s string) string {
	// Find the first key=value where value starts with [ or {
	eqIdx := strings.IndexByte(s, '=')
	if eqIdx < 0 {
		return ""
	}
	key := strings.TrimSpace(s[:eqIdx])
	if key == "" {
		return ""
	}
	valStr := s[eqIdx+1:]

	// Extract the value (balanced brackets/braces)
	val := extractBalanced(strings.TrimSpace(valStr))
	if val == "" {
		return ""
	}

	// Relax the value's JSON (unquoted keys → quoted)
	relaxed := relaxJSON(val)

	// Build the JSON object: {"key": relaxed_value}
	keyJSON, _ := json.Marshal(key)
	result := "{" + string(keyJSON) + ":" + relaxed + "}"

	// Check if there are more key=value pairs after the value
	afterVal := strings.TrimSpace(valStr[len(val):])
	if len(afterVal) > 0 && afterVal[0] == ',' {
		// Parse additional key=value pairs
		extra := parseExtraKV(afterVal[1:])
		if extra != "" {
			// Merge: strip trailing } from result, append extra
			result = result[:len(result)-1] + "," + extra + "}"
		}
	}

	return result
}

// extractBalanced extracts a balanced [...] or {...} from the start of s.
func extractBalanced(s string) string {
	if len(s) == 0 {
		return ""
	}
	open := s[0]
	var close byte
	switch open {
	case '[':
		close = ']'
	case '{':
		close = '}'
	default:
		// Simple value (number, string) — take until comma or newline
		end := strings.IndexAny(s, ",\n")
		if end < 0 {
			return strings.TrimSpace(s)
		}
		return strings.TrimSpace(s[:end])
	}

	depth := 0
	inStr := false
	escaped := false
	for i := 0; i < len(s); i++ {
		c := s[i]
		if escaped {
			escaped = false
			continue
		}
		if c == '\\' && inStr {
			escaped = true
			continue
		}
		if c == '"' {
			inStr = !inStr
			continue
		}
		if inStr {
			continue
		}
		if c == open {
			depth++
		} else if c == close {
			depth--
			if depth == 0 {
				return s[:i+1]
			}
		}
	}
	return ""
}

// parseExtraKV parses "key=value, key2=value2" into JSON fields (without outer braces).
func parseExtraKV(s string) string {
	s = strings.TrimSpace(s)
	if s == "" {
		return ""
	}
	var parts []string
	for s != "" {
		eqIdx := strings.IndexByte(s, '=')
		if eqIdx < 0 {
			break
		}
		key := strings.TrimSpace(s[:eqIdx])
		s = strings.TrimSpace(s[eqIdx+1:])
		val := extractBalanced(s)
		if val == "" {
			break
		}
		s = strings.TrimSpace(s[len(val):])
		if len(s) > 0 && s[0] == ',' {
			s = s[1:]
		}
		keyJSON, _ := json.Marshal(key)
		parts = append(parts, string(keyJSON)+":"+relaxJSON(val))
	}
	return strings.Join(parts, ",")
}

// extractJSONObject finds the first balanced {...} in s.
func extractJSONObject(s string) string {
	start := strings.IndexByte(s, '{')
	if start < 0 {
		return ""
	}
	depth := 0
	inStr := false
	escaped := false
	for i := start; i < len(s); i++ {
		c := s[i]
		if escaped {
			escaped = false
			continue
		}
		if c == '\\' && inStr {
			escaped = true
			continue
		}
		if c == '"' {
			inStr = !inStr
			continue
		}
		if inStr {
			continue
		}
		if c == '{' {
			depth++
		} else if c == '}' {
			depth--
			if depth == 0 {
				return s[start : i+1]
			}
		}
	}
	return ""
}

// relaxJSON converts common relaxed-JSON patterns to valid JSON:
// - unquoted keys (word: → "word":)
// - single-quoted strings → double-quoted
// This is a best-effort heuristic for model output.
func relaxJSON(s string) string {
	var out strings.Builder
	out.Grow(len(s))
	i := 0
	for i < len(s) {
		c := s[i]
		// Single-quoted string → double-quoted
		if c == '\'' {
			out.WriteByte('"')
			i++
			for i < len(s) && s[i] != '\'' {
				if s[i] == '"' {
					out.WriteString(`\"`)
				} else if s[i] == '\\' && i+1 < len(s) {
					out.WriteByte(s[i])
					i++
					out.WriteByte(s[i])
				} else {
					out.WriteByte(s[i])
				}
				i++
			}
			out.WriteByte('"')
			if i < len(s) {
				i++ // skip closing '
			}
			continue
		}
		// Unquoted key before colon: word followed by optional whitespace then ':'
		if c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c == '_' {
			j := i
			for j < len(s) && (s[j] >= 'a' && s[j] <= 'z' || s[j] >= 'A' && s[j] <= 'Z' || s[j] >= '0' && s[j] <= '9' || s[j] == '_') {
				j++
			}
			// Check if followed by optional whitespace then ':'
			k := j
			for k < len(s) && (s[k] == ' ' || s[k] == '\t') {
				k++
			}
			if k < len(s) && s[k] == ':' {
				// It's an unquoted key — quote it
				out.WriteByte('"')
				out.WriteString(s[i:j])
				out.WriteByte('"')
				i = j
				continue
			}
			// Not a key, just copy the word
			out.WriteString(s[i:j])
			i = j
			continue
		}
		out.WriteByte(c)
		i++
	}
	return out.String()
}

// transientWait returns the retry wait for a retryable error and whether it is