package main

import (
  "context"
  "fmt"
  "log"
  "os"
  "strings"

  "github.com/charmbracelet/bubbles/textarea"
  "github.com/charmbracelet/bubbles/viewport"
  tea "github.com/charmbracelet/bubbletea"
  "github.com/charmbracelet/glamour"
  "github.com/charmbracelet/lipgloss"
  "sidekick"
)

var (
  titleStyle = lipgloss.NewStyle().
      Foreground(lipgloss.Color("#FFFDF5")).
      Background(lipgloss.Color("#25A065")).
      Padding(0, 1).
      Bold(true)
  statusStyle = lipgloss.NewStyle().
      Foreground(lipgloss.Color("#FFFDF5")).
      Background(lipgloss.Color("#3C3C3C")).
      Padding(0, 1)
  viewportStyle = lipgloss.NewStyle().
      Border(lipgloss.RoundedBorder()).
      BorderForeground(lipgloss.Color("62")).
      Padding(0, 1)
  textareaStyle = lipgloss.NewStyle().
      Border(lipgloss.RoundedBorder()).
      BorderForeground(lipgloss.Color("240")).
      Padding(0, 1)
  textareaFocusStyle = lipgloss.NewStyle().
        Border(lipgloss.RoundedBorder()).
        BorderForeground(lipgloss.Color("62")).
        Padding(0, 1)
  userStyle   = lipgloss.NewStyle().Foreground(lipgloss.Color("6")).Bold(true)
  agentStyle  = lipgloss.NewStyle().Foreground(lipgloss.Color("5")).Bold(true)
  sysStyle    = lipgloss.NewStyle().Foreground(lipgloss.Color("8")).Italic(true)
  bannerStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("#FFFF00")).Bold(true)

  banner = `        _      __        __     _         __
  ___  (_) ___/ / ___   / /__  (_) ____  / /__
 (_-< / / / _  / / -_) /  '_/ / / / __/ /  '_/
/___//_/  \_,_/  \__/ /_/\_\ /_/  \__/ /_/\_\
----------------------------------------------`
)

type model struct {
  viewport         viewport.Model
  textarea         textarea.Model
  agent            *sidekick.Sidekick
  pool             map[string]*sidekick.Sidekick
  history          []sidekick.Message
  historyIndex     int
  promptDraft      string
  transientSteps   []string
  currentStream    string
  currentReasoning string
  cancel           context.CancelFunc
  isThinking       bool
  err              error
  ready            bool
  width            int
  height           int
  renderer         *glamour.TermRenderer
}

func (m model) formatReasoning(reasoning string) string {
  if reasoning == "" {
    return ""
  }
  lines := strings.Split(strings.TrimSpace(reasoning), "\n")
  for i, line := range lines {
    lines[i] = "> " + line
  }
  return "> *Thinking:*\n>\n" + strings.Join(lines, "\n") + "\n"
}

func (m model) userMessages() []string {
  var msgs []string
  for _, msg := range m.history {
    if msg.Role == sidekick.MessageRoleUser {
      msgs = append(msgs, msg.Content)
    }
  }
  return msgs
}

type agentResponseMsg struct {
  response string
  err      error
}

type streamMsg string
type reasoningMsg string
type intermediateMsg string

var streamChan = make(chan string, 100)
var reasoningChan = make(chan string, 100)
var intermediateChan = make(chan string, 100)

func waitForStream() tea.Cmd {
  return func() tea.Msg {
    return streamMsg(<-streamChan)
  }
}

func waitForReasoning() tea.Cmd {
  return func() tea.Msg {
    return reasoningMsg(<-reasoningChan)
  }
}

func waitForIntermediate() tea.Cmd {
  return func() tea.Msg {
    return intermediateMsg(<-intermediateChan)
  }
}

