package builtin

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

	"git.flytoex.net/yuanwei/flyto-agent/pkg/tools"
	"git.flytoex.net/yuanwei/flyto-agent/pkg/validator"
)

// ValidatedTool wraps a Tool with a validator.Validator gate. Approved
// verdicts pass through unchanged; Block verdicts transform the inner
// Result into an error Result carrying the Validator's identity and
// reason. Warn verdicts pass through unchanged (advisory -- the circuit
// breaker logs the sample but the write proceeds). Inner-tool errors
// (err != nil or Result.IsError) skip validation entirely -- no point
// validating something that already failed.
//
// Integration point for the circuit breaker: the VerdictSink callback
// fires on every validated call (Warn and Block) with the Validator's
// own Verdict; the breaker accumulates rejects and opens when a
// threshold is crossed (upstream TODO in separate commit).
//
// ValidatedTool 把 Tool 挂上 validator.Validator gate. approve 的
// Verdict 原样透传; Block 的 Verdict 把 inner Result 改写为 error
// Result 携带 Validator 身份 + reason. Warn 的 Verdict 原样透传
// (advisory -- 熔断器记录样本但允许 commit). Inner 工具错误 (err !=
// nil 或 Result.IsError) 完全跳过 validation -- 已经失败的操作不审.
//
// 熔断器集成点: VerdictSink 回调在每次 validated 调用 (Warn + Block)
// 以 Validator 自己的 Verdict 触发; 熔断器累积 reject, 越过阈值时 open
// (上游 TODO 另起 commit).
type ValidatedTool struct {
	inner     tools.Tool
	validator validator.Validator
	extractor ResultExtractor
	sink      VerdictSink
}

// ResultExtractor converts a successful tool Result into a DiffInput
// the Validator can consume. toolName matches the inner tool's Name()
// and is routed into DiffInput.SourceTool so the Validator can
// dispatch on it.
//
// ResultExtractor 把成功的工具 Result 转为 Validator 可消费的 DiffInput.
// toolName 与 inner tool 的 Name() 一致, 流入 DiffInput.SourceTool 供
// Validator 分发.
type ResultExtractor func(toolName string, result *tools.Result) validator.DiffInput

// VerdictSink is the observation hook the circuit breaker subscribes
// to. Every ValidatedTool call that reached the Validator fires the
// sink with the inner tool's Name() and the produced Verdict (whether
// Warn or Block, Approved or not). Passing nil disables the hook.
//
// VerdictSink 是熔断器订阅的观测 hook. 每次 ValidatedTool 到达 Validator
// 的调用以 inner tool 的 Name() 和产出的 Verdict (Warn / Block, Approved
// 或否) 触发 sink. 传 nil 关闭 hook.
type VerdictSink func(toolName string, verdict validator.Verdict)

// NewValidatedTool wraps inner with the given Validator. extractor
// converts inner's Result into a DiffInput (use DefaultExtractor if
// only Raw bytes matter, or ExtractorSQLCAS / ExtractorSQLDryRun for
// built-in SQL tools). sink may be nil; when non-nil it fires on
// every validated call for circuit-breaker observation.
//
// Panics at construction if v or extractor is nil: the explicit
// panic surfaces mis-wiring at startup instead of silently deferring
// a nil-deref to the first Execute call. Industries that deliberately
// want an unchecked Tool MUST pass validator.AlwaysApprove{} so the
// opt-out is auditable at the call site and visible in logs.
//
// NewValidatedTool 用给定 Validator 包装 inner. extractor 把 inner 的
// Result 转为 DiffInput (仅需 Raw bytes 用 DefaultExtractor, 内置 SQL
// 工具用 ExtractorSQLCAS / ExtractorSQLDryRun). sink 可为 nil; 非 nil
// 时每次 validated 调用触发供熔断器观测.
//
// 构造期对 nil v / extractor panic: 显式 panic 让错配置在启动期暴露,
// 而非把 nil-deref 静默推迟到首次 Execute. 行业 platform 刻意要 unchecked
// Tool 必须传 validator.AlwaysApprove{}, 让 opt-out 在调用点可审计,
// 日志可见.
func NewValidatedTool(inner tools.Tool, v validator.Validator, extractor ResultExtractor, sink VerdictSink) *ValidatedTool {
	if v == nil {
		panic("builtin.NewValidatedTool: validator is nil; use validator.AlwaysApprove{} for explicit opt-out")
	}
	if extractor == nil {
		panic("builtin.NewValidatedTool: extractor is nil; use DefaultExtractor() if no metadata hints are needed")
	}
	return &ValidatedTool{
		inner:     inner,
		validator: v,
		extractor: extractor,
		sink:      sink,
	}
}

