package tools

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"strings"

	"agentu/pkg/llm"
)

type FileEditTool struct {
	workingDir string
}

func NewFileEditTool(workingDir string) *FileEditTool {
	return &FileEditTool{workingDir: workingDir}
}

func (*FileEditTool) Mutates() bool { return true }

func (t *FileEditTool) Definition() llm.Tool {
	return llm.Tool{
		Type: "function",
		Function: llm.ToolFunction{
			Name:        "file_edit",
			Description: "Edit an existing UTF-8 text file by replacing an inclusive 1-based line range or an exact old text pattern.",
			Parameters: JSONSchema(`{
  "type": "object",
  "properties": {
    "path": {"type": "string", "description": "Existing file path relative to the configured agent working directory."},
    "start_line": {"type": "integer", "description": "1-based first line for a line-range edit."},
    "end_line": {"type": "integer", "description": "Optional 1-based final line; defaults to start_line when omitted."},
    "old": {"type": "string", "description": "Exact text to replace for a pattern edit; not regex."},
    "new": {"type": "string", "description": "Replacement text; may be empty string for deletion."},
    "replace_all": {"type": "boolean", "description": "When true, replace every exact old occurrence; otherwise require exactly one occurrence."}
  },
  "required": ["path", "new"],
  "additionalProperties": false
}`),
		},
	}
}

type fileEditArgs struct {
	Path       string  `json:"path"`
	StartLine  int     `json:"start_line"`
	EndLine    int     `json:"end_line"`
	Old        string  `json:"old"`
	New        *string `json:"new"`
	ReplaceAll bool    `json:"replace_all"`
}

func (t *FileEditTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
	var args fileEditArgs
	if err := json.Unmarshal(raw, &args); err != nil {
		return "", fmt.Errorf("parse arguments: %w", err)
	}
	if strings.TrimSpace(args.Path) == "" {
		return "", errors.New("path is required")
	}
	if args.New == nil {
		return "", errors.New("new is required")
	}
	path, err := resolvePath(t.workingDir, args.Path)
	if err != nil {
		return "", err
	}

	info, err := os.Stat(path)
	if err != nil {
		return "", fmt.Errorf("cannot access %s: %w", args.Path, err)
	}
	if info.IsDir() {
		return "", fmt.Errorf("%s is a directory, not a file", args.Path)
	}
	perm := info.Mode().Perm()

	hasLine := args.StartLine > 0 || args.EndLine > 0
	hasOld := strings.TrimSpace(args.Old) != ""

	if hasLine && hasOld {
		return "", errors.New("provide either start_line/end_line or old, not both")
	}
	if !hasLine && !hasOld {
		return "", errors.New("provide either start_line/end_line or old")
	}

	if err := ctx.Err(); err != nil {
		return "", err
	}

	data, err := os.ReadFile(path)
	if err != nil {
		return "", err
	}
	content := string(data)

	var result string
	if hasOld {
		result, err = editByPattern(content, args.Path, args.Old, *args.New, args.ReplaceAll)
	} else {
		result, err = editByLineRange(content, args.Path, args.StartLine, args.EndLine, *args.New)
	}
	if err != nil {
		return "", err
	}

	if err := ctx.Err(); err != nil {
		return "", err
	}
	if err := os.WriteFile(path, []byte(result), perm); err != nil {
		return "", err
	}

	switch {
	case hasOld:
		count := strings.Count(content, args.Old)
		if args.ReplaceAll {
			return fmt.Sprintf("edited %s: replaced %d occurrences", args.Path, count), nil
		}
		return fmt.Sprintf("edited %s: replaced 1 occurrence", args.Path), nil
	default:
		endLine := args.EndLine
		if endLine == 0 {
			endLine = args.StartLine
		}
		return fmt.Sprintf("edited %s: replaced lines %d-%d", args.Path, args.StartLine, endLine), nil
	}
}

func editByLineRange(content, displayPath string, startLine, endLine int, replacement string) (string, error) {
	if startLine < 1 {
		return "", errors.New("start_line must be >= 1")
	}
	if endLine == 0 {
		endLine = startLine
	}
	if endLine < startLine {
		return "", fmt.Errorf("end_line (%d) must be >= start_line (%d)", endLine, startLine)
	}

	lines := strings.Split(content, "\n")
	if startLine > len(lines) {
		return "", fmt.Errorf("line range %d-%d is outside file with %d lines", startLine, endLine, len(lines))
	}
	if endLine > len(lines) {
		return "", fmt.Errorf("line range %d-%d is outside file with %d lines", startLine, endLine, len(lines))
	}

	var b strings.Builder
	b.WriteString(strings.Join(lines[:startLine-1], "\n"))
	if startLine > 1 {
		b.WriteString("\n")
	}
	b.WriteString(replacement)
	if endLine < len(lines) {
		if !strings.HasSuffix(replacement, "\n") && replacement != "" {
			b.WriteString("\n")
		}
		b.WriteString(strings.Join(lines[endLine:], "\n"))
	}

	return b.String(), nil
}

func editByPattern(content, displayPath, old, replacement string, replaceAll bool) (string, error) {
	count := strings.Count(content, old)
	if count == 0 {
		return "", errors.New("old text not found")
	}
	if !replaceAll && count != 1 {
		return "", fmt.Errorf("old text matched %d occurrences; set replace_all=true or provide a more specific old value", count)
	}
	return strings.ReplaceAll(content, old, replacement), nil
}