func initialModel() model {
  ta := textarea.New()
  ta.Placeholder = "Type a message..."
  ta.Focus()
  ta.Prompt = "┃ "
  ta.CharLimit = 2000
  ta.SetHeight(3)

  // Load TOML config
  cfg, err := sidekick.LoadConfig("config.toml")
  var agent *sidekick.Sidekick
  var pool map[string]*sidekick.Sidekick
  mcpStats := "No MCP servers configured"

  if err != nil {
    ac := sidekick.AgentConfig{ID: "tui", Role: sidekick.RoleSpecialist, MaxIterations: 5}
    agent = sidekick.NewSidekick(ac, sidekick.ModelConfig{}, nil)
  } else {
    // Attempt to load and apply prompts
    tmpl, _ := sidekick.LoadPrompts("prompts.toml")

    sidekick.InitMCPServers(cfg.MCPServers)
    extTools := sidekick.LoadExternalTools(context.Background())
    mcpStats = fmt.Sprintf("Connected to %d MCP servers (%d tools loaded)", len(cfg.MCPServers), len(extTools))

    pool = make(map[string]*sidekick.Sidekick)
    for id, aCfg := range cfg.Agents {
      if tmpl != nil && tmpl.Lookup(id) != nil {
        rendered, rErr := sidekick.RenderPrompt(tmpl, id)
        if rErr == nil {
          aCfg.SystemPrompt = rendered
        }
      }
      mCfg := cfg.Models[aCfg.ModelID]
      pool[id] = sidekick.NewSidekick(aCfg, mCfg, extTools)
    }
    entryAgent := "coordinator"
    if _, ok := pool[entryAgent]; !ok {
      for id := range pool {
        entryAgent = id
        break
      }
    }
    agent = pool[entryAgent]
  }

  // Force configs on all pool agents, including fallback 'agent'
  agentsToUpdate := []*sidekick.Sidekick{agent}
  for _, a := range pool {
    agentsToUpdate = append(agentsToUpdate, a)
  }
  for _, a := range agentsToUpdate {
    if a != nil {
      a.Config.Streaming = true
      a.Config.LogIntermediate = true
      a.StreamHandler = func(s string) {
        streamChan <- s
      }
      a.ReasoningHandler = func(s string) {
        reasoningChan <- s
      }
      a.IntermediateHandler = func(s string) {
        intermediateChan <- s
      }
    }

  }

  return model{
    textarea:     ta,
    agent:        agent,
    pool:         pool,
    historyIndex: -1,
    history: []sidekick.Message{
      {Role: sidekick.MessageRoleSystem, Content: banner + "\n\nSidekick initialized. " + mcpStats + "\nType a message below."},
    },
  }
}

func (m model) Init() tea.Cmd {
  return tea.Batch(textarea.Blink, waitForStream(), waitForReasoning(), waitForIntermediate())
}

func (m model) renderMessage(msg sidekick.Message) string {
  label := ""
  switch msg.Role {
  case sidekick.MessageRoleUser:
    label = userStyle.Render("You")
  case sidekick.MessageRoleAssistant:
    label = agentStyle.Render("Sidekick")
  case sidekick.MessageRoleSystem:
    label = sysStyle.Render("System")
  default:
    label = sysStyle.Render(string(msg.Role))
  }

  content := msg.Content
  reasoning := msg.Reasoning

  // Handle user message specifically to preserve newlines
  if msg.Role == sidekick.MessageRoleUser {
    if m.width > 4 {
      content = lipgloss.NewStyle().Width(m.width - 4).Render(content)
    }
    return label + "\n" + content
  }

  // If this is the banner message, handle it specially to preserve ASCII art
  if strings.Contains(content, "----") && msg.Role == sidekick.MessageRoleSystem {
    if strings.HasPrefix(content, banner) {
      bannerPart := bannerStyle.Render(banner)
      otherPart := strings.TrimPrefix(content, banner)
      if m.renderer != nil && strings.TrimSpace(otherPart) != "" {
        if r, err := m.renderer.Render(otherPart); err == nil {
          otherPart = "\n" + strings.TrimSpace(r)
        }
      }
      return label + "\n" + bannerPart + otherPart
    }
  }

  if m.renderer != nil {
    if reasoning != "" {
      if r, err := m.renderer.Render(m.formatReasoning(reasoning)); err == nil {
        reasoning = strings.TrimSpace(r) + "\n\n"
      }
    }
    if content != "" {
      if r, err := m.renderer.Render(content); err == nil {
        content = strings.TrimSpace(r)
      }
    }
  }

  if reasoning != "" {
    return label + "\n" + reasoning + content
  }
  return label + "\n" + content
}