// Name returns the inner tool's name. Deliberately passes through: the
// ValidatedTool replaces the bare tool at registration without
// changing how LLMs address it.
//
// Name 返回 inner 工具名. 刻意透传: ValidatedTool 在注册点替换裸工具
// 但不改变 LLM 的称呼.
func (v *ValidatedTool) Name() string { return v.inner.Name() }

// Description forwards to inner.Description (context-aware).
//
// Description 透传 inner.Description (context-aware).
func (v *ValidatedTool) Description(ctx context.Context) string {
	return v.inner.Description(ctx)
}

// InputSchema forwards to inner.InputSchema.
//
// InputSchema 透传 inner.InputSchema.
func (v *ValidatedTool) InputSchema() json.RawMessage { return v.inner.InputSchema() }

// Metadata forwards inner's MetadataProvider declaration (or the
// conservative default if inner does not implement it).
//
// Metadata 透传 inner 的 MetadataProvider (或 inner 未实现时的保守
// 默认值).
func (v *ValidatedTool) Metadata() tools.Metadata { return tools.GetMetadata(v.inner) }

// Execute runs the inner tool, then the Validator on its successful
// Result. Sink fires on every Verdict. Block verdicts (or Validator
// errors) rewrite the Result into an IsError=true Result whose Output
// carries Validator identity / reason + the original Output.
//
// Execute 运行 inner 工具, 然后在成功 Result 上跑 Validator. Sink 在
// 每次 Verdict 触发. Block verdict (或 Validator error) 会把 Result
// 改写为 IsError=true 的 Result, Output 含 Validator 身份 / reason +
// 原 Output.
func (v *ValidatedTool) Execute(ctx context.Context, input json.RawMessage, progress tools.ProgressFunc) (*tools.Result, error) {
	result, err := v.inner.Execute(ctx, input, progress)
	if err != nil {
		return result, err
	}
	if result == nil || result.IsError {
		return result, nil
	}

	diff := v.extractor(v.inner.Name(), result)
	verdict, verr := v.validator.Validate(ctx, diff)

	if v.sink != nil {
		v.sink(v.inner.Name(), verdict)
	}

	// Block: rewrite Result. Warn / Approved: pass through unchanged.
	// Validator error always escalates to Block (fail-closed) per
	// Validator interface contract.
	if verdict.Severity == validator.SeverityBlock || verr != nil {
		blocked := formatBlockedOutput(result.Output, verdict, verr)
		return &tools.Result{
			Output:   blocked,
			IsError:  true,
			Data:     result.Data,
			UndoInfo: result.UndoInfo,
		}, nil
	}

	return result, nil
}

// formatBlockedOutput composes the human-readable body placed into
// Result.Output when a Validator blocks. Reads ValidatorName,
// PolicyVersion, and Reason so replay audits can tell which policy
// blocked which diff at which commit-point.
//
// formatBlockedOutput 组装 Validator 阻断时放入 Result.Output 的人类
// 可读正文. 读 ValidatorName / PolicyVersion / Reason 让回放审计能
// 判断哪版策略在哪个 commit 点阻断了哪个 diff.
func formatBlockedOutput(orig string, v validator.Verdict, err error) string {
	var sb strings.Builder
	sb.WriteString("Validator ")
	if v.ValidatorName != "" {
		fmt.Fprintf(&sb, "%q ", v.ValidatorName)
	}
	if v.PolicyVersion != "" {
		fmt.Fprintf(&sb, "(policy %s) ", v.PolicyVersion)
	}
	sb.WriteString("blocked this operation.\n")
	if v.Reason != "" {
		sb.WriteString("Reason: ")
		sb.WriteString(v.Reason)
		sb.WriteString("\n")
	}
	if err != nil {
		sb.WriteString("Error: ")
		sb.WriteString(err.Error())
		sb.WriteString("\n")
	}
	if orig != "" {
		sb.WriteString("\n--- original tool output (for reference) ---\n")
		sb.WriteString(orig)
	}
	return sb.String()
}

