package contrib

import (
	"encoding/json"
	"fmt"
	"sync"

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

// SessionSnapshot is an aggregated view of the current session state.
type SessionSnapshot struct {
	SessionID         string
	ToolCalls         map[string]ToolCallSnapshot
	PlanEntries       []*schema.PlanEntry
	CurrentModeID     string
	AvailableCommands []*schema.AvailableCommand
	AgentMessages     []schema.ContentChunk
	AgentThoughts     []schema.ContentChunk
	UserMessages      []schema.ContentChunk
}

// ToolCallSnapshot is an immutable view of a tool call.
type ToolCallSnapshot struct {
	ToolCallID string
	Title      string
	Kind       schema.ToolKind
	Status     schema.ToolCallStatus
	Content    []*schema.ToolCallContent
	RawInput   json.RawMessage
	RawOutput  json.RawMessage
}

// SessionSubscriber is a callback for session state changes.
type SessionSubscriber func(snapshot SessionSnapshot, notification schema.SessionNotification)

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

// SessionAccumulator merges SessionNotification objects into a session snapshot.
// It is used on the client side to track session state as updates arrive.
type SessionAccumulator struct {
	mu        sync.Mutex
	autoReset bool

	sessionID         string
	toolCalls         map[string]*mutableToolCallState
	planEntries       []*schema.PlanEntry
	currentModeID     string
	availableCommands []*schema.AvailableCommand
	agentMessages     []schema.ContentChunk
	agentThoughts     []schema.ContentChunk
	userMessages      []schema.ContentChunk

	subscribers []SessionSubscriber
}

// NewSessionAccumulator creates a new session state accumulator.
func NewSessionAccumulator(opts ...SessionAccumulatorOption) *SessionAccumulator {
	acc := &SessionAccumulator{
		autoReset: true,
		toolCalls: make(map[string]*mutableToolCallState),
	}
	for _, opt := range opts {
		opt(acc)
	}
	return acc
}

// SessionAccumulatorOption configures a SessionAccumulator.
type SessionAccumulatorOption func(*SessionAccumulator)

// AutoResetOnSessionChange controls whether the accumulator resets when
// it receives a notification for a different session.
func AutoResetOnSessionChange(autoReset bool) SessionAccumulatorOption {
	return func(acc *SessionAccumulator) { acc.autoReset = autoReset }
}

// Apply merges a session notification into the accumulator and returns a snapshot.
func (acc *SessionAccumulator) Apply(notification *schema.SessionNotification) (SessionSnapshot, error) {
	acc.mu.Lock()
	defer acc.mu.Unlock()

	// Validate session ID
	if acc.sessionID == "" {
		acc.sessionID = *notification.SessionId
	} else if *notification.SessionId != acc.sessionID {
		if !acc.autoReset {
			return SessionSnapshot{}, fmt.Errorf(
				"session mismatch: expected %s, got %s",
				acc.sessionID, *notification.SessionId,
			)
		}
		acc.reset()
		acc.sessionID = *notification.SessionId
	}

	acc.applyUpdate(*notification.Update)
	snap := acc.snapshot()
	acc.notifySubscribers(snap, *notification)
	return snap, nil
}

// Snapshot returns the current session state snapshot.
func (acc *SessionAccumulator) Snapshot() (SessionSnapshot, error) {
	acc.mu.Lock()
	defer acc.mu.Unlock()

	if acc.sessionID == "" {
		return SessionSnapshot{}, fmt.Errorf("no session data accumulated yet")
	}
	return acc.snapshot(), nil
}

// Subscribe registers a callback for state changes.
// Returns an unsubscribe function.
func (acc *SessionAccumulator) Subscribe(callback SessionSubscriber) func() {
	acc.mu.Lock()
	defer acc.mu.Unlock()

	acc.subscribers = append(acc.subscribers, callback)
	return func() {
		acc.mu.Lock()
		defer acc.mu.Unlock()
		for i, sub := range acc.subscribers {
			if fmt.Sprintf("%p", sub) == fmt.Sprintf("%p", callback) {
				acc.subscribers = append(acc.subscribers[:i], acc.subscribers[i+1:]...)
				break
			}
		}
	}
}

// Reset clears all accumulated state.
func (acc *SessionAccumulator) Reset() {
	acc.mu.Lock()
	defer acc.mu.Unlock()
	acc.reset()
}

func (acc *SessionAccumulator) reset() {
	acc.sessionID = ""
	acc.toolCalls = make(map[string]*mutableToolCallState)
	acc.planEntries = nil
	acc.currentModeID = ""
	acc.availableCommands = nil
	acc.agentMessages = nil
	acc.agentThoughts = nil
	acc.userMessages = nil
}

func (acc *SessionAccumulator) applyUpdate(update schema.SessionUpdate) {
	switch update.SessionUpdate {
	case schema.SessionUpdateKindToolCall:
		if update.ToolCall != nil {
			tc := update.ToolCall
			state := acc.toolCalls[string(*tc.ToolCallId)]
			if state == nil {
				state = &mutableToolCallState{toolCallID: string(*tc.ToolCallId)}
				acc.toolCalls[string(*tc.ToolCallId)] = state
			}
			state.title = tc.Title
			if tc.Kind != nil {
				state.kind = *tc.Kind
			}
			if tc.Status != nil {
				state.status = *tc.Status
			}
			state.content = tc.Content
			state.locations = tc.Locations
			if tc.RawInput != nil {
				state.rawInput, _ = tc.RawInput.(json.RawMessage)
			}
			if tc.RawOutput != nil {
				state.rawOutput, _ = tc.RawOutput.(json.RawMessage)
			}
		}

	case schema.SessionUpdateKindToolCallUpdate:
		if update.ToolCallUpdate != nil {
			tcID := string(*update.ToolCallUpdate.ToolCallId)
			state := acc.toolCalls[tcID]
			if state == nil {
				state = &mutableToolCallState{toolCallID: tcID}
				acc.toolCalls[tcID] = state
			}
			tcu := update.ToolCallUpdate
			if tcu.Title != nil {
				state.title = *tcu.Title
			}
			if tcu.Kind != nil {
				state.kind = *tcu.Kind
			}
			if tcu.Status != nil {
				state.status = *tcu.Status
			}
			if tcu.Content != nil {
				state.content = tcu.Content
			}
			if tcu.RawInput != nil {
				state.rawInput, _ = tcu.RawInput.(json.RawMessage)
			}
			if tcu.RawOutput != nil {
				state.rawOutput, _ = tcu.RawOutput.(json.RawMessage)
			}
		}

	case schema.SessionUpdateKindPlan:
		if update.Plan != nil {
			acc.planEntries = update.Plan.Entries
		}

	case schema.SessionUpdateKindCurrentModeUpdate:
		if update.CurrentModeUpdate != nil {
			acc.currentModeID = string(*update.CurrentModeUpdate.CurrentModeId)
		}

	case schema.SessionUpdateKindAvailableCommandsUpdate:
		if update.AvailableCommandsUpdate != nil {
			acc.availableCommands = update.AvailableCommandsUpdate.AvailableCommands
		}

	case schema.SessionUpdateKindAgentMessageChunk:
		if update.AgentMessageChunk != nil {
			acc.agentMessages = append(acc.agentMessages, *update.AgentMessageChunk)
		}

	case schema.SessionUpdateKindAgentThoughtChunk:
		if update.AgentThoughtChunk != nil {
			acc.agentThoughts = append(acc.agentThoughts, *update.AgentThoughtChunk)
		}

	case schema.SessionUpdateKindUserMessageChunk:
		if update.UserMessageChunk != nil {
			acc.userMessages = append(acc.userMessages, *update.UserMessageChunk)
		}
	}
}

func (acc *SessionAccumulator) snapshot() SessionSnapshot {
	toolCalls := make(map[string]ToolCallSnapshot, len(acc.toolCalls))
	for id, state := range acc.toolCalls {
		toolCalls[id] = ToolCallSnapshot{
			ToolCallID: state.toolCallID,
			Title:      state.title,
			Kind:       state.kind,
			Status:     state.status,
			Content:    state.content,
			RawInput:   state.rawInput,
			RawOutput:  state.rawOutput,
		}
	}

	return SessionSnapshot{
		SessionID:         acc.sessionID,
		ToolCalls:         toolCalls,
		PlanEntries:       acc.planEntries,
		CurrentModeID:     acc.currentModeID,
		AvailableCommands: acc.availableCommands,
		AgentMessages:     acc.agentMessages,
		AgentThoughts:     acc.agentThoughts,
		UserMessages:      acc.userMessages,
	}
}

func (acc *SessionAccumulator) notifySubscribers(snap SessionSnapshot, notif schema.SessionNotification) {
	for _, sub := range acc.subscribers {
		sub(snap, notif)
	}
}