package main

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

  "github.com/charmbracelet/bubbles/textarea"
  "github.com/charmbracelet/bubbles/viewport"
  tea "github.com/charmbracelet/bubbletea"
  "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)
)

type model struct {
  viewport    viewport.Model
  textarea    textarea.Model
  agent       *sidekick.Sidekick
  pool        map[string]*sidekick.Sidekick
  messages    []string
  history     []sidekick.Message
  isThinking  bool
  err         error
  ready       bool
  width       int
  height      int
}

type agentResponseMsg struct {
  response string
  err      error
}

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

  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)
    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, nil)
    }
    entryAgent := "coordinator"
    if _, ok := pool[entryAgent]; !ok {
      for id := range pool {
        entryAgent = id
        break
      }
    }
    agent = pool[entryAgent]
  }

  return model{
    textarea: ta,
    agent:    agent,
    pool:     pool,
    messages: []string{sysStyle.Render("System: Sidekick initialized. Type a message below.")},
    history:  []sidekick.Message{},
  }
}

func (m model) Init() tea.Cmd { return textarea.Blink }

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 !m.ready {
      m.viewport = viewport.New(vpWidth, vpHeight)
      m.viewport.SetContent(strings.Join(m.messages, "\n\n"))
      m.ready = true
    } else {
      m.viewport.Width = vpWidth
      m.viewport.Height = vpHeight
    }
    
    m.textarea.SetWidth(msg.Width - 4)

  case tea.KeyMsg:
    switch msg.Type {
    case tea.KeyCtrlC, tea.KeyEsc:
      return m, tea.Quit
    case tea.KeyEnter:
      if msg.Alt {
        break
      }
      v := m.textarea.Value()
      if strings.TrimSpace(v) == "" || m.isThinking {
        return m, nil
      }
      
      wrappedUser := lipgloss.NewStyle().Width(m.viewport.Width).Render(userStyle.Render("You") + "\n" + v)
      m.messages = append(m.messages, wrappedUser)
      m.textarea.Reset()
      m.viewport.SetContent(strings.Join(m.messages, "\n\n"))
      m.viewport.GotoBottom()
      m.isThinking = true
      return m, m.runAgent(v)
    }

  case agentResponseMsg:
    m.isThinking = false
    var newMsg string
    if msg.err != nil {
      newMsg = sysStyle.Render(fmt.Sprintf("Error: %v", msg.err))
    } else {
      newMsg = agentStyle.Render("Sidekick") + "\n" + msg.response
    }
    wrappedMsg := lipgloss.NewStyle().Width(m.viewport.Width).Render(newMsg)
    m.messages = append(m.messages, wrappedMsg)
    m.viewport.SetContent(strings.Join(m.messages, "\n\n"))
    m.viewport.GotoBottom()
    return m, nil
  }

  m.textarea, tiCmd = m.textarea.Update(msg)
  m.viewport, vpCmd = m.viewport.Update(msg)
  return m, tea.Batch(tiCmd, vpCmd)
}

func (m model) runAgent(prompt string) tea.Cmd {
  return func() tea.Msg {
    res, err := m.agent.Run(context.Background(), []sidekick.Message{{Role: sidekick.MessageRoleUser, Content: prompt}}, 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 NG "),
    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() {
  if _, err := tea.NewProgram(initialModel(), tea.WithAltScreen()).Run(); err != nil {
    log.Fatal(err)
  }
}