// DefaultExtractor produces a minimal DiffInput: SourceTool = toolName,
// Raw = json.Marshal(result.Data), no metadata hints. Use when the
// Validator inspects Raw bytes only (LLMValidator); Rules needing
// numeric / string hints (DiffSizeRule / TableWhitelistRule) require a
// tool-specific extractor that populates Metadata.
//
// DefaultExtractor 产出最小 DiffInput: SourceTool = toolName, Raw =
// json.Marshal(result.Data), 无 metadata 提示. 适合 Validator 只检查
// Raw bytes 的场景 (LLMValidator). 需要数值 / 字符串提示的 Rule
// (DiffSizeRule / TableWhitelistRule) 要用具体工具的 extractor 填
// Metadata.
func DefaultExtractor() ResultExtractor {
	return func(toolName string, result *tools.Result) validator.DiffInput {
		raw, err := json.Marshal(result.Data)
		if err != nil {
			raw = []byte("{}")
		}
		return validator.DiffInput{
			SourceTool: toolName,
			Raw:        raw,
		}
	}
}

// ExtractorSQLCAS builds a ResultExtractor for the SQLCAS tool. It
// marshals Result.Data as Raw and populates Metadata with the
// conventional keys "affected_rows" and "table_name" (read by
// DiffSizeRule / TableWhitelistRule respectively) when present in
// Data's JSON shape.
//
// ExtractorSQLCAS 为 SQLCAS 工具构造 ResultExtractor. marshal Result.Data
// 作 Raw, 当 Data 的 JSON 形态含 "affected_rows" / "table_name" 约定
// key 时填入 Metadata (DiffSizeRule / TableWhitelistRule 会读).
func ExtractorSQLCAS() ResultExtractor {
	return extractWithMetaFields(map[string]string{
		"affected_rows": "affected_rows",
		"table_name":    "table_name",
	})
}

// ExtractorSQLDryRun builds a ResultExtractor for the SQLDryRun tool.
// DryRun's Data typically exposes "affected_row_count" (not
// "affected_rows") which this extractor normalises to the common key
// "affected_rows" so Rules keyed on the common name work uniformly
// across SQL tools. Also surfaces "after_predicate_mismatch" as a
// metadata hint for consistency checks.
//
// ExtractorSQLDryRun 为 SQLDryRun 工具构造 ResultExtractor. DryRun 的
// Data 通常暴露 "affected_row_count" (非 "affected_rows"), 此 extractor
// 归一化为通用 key "affected_rows", Rule 按通用 key 匹配即可跨 SQL 工具
// 统一. 同时把 "after_predicate_mismatch" 作为一致性检查 hint 暴露.
func ExtractorSQLDryRun() ResultExtractor {
	return extractWithMetaFields(map[string]string{
		"affected_row_count":       "affected_rows",
		"affected_rows":            "affected_rows",
		"table_name":               "table_name",
		"after_predicate_mismatch": "after_predicate_mismatch",
	})
}

// extractWithMetaFields is the shared implementation behind built-in
// extractors. keyMap maps Data JSON keys to the conventional Metadata
// keys Rules expect (many-to-one normalisation is allowed).
//
// extractWithMetaFields 是内置 extractor 的共享实现. keyMap 把 Data 的
// JSON key 映射到 Rule 期望的约定 Metadata key (允许多对一归一化).
func extractWithMetaFields(keyMap map[string]string) ResultExtractor {
	return func(toolName string, result *tools.Result) validator.DiffInput {
		raw, err := json.Marshal(result.Data)
		if err != nil {
			raw = []byte("{}")
		}
		meta := make(map[string]any)
		var parsed map[string]any
		if jerr := json.Unmarshal(raw, &parsed); jerr == nil {
			for dataKey, metaKey := range keyMap {
				if v, ok := parsed[dataKey]; ok {
					meta[metaKey] = v
				}
			}
		}
		return validator.DiffInput{
			SourceTool: toolName,
			Raw:        raw,
			Metadata:   meta,
		}
	}
}