package main

import (
	"fmt"
	"sort"
	"strings"
	"time"

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

// agentInfo tracks metadata about a discovered agent.
type agentInfo struct {
	id       string
	slug     string
	name     string // team agent name (from config or initial message)
	colorIdx int
	lastSeen time.Time
}

type model struct {
	// Directory being watched
	dir string

	// File watcher
	watcher *Watcher

	// Session tracking
	sessions       []SessionInfo // all sessions, sorted most recent first
	activeSession  string        // currently displayed session ID
	sessionIdx     int           // index into sessions slice

	// All events received, in order (unfiltered)
	events []Event

	// Agent tracking (for current session only)
	agents     []agentInfo         // sorted by most recent activity
	agentIndex map[string]int      // agentID -> index into agents
	agentLines map[string][]string // agentID -> rendered lines for that agent
	agentNames map[string]string   // agentID -> resolved name (from team config)

	// Stable color assignment (doesn't change when tabs reorder)
	agentColorMap map[string]int // agentID -> color index
	nextColor     int

	// All rendered lines (for "All" tab, current session)
	allLines []string

	// Viewport for scrollable content
	vp viewport.Model

	// UI state
	activeTabID string // agent ID of selected tab ("" = All tab)
	follow      bool   // auto-scroll to bottom
	showTokens  bool
	width       int
	height      int
	ready       bool
}

func newModel(dir string, sessions []SessionInfo, activeSession string) model {
	sessionIdx := 0
	for i, s := range sessions {
		if s.ID == activeSession {
			sessionIdx = i
			break
		}
	}

	// Try to load agent names from team config files
	agentNames := loadTeamAgentNames()

	return model{
		dir:           dir,
		sessions:      sessions,
		activeSession: activeSession,
		sessionIdx:    sessionIdx,
		agentIndex:    make(map[string]int),
		agentLines:    make(map[string][]string),
		agentNames:    agentNames,
		agentColorMap: make(map[string]int),
		follow:        true,
		showTokens:    false,
	}
}

// waitForEvent returns a tea.Cmd that waits for the next event from the watcher.
func (m model) waitForEvent() tea.Cmd {
	return func() tea.Msg {
		evt, ok := <-m.watcher.Events()
		if !ok {
			return nil
		}
		return evt
	}
}

func (m model) Init() tea.Cmd {
	// Watcher is already started by main.go before tea.NewProgram.
	// Just start listening for events.
	return m.waitForEvent()
}

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	var cmds []tea.Cmd

	switch msg := msg.(type) {
	case tea.KeyMsg:
		switch msg.String() {
		case "q", "ctrl+c":
			if m.watcher != nil {
				m.watcher.Stop()
			}
			return m, tea.Quit

		case "f":
			m.follow = !m.follow
			if m.follow {
				m.vp.GotoBottom()
			}

		case "t":
			m.showTokens = !m.showTokens
			m.rebuildAllLines()
			m.syncViewport()

		case "s":
			// Next session
			if len(m.sessions) > 1 {
				m.sessionIdx = (m.sessionIdx + 1) % len(m.sessions)
				m.switchSession(m.sessions[m.sessionIdx].ID)
			}

		case "S":
			// Previous session
			if len(m.sessions) > 1 {
				m.sessionIdx = (m.sessionIdx - 1 + len(m.sessions)) % len(m.sessions)
				m.switchSession(m.sessions[m.sessionIdx].ID)
			}

		case "left", "h":
			tabIdx := m.activeTabIndex()
			if tabIdx > 0 {
				tabIdx--
				if tabIdx == 0 {
					m.activeTabID = ""
				} else {
					m.activeTabID = m.agents[tabIdx-1].id
				}
				m.syncViewport()
			}

		case "right", "l":
			tabIdx := m.activeTabIndex()
			if tabIdx < len(m.agents) {
				tabIdx++
				m.activeTabID = m.agents[tabIdx-1].id
				m.syncViewport()
			}

		case "0", "a":
			m.activeTabID = ""
			m.syncViewport()

		case "1", "2", "3", "4", "5", "6", "7", "8", "9":
			n := int(msg.String()[0] - '0')
			if n <= len(m.agents) {
				m.activeTabID = m.agents[n-1].id
				m.syncViewport()
			}

		case "G", "end":
			m.follow = true
			m.vp.GotoBottom()

		case "g", "home":
			m.follow = false
			m.vp.GotoTop()
		}

	case tea.MouseMsg:
		// Let viewport handle mouse scrolling
		wasAtBottom := m.vp.AtBottom()
		var cmd tea.Cmd
		m.vp, cmd = m.vp.Update(msg)
		cmds = append(cmds, cmd)

		// If user scrolled up, disable follow
		if wasAtBottom && !m.vp.AtBottom() {
			m.follow = false
		}
		// If user scrolled to bottom, enable follow
		if m.vp.AtBottom() {
			m.follow = true
		}
		return m, tea.Batch(cmds...)

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

		headerHeight := 3 // session bar + tab bar + divider
		footerHeight := 1 // status bar
		vpHeight := m.height - headerHeight - footerHeight

		if !m.ready {
			m.vp = viewport.New(m.width, vpHeight)
			m.vp.MouseWheelEnabled = true
			m.vp.MouseWheelDelta = 3
			m.ready = true
		} else {
			m.vp.Width = m.width
			m.vp.Height = vpHeight
		}
		m.rebuildAllLines()
		m.syncViewport()

	case EventMsg:
		m.handleNewEvent(msg.Event)
		cmds = append(cmds, m.waitForEvent())

	case WatcherErrorMsg:
		// Could display error, for now just ignore
	}

	// Pass remaining messages to viewport (keyboard scrolling)
	if m.ready {
		wasAtBottom := m.vp.AtBottom()
		var cmd tea.Cmd
		m.vp, cmd = m.vp.Update(msg)
		cmds = append(cmds, cmd)

		if wasAtBottom && !m.vp.AtBottom() {
			m.follow = false
		}
	}

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

// switchSession changes the active session and rebuilds the view.
func (m *model) switchSession(sessionID string) {
	m.activeSession = sessionID
	m.activeTabID = ""

	// Reset agent tracking
	m.agents = nil
	m.agentIndex = make(map[string]int)
	m.agentLines = make(map[string][]string)
	m.agentColorMap = make(map[string]int)
	m.nextColor = 0
	m.allLines = nil

	// Re-register agents and rebuild lines from stored events
	for i := range m.events {
		evt := &m.events[i]
		if evt.SessionID != m.activeSession {
			continue
		}
		m.registerAgent(evt)
		m.updateAgentActivity(evt.AgentID, evt.parsed.time)
	}

	m.rebuildAllLines()
	m.follow = true
	m.syncViewport()
}

func (m *model) registerAgent(evt *Event) {
	if evt.AgentID == "" {
		return
	}
	if _, ok := m.agentIndex[evt.AgentID]; !ok {
		// Assign a stable color
		if _, hasColor := m.agentColorMap[evt.AgentID]; !hasColor {
			m.agentColorMap[evt.AgentID] = m.nextColor
			m.nextColor++
		}

		// Resolve name: team config > "You are X" in initial message > slug
		name := ""
		if n, ok := m.agentNames[evt.AgentID]; ok {
			name = n
		}
		if name == "" {
			name = resolveAgentName(evt.AgentID, m.agentNames, m.events)
		}

		m.agents = append(m.agents, agentInfo{
			id:       evt.AgentID,
			slug:     evt.Slug,
			name:     name,
			colorIdx: m.agentColorMap[evt.AgentID],
			lastSeen: evt.parsed.time,
		})
		m.rebuildAgentIndex()
	}
}

// updateAgentActivity updates lastSeen and re-sorts tabs by most recent activity.
func (m *model) updateAgentActivity(agentID string, t time.Time) {
	idx, ok := m.agentIndex[agentID]
	if !ok {
		return
	}
	if t.After(m.agents[idx].lastSeen) {
		m.agents[idx].lastSeen = t
	}
	m.sortAgents()
}

// sortAgents sorts agents by lastSeen descending (most recent first).
func (m *model) sortAgents() {
	sort.SliceStable(m.agents, func(i, j int) bool {
		return m.agents[i].lastSeen.After(m.agents[j].lastSeen)
	})
	m.rebuildAgentIndex()
}

// rebuildAgentIndex rebuilds the agentID -> slice index map after sorting.
func (m *model) rebuildAgentIndex() {
	m.agentIndex = make(map[string]int, len(m.agents))
	for i, a := range m.agents {
		m.agentIndex[a.id] = i
	}
}

// activeTabIndex returns the current tab index (0 = All, 1+ = agent position).
func (m *model) activeTabIndex() int {
	if m.activeTabID == "" {
		return 0
	}
	if idx, ok := m.agentIndex[m.activeTabID]; ok {
		return idx + 1
	}
	return 0
}

func (m *model) handleNewEvent(evt Event) {
	m.events = append(m.events, evt)

	// Register agent if it belongs to active session
	if evt.SessionID == m.activeSession {
		m.registerAgent(&evt)
		m.updateAgentActivity(evt.AgentID, evt.parsed.time)

		// Render this event
		colorIdx := 0
		if idx, ok := m.agentIndex[evt.AgentID]; ok {
			colorIdx = m.agents[idx].colorIdx
		}

		lines := renderEvent(&evt, colorIdx, m.showTokens, m.width)

		// Add to per-agent and all lines
		for _, line := range lines {
			m.allLines = append(m.allLines, line)
			if evt.AgentID != "" {
				m.agentLines[evt.AgentID] = append(m.agentLines[evt.AgentID], line)
			}
		}

		m.syncViewport()
	}

	// Update session metadata if this is a new session we haven't seen
	knownSession := false
	for _, s := range m.sessions {
		if s.ID == evt.SessionID {
			knownSession = true
			break
		}
	}
	if !knownSession && evt.SessionID != "" {
		m.sessions = append([]SessionInfo{{
			ID:         evt.SessionID,
			AgentCount: 1,
			EventCount: 1,
			FirstSeen:  evt.parsed.time,
			LastSeen:   evt.parsed.time,
			Agents:     map[string]string{evt.AgentID: evt.Slug},
		}}, m.sessions...)
		// Adjust sessionIdx since we prepended
		m.sessionIdx++
	}
}

// rebuildAllLines re-renders all events for the current session.
func (m *model) rebuildAllLines() {
	m.allLines = nil
	m.agentLines = make(map[string][]string)

	for i := range m.events {
		evt := &m.events[i]
		if evt.SessionID != m.activeSession {
			continue
		}
		colorIdx := 0
		if idx, ok := m.agentIndex[evt.AgentID]; ok {
			colorIdx = m.agents[idx].colorIdx
		}
		lines := renderEvent(evt, colorIdx, m.showTokens, m.width)
		for _, line := range lines {
			m.allLines = append(m.allLines, line)
			if evt.AgentID != "" {
				m.agentLines[evt.AgentID] = append(m.agentLines[evt.AgentID], line)
			}
		}
	}
}

// syncViewport updates the viewport content from the current active tab.
func (m *model) syncViewport() {
	if !m.ready {
		return
	}

	var lines []string
	if m.activeTabID == "" {
		lines = m.allLines
	} else {
		lines = m.agentLines[m.activeTabID]
	}

	content := strings.Join(lines, "\n")
	m.vp.SetContent(content)

	if m.follow {
		m.vp.GotoBottom()
	}
}

func (m model) View() string {
	if !m.ready {
		return "Loading..."
	}

	// Session bar
	sessionBar := m.renderSessionBar()

	// Tab bar
	tabs := m.renderTabs()

	// Viewport
	body := m.vp.View()

	// Status bar
	status := m.renderStatusBar()

	return lipgloss.JoinVertical(lipgloss.Left, sessionBar, tabs, body, status)
}

func (m model) renderSessionBar() string {
	cur := m.sessions[m.sessionIdx]

	label := fmt.Sprintf(" session %s  (%d agents, %s – %s)",
		shortSessionID(cur.ID),
		cur.AgentCount,
		cur.FirstSeen.Format("15:04:05"),
		cur.LastSeen.Format("15:04:05"),
	)

	nav := ""
	if len(m.sessions) > 1 {
		nav = fmt.Sprintf("  %d/%d ", m.sessionIdx+1, len(m.sessions))
	}

	labelWidth := lipgloss.Width(label)
	navWidth := lipgloss.Width(nav)
	gap := m.width - labelWidth - navWidth
	if gap < 0 {
		gap = 0
	}

	return sessionBarStyle.Width(m.width).Render(
		label + strings.Repeat(" ", gap) + nav)
}

func (m model) renderTabs() string {
	var tabs []string

	// "All" tab
	if m.activeTabID == "" {
		tabs = append(tabs, activeTabStyle.Render("All"))
	} else {
		tabs = append(tabs, inactiveTabStyle.Render("All"))
	}

	// Agent tabs (sorted by most recent activity)
	for _, agent := range m.agents {
		label := m.agentLabel(agent)

		isActive := m.activeTabID == agent.id
		if isActive {
			style := lipgloss.NewStyle().
				Bold(true).
				Foreground(lipgloss.Color("#282C34")).
				Background(agentColors[agent.colorIdx%len(agentColors)]).
				Padding(0, 1)
			tabs = append(tabs, style.Render(label))
		} else {
			style := lipgloss.NewStyle().
				Foreground(agentColors[agent.colorIdx%len(agentColors)]).
				Padding(0, 1)
			tabs = append(tabs, style.Render(label))
		}
	}

	tabBar := lipgloss.JoinHorizontal(lipgloss.Top, tabs...)

	// Divider
	divider := dimStyle.Render(strings.Repeat("─", m.width))

	return tabBar + "\n" + divider
}

// agentLabel returns the display label for an agent tab.
func (m model) agentLabel(agent agentInfo) string {
	shortID := agent.id[:min(7, len(agent.id))]
	if agent.name != "" {
		return agent.name
	}
	if agent.slug != "" {
		parts := strings.SplitN(agent.slug, "-", 4)
		if len(parts) >= 3 {
			return shortID + " " + parts[0] + "-" + parts[1]
		}
	}
	return shortID
}

func (m model) renderStatusBar() string {
	var followStr string
	if m.follow {
		followStr = followOnStyle.Render("● FOLLOW")
	} else {
		followStr = followOffStyle.Render("○ PAUSED")
	}

	var tokensStr string
	if m.showTokens {
		tokensStr = "tok:on"
	} else {
		tokensStr = "tok:off"
	}

	eventCount := len(m.allLines)
	agentCount := len(m.agents)

	info := fmt.Sprintf(" %d agents │ %d lines │ %s │ %s",
		agentCount, eventCount, followStr, tokensStr)

	keys := statusKeyStyle.Render("←→") + " tabs " +
		statusKeyStyle.Render("s/S") + " session " +
		statusKeyStyle.Render("f") + " follow " +
		statusKeyStyle.Render("t") + " tokens " +
		statusKeyStyle.Render("q") + " quit "

	infoWidth := lipgloss.Width(info)
	keysWidth := lipgloss.Width(keys)
	gap := m.width - infoWidth - keysWidth
	if gap < 0 {
		gap = 0
	}

	return statusBarStyle.Width(m.width).Render(
		info + strings.Repeat(" ", gap) + keys)
}