package main

import (
	"fmt"
	"strings"

	"github.com/charmbracelet/bubbles/textinput"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"

	llm "gitea.stevedudenhoeffer.com/steve/go-llm/v2"
)

// pendingToolCalls stores the last response's tool calls so we can pair them
// with tool execution results for display.
var pendingToolCalls []llm.ToolCall

// Update handles messages and updates the model.
func (m Model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	var cmd tea.Cmd
	var cmds []tea.Cmd

	switch msg := msg.(type) {
	case tea.KeyMsg:
		return m.handleKeyMsg(msg)

	case tea.WindowSizeMsg:
		m.width = msg.Width
		m.height = msg.Height

		headerHeight := 3
		footerHeight := 4
		verticalMargins := headerHeight + footerHeight

		if !m.viewportReady {
			m.viewport = viewport.New(msg.Width-4, msg.Height-verticalMargins)
			m.viewport.HighPerformanceRendering = false
			m.viewportReady = true
		} else {
			m.viewport.Width = msg.Width - 4
			m.viewport.Height = msg.Height - verticalMargins
		}

		m.input.Width = msg.Width - 6
		m.apiKeyInput.Width = msg.Width - 6

		m.viewport.SetContent(m.renderMessages())

	case ChatResponseMsg:
		m.loading = false
		if msg.Err != nil {
			m.err = msg.Err
			return m, nil
		}

		resp := msg.Response

		// Add the assistant message to the conversation history.
		m.conversation = append(m.conversation, resp.Message())

		// Show any text the assistant produced alongside tool calls.
		if resp.Text != "" {
			m.addAssistantMessage(resp.Text)
		}

		if resp.HasToolCalls() && m.toolsEnabled {
			pendingToolCalls = resp.ToolCalls

			for _, call := range resp.ToolCalls {
				m.addToolCallMessage(call.Name, call.Arguments)
			}

			m.viewport.SetContent(m.renderMessages())
			m.viewport.GotoBottom()

			m.loading = true
			return m, executeTools(m.toolbox, resp.ToolCalls)
		}

		m.viewport.SetContent(m.renderMessages())
		m.viewport.GotoBottom()

	case ToolExecutionMsg:
		if msg.Err != nil {
			m.loading = false
			m.err = msg.Err
			return m, nil
		}

		// Display results paired with the tool calls that produced them.
		for i, result := range msg.Results {
			name := ""
			if i < len(pendingToolCalls) {
				name = pendingToolCalls[i].Name
			}
			m.addToolResultMessage(name, result.Content.Text)
		}

		// Append the raw tool result messages to the conversation so the
		// assistant can reference them on the next turn.
		m.conversation = append(m.conversation, msg.Results...)

		m.viewport.SetContent(m.renderMessages())
		m.viewport.GotoBottom()

		// Ask the model to continue given the tool results.
		return m, sendChatRequest(m.chat, m.conversation, m.toolbox, m.toolsEnabled, m.temperature)

	case ImageLoadedMsg:
		if msg.Err != nil {
			m.err = msg.Err
			m.state = m.previousState
			return m, nil
		}

		m.pendingImages = append(m.pendingImages, msg.Image)
		m.state = m.previousState
		m.err = nil

	default:
		// Update text input.
		if m.state == StateChat {
			m.input, cmd = m.input.Update(msg)
			cmds = append(cmds, cmd)
		} else if m.state == StateAPIKeyInput {
			m.apiKeyInput, cmd = m.apiKeyInput.Update(msg)
			cmds = append(cmds, cmd)
		}
	}

	return m, tea.Batch(cmds...)
}

// handleKeyMsg handles keyboard input.
func (m Model) handleKeyMsg(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch msg.String() {
	case "ctrl+c":
		return m, tea.Quit
	case "esc":
		if m.state != StateChat {
			m.state = StateChat
			m.input.Focus()
			return m, nil
		}
		return m, tea.Quit
	}

	switch m.state {
	case StateChat:
		return m.handleChatKeys(msg)
	case StateProviderSelect:
		return m.handleProviderSelectKeys(msg)
	case StateModelSelect:
		return m.handleModelSelectKeys(msg)
	case StateImageInput:
		return m.handleImageInputKeys(msg)
	case StateToolsPanel:
		return m.handleToolsPanelKeys(msg)
	case StateSettings:
		return m.handleSettingsKeys(msg)
	case StateAPIKeyInput:
		return m.handleAPIKeyInputKeys(msg)
	}

	return m, nil
}

