package tools

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

	"agentu/pkg/llm"
)

// ────────────────────────────────────────────────────────────────────────────
// ApplyPatchTool — multi-file patch using V4 format
// ────────────────────────────────────────────────────────────────────────────

type ApplyPatchTool struct {
	workingDir string
}

func NewApplyPatchTool(workingDir string) *ApplyPatchTool {
	return &ApplyPatchTool{workingDir: workingDir}
}

// MutatingTool interface.

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

// CallMutatingTool — dry_run is non-mutating.
func (t *ApplyPatchTool) MutatesCall(args json.RawMessage) bool {
	var a applyPatchArgs
	if err := json.Unmarshal(args, &a); err != nil {
		return true // conservative
	}
	return !a.DryRun
}

// PreviewableTool interface.

func (t *ApplyPatchTool) Preview(ctx context.Context, args json.RawMessage) (string, error) {
	var a applyPatchArgs
	if err := json.Unmarshal(args, &a); err != nil {
		return "", fmt.Errorf("parse arguments: %w", err)
	}
	return summarizePatch(a.Patch)
}

// TargetPathTool interface.

func (t *ApplyPatchTool) TargetPaths(args json.RawMessage) ([]string, error) {
	var a applyPatchArgs
	if err := json.Unmarshal(args, &a); err != nil {
		return nil, fmt.Errorf("parse arguments: %w", err)
	}
	return extractTargetPaths(t.workingDir, a.Patch)
}

// Definition.

func (t *ApplyPatchTool) Definition() llm.Tool {
	return llm.Tool{
		Type: "function",
		Function: llm.ToolFunction{
			Name:        "apply_patch",
			Description: "Apply a multi-file patch to create, update, or delete files. Uses a V4 diff format with *** Begin Patch / *** Add File / *** Update File / *** Delete File markers. Prefer this over multiple file_edit calls when changing 2+ files.",
			Parameters: JSONSchema(`{
  "type": "object",
  "properties": {
    "patch": {"type": "string", "description": "V4 format patch text. Must start with *** Begin Patch and end with *** End Patch."},
    "dry_run": {"type": "boolean", "description": "Preview the patch without writing any files."}
  },
  "required": ["patch"],
  "additionalProperties": false
}`),
		},
	}
}

type applyPatchArgs struct {
	Patch  string `json:"patch"`
	DryRun bool   `json:"dry_run"`
}

// Execute.

func (t *ApplyPatchTool) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
	var args applyPatchArgs
	if err := json.Unmarshal(raw, &args); err != nil {
		return "", fmt.Errorf("parse arguments: %w", err)
	}

	if strings.TrimSpace(args.Patch) == "" {
		return "", fmt.Errorf("patch is empty")
	}

	pp, err := parsePatch(args.Patch)
	if err != nil {
		return "", err
	}

	if args.DryRun {
		summary, err := summarizePatch(args.Patch)
		if err != nil {
			return "", err
		}
		return "dry run — " + summary, nil
	}

	rollback := NewRollbackStore()
	var snapshotIDs []string
	for _, pf := range pp.files {
		path, err := resolvePathForWrite(t.workingDir, pf.path)
		if err != nil {
			return "", err
		}
		id, err := rollback.Capture("apply_patch", pf.path, path)
		if err != nil {
			return "", err
		}
		snapshotIDs = append(snapshotIDs, id)
	}

	var results []string
	for _, pf := range pp.files {
		if err := ctx.Err(); err != nil {
			if rollbackErr := restoreSnapshots(rollback, snapshotIDs); rollbackErr != nil {
				return "", fmt.Errorf("%w; rollback failed: %v", err, rollbackErr)
			}
			return "", err
		}
		result, err := applyPatchFile(t.workingDir, pf)
		if err != nil {
			if rollbackErr := restoreSnapshots(rollback, snapshotIDs); rollbackErr != nil {
				return "", fmt.Errorf("apply patch: %w; rollback failed: %v", err, rollbackErr)
			}
			return "", fmt.Errorf("apply patch: %w", err)
		}
		results = append(results, result)
	}

	summary := fmt.Sprintf("applied patch: %d file(s) changed\n%s", len(results), strings.Join(results, "\n"))
	return summary, nil
}

func restoreSnapshots(store *RollbackStore, ids []string) error {
	var errs []error
	for i := len(ids) - 1; i >= 0; i-- {
		if _, err := store.Restore(ids[i]); err != nil {
			errs = append(errs, err)
		}
	}
	return errors.Join(errs...)
}