// Package tools defines the tool interface the agent loop dispatches against.
//
// Tools return (string, error). The error is text for the model to read, not an
// exception for Codey to swallow: a failed edit is a normal event in the
// transcript.
package tools

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"sync"
	"time"

	"github.com/johannSo/codey/internal/llm"
	"github.com/johannSo/codey/internal/metrics"

	"github.com/invopop/jsonschema"
)

// Tool is everything the loop needs to advertise and run one tool.
type Tool interface {
	Name() string
	Description() string
	InputSchema() json.RawMessage
	Execute(ctx context.Context, raw json.RawMessage) (string, error)
}

// Output is what a tool call produced. Text is always set: it is what the
// transcript shows, the log measures and — for most tools — what the model
// reads. Blocks is set only when text cannot carry the result, an image above
// all; then the model is sent the blocks and Text is their caption.
type Output struct {
	Text   string
	Blocks []llm.ContentBlock
}

// BlockTool is a Tool whose result may be content blocks rather than a string.
// The registry prefers ExecuteBlocks when a tool offers it; everything else
// keeps returning text and never knows the difference.
type BlockTool interface {
	Tool
	ExecuteBlocks(ctx context.Context, raw json.RawMessage) (Output, error)
}

// typed adapts a func taking a concrete argument struct to the Tool interface,
// so individual tools never touch map[string]any.
type typed[T any] struct {
	name, desc string
	schema     json.RawMessage
	fn         func(context.Context, T) (string, error)
}

func (t typed[T]) Name() string                 { return t.name }
func (t typed[T]) Description() string          { return t.desc }
func (t typed[T]) InputSchema() json.RawMessage { return t.schema }

func (t typed[T]) Execute(ctx context.Context, raw json.RawMessage) (string, error) {
	if len(raw) == 0 {
		raw = json.RawMessage("{}")
	}
	var in T
	if err := json.Unmarshal(raw, &in); err != nil {
		return "", fmt.Errorf("invalid arguments: %w", err)
	}
	return t.fn(ctx, in)
}

// New builds a Tool whose input schema is derived from T's struct tags. It
// panics if T cannot be reflected, which is a programming error at wiring time.
func New[T any](name, desc string, fn func(context.Context, T) (string, error)) Tool {
	return typed[T]{name: name, desc: desc, schema: SchemaFor[T](), fn: fn}
}

// reflector inlines everything: the Messages API wants a self-contained schema
// object, not one stitched together from $defs.
var reflector = &jsonschema.Reflector{
	ExpandedStruct: true,
	DoNotReference: true,
}

// SchemaFor derives a JSON Schema from T. Fields tagged omitempty are optional;
// everything else is required. Descriptions come from `jsonschema_description:"..."`.
func SchemaFor[T any]() json.RawMessage {
	var zero T
	b, err := json.Marshal(reflector.Reflect(zero))
	if err != nil {
		panic(fmt.Sprintf("tools: cannot reflect schema for %T: %v", zero, err))
	}
	var m map[string]any
	if err := json.Unmarshal(b, &m); err != nil {
		panic(fmt.Sprintf("tools: cannot reflect schema for %T: %v", zero, err))
	}
	// Metadata the API has no use for, and which would only add bytes to the
	// cached prompt prefix.
	delete(m, "$schema")
	delete(m, "$id")
	delete(m, "version")
	if _, ok := m["type"]; !ok {
		m["type"] = "object"
	}
	if _, ok := m["properties"]; !ok {
		m["properties"] = map[string]any{}
	}
	out, err := json.Marshal(m)
	if err != nil {
		panic(fmt.Sprintf("tools: cannot reflect schema for %T: %v", zero, err))
	}
	return out
}

// Registry holds the tools available to one session and is the single place
// tool dispatch gets logged, so no individual tool has to remember to do it.
type Registry struct {
	mu     sync.RWMutex
	byName map[string]Tool
	order  []string
	log    *metrics.Logger
	// drop is what Exclude was told to keep out.
	drop func(name string) bool
}

// NewRegistry returns an empty registry. log may be nil.
func NewRegistry(log *metrics.Logger) *Registry {
	return &Registry{byName: map[string]Tool{}, log: log}
}

// Register adds a tool, panicking on a duplicate name.
func (r *Registry) Register(ts ...Tool) *Registry {
	r.mu.Lock()
	defer r.mu.Unlock()
	for _, t := range ts {
		if r.drop != nil && r.drop(t.Name()) {
			continue
		}
		if _, dup := r.byName[t.Name()]; dup {
			panic(fmt.Sprintf("tools: duplicate registration of %q", t.Name()))
		}
		r.byName[t.Name()] = t
		r.order = append(r.order, t.Name())
	}
	return r
}

// Exclude keeps every tool drop says yes to out of the registry: the ones
// already in it are taken out, and the ones registered later are skipped. It
// is how --disallowedTools works — a tool the model is never shown is one it
// cannot call, which beats refusing it every time it tries — and it works on
// the registry rather than on a list of names because the tools of a session
// are registered from several places, and every one of them has to hear it.
func (r *Registry) Exclude(drop func(name string) bool) *Registry {
	r.mu.Lock()
	defer r.mu.Unlock()
	prev := r.drop
	r.drop = func(name string) bool { return (prev != nil && prev(name)) || drop(name) }
	kept := r.order[:0]
	for _, n := range r.order {
		if drop(n) {
			delete(r.byName, n)
			continue
		}
		kept = append(kept, n)
	}
	r.order = kept
	return r
}

// Get returns the named tool.
func (r *Registry) Get(name string) (Tool, bool) {
	r.mu.RLock()
	defer r.mu.RUnlock()
	t, ok := r.byName[name]
	return t, ok
}

// List returns the tools in a stable order, so the tool-definition block of the
// prompt stays byte-identical between turns and keeps the cache prefix warm.
func (r *Registry) List() []Tool {
	r.mu.RLock()
	defer r.mu.RUnlock()
	names := append([]string(nil), r.order...)
	sort.Strings(names)
	out := make([]Tool, 0, len(names))
	for _, n := range names {
		out = append(out, r.byName[n])
	}
	return out
}

// Execute dispatches by name and logs the outcome. An unknown tool is reported
// back to the model as an error rather than failing the turn.
func (r *Registry) Execute(ctx context.Context, name string, raw json.RawMessage) (Output, error) {
	t, ok := r.Get(name)
	if !ok {
		err := fmt.Errorf("unknown tool %q", name)
		r.log.ToolCall(name, 0, err, nil)
		return Output{}, err
	}
	start := time.Now()
	var out Output
	var err error
	if bt, ok := t.(BlockTool); ok {
		out, err = bt.ExecuteBlocks(ctx, raw)
	} else {
		out.Text, err = t.Execute(ctx, raw)
	}
	extra := map[string]any{"output_bytes": len(out.Text)}
	if len(out.Blocks) > 0 {
		extra["blocks"] = len(out.Blocks)
	}
	r.log.ToolCall(name, time.Since(start), err, extra)
	return out, err
}