// handleChatKeys handles keys in chat state.
func (m Model) handleChatKeys(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch msg.String() {
	case "enter":
		if m.loading {
			return m, nil
		}

		text := strings.TrimSpace(m.input.Value())
		if text == "" {
			return m, nil
		}

		if m.chat == nil {
			m.err = fmt.Errorf("no model selected - press Ctrl+P to select a provider")
			return m, nil
		}

		// Ensure a system message is at the head of the conversation.
		if len(m.conversation) == 0 && m.systemPrompt != "" {
			m.conversation = append(m.conversation, llm.SystemMessage(m.systemPrompt))
		}

		m.addUserMessage(text, m.pendingImages)

		m.input.Reset()
		m.pendingImages = nil
		m.err = nil
		m.loading = true

		m.viewport.SetContent(m.renderMessages())
		m.viewport.GotoBottom()

		return m, sendChatRequest(m.chat, m.conversation, m.toolbox, m.toolsEnabled, m.temperature)

	case "ctrl+i":
		m.previousState = StateChat
		m.state = StateImageInput
		m.input.SetValue("")
		m.input.Placeholder = "Enter image path or URL..."
		return m, nil

	case "ctrl+t":
		m.state = StateToolsPanel
		return m, nil

	case "ctrl+p":
		m.state = StateProviderSelect
		m.listIndex = m.providerIndex
		return m, nil

	case "ctrl+m":
		if m.client == nil {
			m.err = fmt.Errorf("select a provider first")
			return m, nil
		}
		m.state = StateModelSelect
		m.listItems = m.providers[m.providerIndex].Info.Models
		m.listIndex = m.providers[m.providerIndex].ModelIndex
		return m, nil

	case "ctrl+s":
		m.state = StateSettings
		return m, nil

	case "ctrl+n":
		m.newConversation()
		m.viewport.SetContent(m.renderMessages())
		return m, nil

	case "up", "down", "pgup", "pgdown":
		var cmd tea.Cmd
		m.viewport, cmd = m.viewport.Update(msg)
		return m, cmd

	default:
		var cmd tea.Cmd
		m.input, cmd = m.input.Update(msg)
		return m, cmd
	}
}

// handleProviderSelectKeys handles keys in provider selection state.
func (m Model) handleProviderSelectKeys(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch msg.String() {
	case "up", "k":
		if m.listIndex > 0 {
			m.listIndex--
		}
	case "down", "j":
		if m.listIndex < len(m.providers)-1 {
			m.listIndex++
		}
	case "enter":
		p := m.providers[m.listIndex]
		if !p.HasAPIKey {
			m.state = StateAPIKeyInput
			m.apiKeyInput.Focus()
			m.apiKeyInput.SetValue("")
			return m, textinput.Blink
		}

		if err := m.selectProvider(m.listIndex); err != nil {
			m.err = err
			return m, nil
		}

		m.state = StateChat
		m.input.Focus()
		m.newConversation()
		return m, nil
	}
	return m, nil
}

// handleAPIKeyInputKeys handles keys in API key input state.
func (m Model) handleAPIKeyInputKeys(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch msg.String() {
	case "enter":
		key := strings.TrimSpace(m.apiKeyInput.Value())
		if key == "" {
			return m, nil
		}

		p := m.providers[m.listIndex]
		m.apiKeys[p.Info.Name] = key
		m.providers[m.listIndex].HasAPIKey = true

		for i, prov := range m.providers {
			status := " (no key)"
			if prov.HasAPIKey {
				status = " (ready)"
				if prov.Info.EnvKey == "" {
					status = " (local)"
				}
			}
			m.listItems[i] = prov.Info.DisplayName + status
		}

		if err := m.selectProvider(m.listIndex); err != nil {
			m.err = err
			return m, nil
		}

		m.state = StateChat
		m.input.Focus()
		m.newConversation()
		return m, nil

	default:
		var cmd tea.Cmd
		m.apiKeyInput, cmd = m.apiKeyInput.Update(msg)
		return m, cmd
	}
}

// handleModelSelectKeys handles keys in model selection state.
func (m Model) handleModelSelectKeys(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch msg.String() {
	case "up", "k":
		if m.listIndex > 0 {
			m.listIndex--
		}
	case "down", "j":
		if m.listIndex < len(m.listItems)-1 {
			m.listIndex++
		}
	case "enter":
		if err := m.selectModel(m.listIndex); err != nil {
			m.err = err
			return m, nil
		}
		m.state = StateChat
		m.input.Focus()
	}
	return m, nil
}

// handleImageInputKeys handles keys in image input state.
func (m Model) handleImageInputKeys(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch msg.String() {
	case "enter":
		input := strings.TrimSpace(m.input.Value())
		if input == "" {
			m.state = m.previousState
			m.input.Placeholder = "Type your message..."
			return m, nil
		}

		m.input.Placeholder = "Type your message..."

		switch {
		case strings.HasPrefix(input, "http://") || strings.HasPrefix(input, "https://"):
			return m, loadImageFromURL(input)
		case strings.HasPrefix(input, "data:") || (len(input) > 100 && !strings.Contains(input, "/") && !strings.Contains(input, "\\")):
			return m, loadImageFromBase64(input)
		default:
			return m, loadImageFromPath(input)
		}

	default:
		var cmd tea.Cmd
		m.input, cmd = m.input.Update(msg)
		return m, cmd
	}
}

// handleToolsPanelKeys handles keys in tools panel state.
func (m Model) handleToolsPanelKeys(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch msg.String() {
	case "t":
		m.toolsEnabled = !m.toolsEnabled
	case "enter", "q":
		m.state = StateChat
		m.input.Focus()
	}
	return m, nil
}

// handleSettingsKeys handles keys in settings state.
func (m Model) handleSettingsKeys(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
	switch msg.String() {
	case "1":
		m.temperature = nil
	case "2":
		t := 0.0
		m.temperature = &t
	case "3":
		t := 0.5
		m.temperature = &t
	case "4":
		t := 0.7
		m.temperature = &t
	case "5":
		t := 1.0
		m.temperature = &t
	case "enter", "q":
		m.state = StateChat
		m.input.Focus()
	}
	return m, nil
}