package main

import tea "github.com/charmbracelet/bubbletea"

type InputState int

const (
	StateReady InputState = iota
	StateCount
	StateOperatorPending
	StateMotionCount
)

type InputHandler struct {
	state       InputState
	count1      int
	count2      int
	operator    Operator
	operatorKey string // track which key started operator (for dd, yy, cc)
	buffer      string // for display (what user has typed)
	pending     string // partial key sequence (e.g., "g" waiting for second key)
	keymap      *Keymap
}

func NewInputHandler() *InputHandler {
	return &InputHandler{
		keymap: NewNormalKeymap(),
	}
}

func (h *InputHandler) Handle(m *model, key string) tea.Cmd {
	// ESC always resets everything
	if key == "esc" {
		h.Reset()
		return nil
	}

	// Try to accumulate count (only if no pending sequence)
	if h.pending == "" && h.tryAccumulateCount(key) {
		return nil
	}

	// Build the sequence (pending + new key)
	sequence := h.pending + key

	// Check for exact match with full sequence
	kind, binding := h.keymap.Lookup(sequence)
	if kind != "" {
		h.pending = ""
		h.buffer += key
		return h.dispatch(m, kind, binding, sequence)
	}

	// No exact match - could this be a prefix of something?
	if h.keymap.HasPrefix(sequence) {
		h.pending = sequence
		h.buffer += key
		return nil // wait for more keys
	}

	// Not a prefix either - if we had pending, try just the new key
	if h.pending != "" {
		h.pending = ""
		kind, binding = h.keymap.Lookup(key)
		if kind != "" {
			h.buffer = key
			return h.dispatch(m, kind, binding, key)
		}
	}

	// Nothing matched
	h.Reset()
	return nil
}

// dispatch routes to the right handler based on current state
func (h *InputHandler) dispatch(m *model, kind string, binding any, key string) tea.Cmd {
	switch h.state {
	case StateReady, StateCount:
		return h.handleInitial(m, kind, binding, key)
	case StateOperatorPending, StateMotionCount:
		return h.handleAfterOperator(m, kind, binding, key)
	}
	h.Reset()
	return nil
}

func (h *InputHandler) handleInitial(m *model, kind string, binding any, key string) tea.Cmd {
	count := h.effectiveCount()

	switch kind {
	case "motion":
		motion := binding.(Motion)
		if r, ok := motion.(Repeatable); ok {
			motion = r.WithCount(count).(Motion)
		}
		cmd := motion.Execute(m)
		h.Reset()
		return cmd

	case "operator":
		h.operator = binding.(Operator)
		h.operatorKey = key
		h.state = StateOperatorPending
		return nil

	case "action":
		action := binding.(Action)
		if r, ok := action.(Repeatable); ok {
			action = r.WithCount(count)
		}
		cmd := action.Execute(m)
		h.Reset()
		return cmd
	}

	h.Reset()
	return nil
}

func (h *InputHandler) handleAfterOperator(m *model, kind string, binding any, key string) tea.Cmd {
	count := h.effectiveCount()

	// dd, yy, cc - same operator key pressed twice
	if kind == "operator" && key == h.operatorKey {
		cmd := h.operator.DoublePress(m)
		h.Reset()
		return cmd
	}

	// Motion after operator
	if kind == "motion" {
		motion := binding.(Motion)
		if r, ok := motion.(Repeatable); ok {
			motion = r.WithCount(count).(Motion)
		}
		// Get range here
		start := m.getCursorPosition()
		motion.Execute(m)
		end := m.getCursorPosition()
		cmd := h.operator.Operate(m, start, end)
		h.Reset()
		return cmd
	}

	h.Reset()
	return nil
}

func (h *InputHandler) tryAccumulateCount(key string) bool {
	if len(key) != 1 || key[0] < '0' || key[0] > '9' {
		return false
	}

	digit := int(key[0] - '0')

	// 0 at start is a motion, not a count
	if digit == 0 && h.currentCount() == 0 {
		return false
	}

	switch h.state {
	case StateReady, StateCount:
		h.count1 = h.count1*10 + digit
		h.state = StateCount
	case StateOperatorPending, StateMotionCount:
		h.count2 = h.count2*10 + digit
		h.state = StateMotionCount
	}

	h.buffer += key
	return true
}

func (h *InputHandler) currentCount() int {
	if h.state == StateOperatorPending || h.state == StateMotionCount {
		return h.count2
	}
	return h.count1
}

func (h *InputHandler) effectiveCount() int {
	c1, c2 := h.count1, h.count2
	if c1 == 0 {
		c1 = 1
	}
	if c2 == 0 {
		c2 = 1
	}
	return c1 * c2
}

func (h *InputHandler) Reset() {
	h.state = StateReady
	h.count1 = 0
	h.count2 = 0
	h.operator = nil
	h.operatorKey = ""
	h.buffer = ""
	h.pending = ""
}

func (h *InputHandler) Pending() string {
	return h.buffer
}