package contrib

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"sync"

	"github.com/keepmind9/acp-sdk-go/schema"
)

// ToolCallTracker manages tool call lifecycle and generates ACP notifications.
// It tracks tool calls by external ID and produces SessionUpdate values
// suitable for sending via AgentSideConnection.SessionUpdate.
type ToolCallTracker struct {
	mu         sync.Mutex
	idFactory  func() string
	calls      map[string]*trackedToolCall
	streamBufs map[string]string
}

// trackedToolCall holds the mutable state of a tool call.
type trackedToolCall struct {
	toolCallID string
	title      string
	kind       schema.ToolKind
	status     schema.ToolCallStatus
	content    []*schema.ToolCallContent
	locations  []*schema.ToolCallLocation
	rawInput   json.RawMessage
	rawOutput  json.RawMessage
}

func defaultIDFactory() func() string {
	return func() string {
		b := make([]byte, 16)
		rand.Read(b)
		return hex.EncodeToString(b)
	}
}

// NewToolCallTracker creates a new ToolCallTracker.
// If idFactory is nil, random hex IDs are generated.
func NewToolCallTracker(idFactory func() string) *ToolCallTracker {
	if idFactory == nil {
		idFactory = defaultIDFactory()
	}
	return &ToolCallTracker{
		idFactory:  idFactory,
		calls:      make(map[string]*trackedToolCall),
		streamBufs: make(map[string]string),
	}
}

// Start registers a new tool call and returns a SessionUpdate for the tool_call notification.
func (t *ToolCallTracker) Start(externalID string, title string, opts ...ToolCallOption) schema.SessionUpdate {
	t.mu.Lock()
	defer t.mu.Unlock()

	callID := t.idFactory()
	tc := &trackedToolCall{
		toolCallID: callID,
		title:      title,
		status:     schema.ToolCallStatusInProgress,
	}
	for _, opt := range opts {
		opt(tc)
	}
	t.calls[externalID] = tc

	return t.sessionUpdateForToolCall(tc)
}

// Progress updates a tracked tool call and returns a SessionUpdate.
func (t *ToolCallTracker) Progress(externalID string, opts ...ToolCallOption) (schema.SessionUpdate, error) {
	t.mu.Lock()
	defer t.mu.Unlock()

	tc, ok := t.calls[externalID]
	if !ok {
		return schema.SessionUpdate{}, fmt.Errorf("unknown tool call: %s", externalID)
	}

	for _, opt := range opts {
		opt(tc)
	}

	return t.sessionUpdateForUpdate(tc), nil
}

// AppendStreamText appends text to the tool call's streaming content buffer.
func (t *ToolCallTracker) AppendStreamText(externalID string, text string, opts ...ToolCallOption) (schema.SessionUpdate, error) {
	t.mu.Lock()
	defer t.mu.Unlock()

	tc, ok := t.calls[externalID]
	if !ok {
		return schema.SessionUpdate{}, fmt.Errorf("unknown tool call: %s", externalID)
	}

	t.streamBufs[externalID] += text
	tc.content = []*schema.ToolCallContent{
		{Type: schema.TypeToolCallContentKindContent, Content: &schema.Content{Content: &schema.ContentBlock{Type: schema.TypeContentBlockKindText, Text: &schema.TextContent{Text: t.streamBufs[externalID]}}}},
	}

	for _, opt := range opts {
		opt(tc)
	}

	return t.sessionUpdateForUpdate(tc), nil
}

// Forget removes a tracked tool call.
func (t *ToolCallTracker) Forget(externalID string) {
	t.mu.Lock()
	defer t.mu.Unlock()
	delete(t.calls, externalID)
	delete(t.streamBufs, externalID)
}

// ToolCallUpdate returns a copy of the tool call as a ToolCallUpdate
// suitable for permission requests.
func (t *ToolCallTracker) ToolCallUpdate(externalID string) (schema.ToolCallUpdate, error) {
	t.mu.Lock()
	defer t.mu.Unlock()

	tc, ok := t.calls[externalID]
	if !ok {
		return schema.ToolCallUpdate{}, fmt.Errorf("unknown tool call: %s", externalID)
	}

	return t.toolCallUpdateFor(tc), nil
}

func (t *ToolCallTracker) sessionUpdateForToolCall(tc *trackedToolCall) schema.SessionUpdate {
	return schema.SessionUpdate{
		SessionUpdate: schema.SessionUpdateKindToolCall,
		ToolCall:      t.toolCallFor(tc),
	}
}

func (t *ToolCallTracker) sessionUpdateForUpdate(tc *trackedToolCall) schema.SessionUpdate {
	u := t.toolCallUpdateFor(tc)
	return schema.SessionUpdate{
		SessionUpdate:  schema.SessionUpdateKindToolCallUpdate,
		ToolCallUpdate: &u,
	}
}

func (t *ToolCallTracker) toolCallFor(tc *trackedToolCall) *schema.ToolCall {
	tcid := schema.ToolCallId(tc.toolCallID)
	title := tc.title
	var kind *schema.ToolKind
	if tc.kind != "" {
		kind = &tc.kind
	}
	var status *schema.ToolCallStatus
	if tc.status != "" {
		status = &tc.status
	}
	return &schema.ToolCall{
		ToolCallId: &tcid,
		Title:      title,
		Kind:       kind,
		Status:     status,
		Content:    tc.content,
		Locations:  tc.locations,
		RawInput:   tc.rawInput,
		RawOutput:  tc.rawOutput,
	}
}

func (t *ToolCallTracker) toolCallUpdateFor(tc *trackedToolCall) schema.ToolCallUpdate {
	tcid := schema.ToolCallId(tc.toolCallID)
	title := tc.title
	var kind *schema.ToolKind
	if tc.kind != "" {
		kind = &tc.kind
	}
	var status *schema.ToolCallStatus
	if tc.status != "" {
		status = &tc.status
	}
	return schema.ToolCallUpdate{
		ToolCallId: &tcid,
		Title:      &title,
		Kind:       kind,
		Status:     status,
		Content:    tc.content,
		RawInput:   tc.rawInput,
		RawOutput:  tc.rawOutput,
	}
}

// ToolCallOption configures a tracked tool call.
type ToolCallOption func(*trackedToolCall)

// WithKind sets the tool kind.
func WithKind(kind schema.ToolKind) ToolCallOption {
	return func(tc *trackedToolCall) { tc.kind = kind }
}

// WithStatus sets the tool call status.
func WithStatus(status schema.ToolCallStatus) ToolCallOption {
	return func(tc *trackedToolCall) { tc.status = status }
}

// WithContent sets the tool call content.
func WithContent(content []*schema.ToolCallContent) ToolCallOption {
	return func(tc *trackedToolCall) { tc.content = content }
}

// WithLocations sets the tool call locations.
func WithLocations(locations []*schema.ToolCallLocation) ToolCallOption {
	return func(tc *trackedToolCall) { tc.locations = locations }
}

// WithRawInput sets the raw input.
func WithRawInput(input json.RawMessage) ToolCallOption {
	return func(tc *trackedToolCall) { tc.rawInput = input }
}

// WithRawOutput sets the raw output.
func WithRawOutput(output json.RawMessage) ToolCallOption {
	return func(tc *trackedToolCall) { tc.rawOutput = output }
}