func (m *model) updateViewportContent() {
  var rendered []string
  for _, msg := range m.history {
    rendered = append(rendered, m.renderMessage(msg))
  }
  for _, step := range m.transientSteps {
    wrapped := step
    if m.renderer != nil {
      if r, err := m.renderer.Render(step); err == nil {
        wrapped = strings.TrimSpace(r)
      }
    }
    rendered = append(rendered, sysStyle.Render("System")+"\n"+wrapped)
  }
  if m.currentReasoning != "" || m.currentStream != "" {
    content := m.currentStream
    reasoning := m.currentReasoning
    if m.renderer != nil {
      if reasoning != "" {
        if r, err := m.renderer.Render(m.formatReasoning(reasoning)); err == nil {
          reasoning = strings.TrimSpace(r) + "\n\n"
        }
      }
      if content != "" {
        if r, err := m.renderer.Render(content); err == nil {
          content = strings.TrimSpace(r)
        }
      }
    }
    rendered = append(rendered, agentStyle.Render("Sidekick")+"\n"+reasoning+content)
  }
  m.viewport.SetContent(strings.Join(rendered, "\n\n"))
}

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
  var (
    tiCmd tea.Cmd
    vpCmd tea.Cmd
  )

  switch msg := msg.(type) {
  case tea.WindowSizeMsg:
    m.width = msg.Width
    m.height = msg.Height

    headerHeight := 1
    inputHeight := 5 // Textarea height(3) + border/padding

    vpWidth := msg.Width - 4 // Account for viewportStyle padding and border
    vpHeight := msg.Height - headerHeight - inputHeight - 2

    if vpWidth > 0 && vpHeight > 0 {
      r, err := glamour.NewTermRenderer(
        glamour.WithStandardStyle("dark"),
        glamour.WithWordWrap(vpWidth-2),
      )
      if err == nil {
        m.renderer = r
      }

      if !m.ready {
        m.viewport = viewport.New(vpWidth, vpHeight)
        m.ready = true
      } else {
        m.viewport.Width = vpWidth
        m.viewport.Height = vpHeight
      }
      m.updateViewportContent()
      m.viewport.GotoBottom()
    }

    m.textarea.SetWidth(msg.Width - 4)

  case tea.KeyMsg:
    switch msg.Type {
    case tea.KeyCtrlC, tea.KeyEsc:
      return m, tea.Quit
    case tea.KeyPgUp, tea.KeyCtrlU:
      m.viewport.HalfViewUp()
      return m, nil
    case tea.KeyPgDown, tea.KeyCtrlD:
      if msg.Type == tea.KeyCtrlD && m.isThinking {
        if m.cancel != nil {
          m.cancel()
        }
        m.isThinking = false
        m.currentReasoning = ""
        m.currentStream = ""
        m.history = append(m.history, sidekick.Message{
          Role:    sidekick.MessageRoleSystem,
          Content: "Request cancelled.",
        })
        m.updateViewportContent()
        return m, nil
      }
      m.viewport.HalfViewDown()
      return m, nil
    case tea.KeyCtrlJ:
      m.textarea.InsertString("\n")
      return m, nil
    case tea.KeyUp:
      if m.textarea.Line() == 0 {
        userMsgs := m.userMessages()
        if len(userMsgs) == 0 {
          return m, nil
        }
        if m.historyIndex == -1 {
          m.promptDraft = m.textarea.Value()
          m.historyIndex = len(userMsgs) - 1
        } else if m.historyIndex > 0 {
          m.historyIndex--
        } else {
          return m, nil
        }
        m.textarea.SetValue(userMsgs[m.historyIndex])
        m.textarea.CursorEnd()
        return m, nil
      }
    case tea.KeyDown:
      if m.textarea.Line() == m.textarea.LineCount()-1 {
        if m.historyIndex == -1 {
          return m, nil
        }
        userMsgs := m.userMessages()
        if m.historyIndex < len(userMsgs)-1 {
          m.historyIndex++
          m.textarea.SetValue(userMsgs[m.historyIndex])
          m.textarea.CursorStart()
        } else {
          m.historyIndex = -1
          m.textarea.SetValue(m.promptDraft)
          m.textarea.CursorStart()
        }
        return m, nil
      }
    case tea.KeyEnter:
      v := m.textarea.Value()
      if strings.TrimSpace(v) == "" || m.isThinking {
        return m, nil
      }

      trimmed := strings.TrimSpace(v)
      if trimmed == "/clear" {
        m.history = []sidekick.Message{
          {Role: sidekick.MessageRoleSystem, Content: banner + "\n\nContext cleared."},
        }
        m.currentStream = ""
        m.currentReasoning = ""
        m.transientSteps = nil
        m.textarea.Reset()
        m.historyIndex = -1
        m.promptDraft = ""
        m.updateViewportContent()
        m.viewport.GotoBottom()
        return m, nil
      }

      m.history = append(m.history, sidekick.Message{
        Role:    sidekick.MessageRoleUser,
        Content: v,
      })
      m.historyIndex = -1
      m.promptDraft = ""
      m.currentStream = ""
      m.currentReasoning = ""
      m.transientSteps = nil
      m.updateViewportContent()
      m.textarea.Reset()
      m.viewport.GotoBottom()
      m.isThinking = true
      return m, m.runAgent()
    }

  case streamMsg:
    if !m.isThinking {
      return m, waitForStream()
    }
    m.currentStream += string(msg)
    m.updateViewportContent()
    m.viewport.GotoBottom()
    return m, waitForStream()
  case reasoningMsg:
    if !m.isThinking {
      return m, waitForReasoning()
    }
    m.currentReasoning += string(msg)
    m.updateViewportContent()
    m.viewport.GotoBottom()
    return m, waitForReasoning()
  case intermediateMsg:
    if !m.isThinking {
      return m, waitForIntermediate()
    }
    m.transientSteps = append(m.transientSteps, string(msg))
    m.updateViewportContent()
    m.viewport.GotoBottom()
    return m, waitForIntermediate()
  case agentResponseMsg:
    m.isThinking = false
    m.currentStream = ""

    m.currentReasoning = ""
    m.transientSteps = nil
    if msg.err != nil {
      m.history = append(m.history, sidekick.Message{
        Role:    sidekick.MessageRoleSystem,
        Content: fmt.Sprintf("Error: %v", msg.err),
      })
    } else {
      m.history = append(m.history, sidekick.Message{
        Role:    sidekick.MessageRoleAssistant,
        Content: msg.response,
      })
    }
    m.updateViewportContent()
    m.viewport.GotoBottom()
    return m, nil
  }

  m.textarea, tiCmd = m.textarea.Update(msg)

  if _, ok := msg.(tea.KeyMsg); !ok {
    m.viewport, vpCmd = m.viewport.Update(msg)
  }

  return m, tea.Batch(tiCmd, vpCmd)
}

func (m *model) runAgent() tea.Cmd {
  var ctx context.Context
  ctx, m.cancel = context.WithCancel(context.Background())
  return func() tea.Msg {
    res, err := m.agent.Run(ctx, m.history, m.pool)
    return agentResponseMsg{response: res, err: err}
  }
}

func (m model) View() string {
  if !m.ready {
    return "\n  Initializing..."
  }

  status := "IDLE"
  if m.isThinking {
    status = "THINKING"
  }

  header := lipgloss.JoinHorizontal(lipgloss.Top,
    titleStyle.Render(" SIDEKICK TUI "),
    statusStyle.Render(" "+status+" "),
  )

  vpView := viewportStyle.Width(m.width - 2).Render(m.viewport.View())

  taStyle := textareaStyle
  if m.textarea.Focused() {
    taStyle = textareaFocusStyle
  }
  taView := taStyle.Width(m.width - 2).Render(m.textarea.View())

  return lipgloss.JoinVertical(lipgloss.Left,
    header,
    vpView,
    taView,
  )
}

func main() {
  sidekick.SetupLogging(os.Stderr)
  if _, err := tea.NewProgram(initialModel(), tea.WithAltScreen()).Run(); err != nil {
    log.Fatal(err)
  }
}