package ui

import (
	"fmt"
	"strings"
	"time"

	"github.com/Armur-Ai/Pentest-Swarm-AI/internal/pipeline"
	"github.com/charmbracelet/bubbles/spinner"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
)

// Styles
var (
	titleStyle = lipgloss.NewStyle().
			Bold(true).
			Foreground(lipgloss.Color("#60A5FA")).
			Padding(0, 1)

	dimStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("#555555"))

	agentActiveStyle = lipgloss.NewStyle().
				Border(lipgloss.RoundedBorder()).
				BorderForeground(lipgloss.Color("#60A5FA")).
				Padding(0, 1).
				Width(40)

	agentIdleStyle = lipgloss.NewStyle().
			Border(lipgloss.RoundedBorder()).
			BorderForeground(lipgloss.Color("#333333")).
			Padding(0, 1).
			Width(40)

	findingCritical = lipgloss.NewStyle().Foreground(lipgloss.Color("#EF4444")).Bold(true)
	findingHigh     = lipgloss.NewStyle().Foreground(lipgloss.Color("#F97316")).Bold(true)
	findingMedium   = lipgloss.NewStyle().Foreground(lipgloss.Color("#EAB308"))
	findingLow      = lipgloss.NewStyle().Foreground(lipgloss.Color("#22C55E"))

	phaseActive  = lipgloss.NewStyle().Background(lipgloss.Color("#60A5FA")).Foreground(lipgloss.Color("#000")).Padding(0, 1)
	phaseDone    = lipgloss.NewStyle().Background(lipgloss.Color("#22C55E")).Foreground(lipgloss.Color("#000")).Padding(0, 1)
	phasePending = lipgloss.NewStyle().Foreground(lipgloss.Color("#555555")).Padding(0, 1)

	footerStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("#555555")).
			Padding(0, 1)
)

// EventMsg delivers a campaign event to the TUI.
type EventMsg pipeline.CampaignEvent

// TickMsg triggers periodic UI updates.
type TickMsg time.Time

// Model is the bubbletea model for the campaign watch TUI.
type Model struct {
	// Campaign info
	campaignID string
	target     string
	objective  string
	startTime  time.Time

	// Agent status
	agents map[string]AgentStatus

	// Events log
	events   []pipeline.CampaignEvent
	viewport viewport.Model

	// Findings
	findings    []FindingDisplay
	severityMap map[pipeline.Severity]int

	// Phase tracking
	currentPhase string
	phases       []PhaseInfo

	// UI
	spinner  spinner.Model
	width    int
	height   int
	quitting bool
}

// AgentStatus tracks an agent's display state.
type AgentStatus struct {
	Name   string
	Status string // idle, active, complete, error
	Detail string
}

// FindingDisplay is a finding formatted for display.
type FindingDisplay struct {
	Severity pipeline.Severity
	Title    string
	Target   string
}

// PhaseInfo tracks phase progress.
type PhaseInfo struct {
	Name   string
	Status string // pending, active, done
}

// NewModel creates the TUI model.
func NewModel(campaignID, target, objective string) Model {
	s := spinner.New()
	s.Spinner = spinner.Dot
	s.Style = lipgloss.NewStyle().Foreground(lipgloss.Color("#60A5FA"))

	vp := viewport.New(80, 20)

	return Model{
		campaignID: campaignID,
		target:     target,
		objective:  objective,
		startTime:  time.Now(),
		agents: map[string]AgentStatus{
			"orchestrator": {Name: "Orchestrator", Status: "active", Detail: "Initializing..."},
			"recon":        {Name: "Recon Agent", Status: "idle", Detail: "Waiting"},
			"classifier":   {Name: "Classifier", Status: "idle", Detail: "Waiting"},
			"exploit":      {Name: "Exploit Agent", Status: "idle", Detail: "Waiting"},
			"report":       {Name: "Report Agent", Status: "idle", Detail: "Waiting"},
		},
		severityMap: make(map[pipeline.Severity]int),
		phases: []PhaseInfo{
			{Name: "Recon", Status: "pending"},
			{Name: "Classify", Status: "pending"},
			{Name: "Plan", Status: "pending"},
			{Name: "Execute", Status: "pending"},
			{Name: "Report", Status: "pending"},
		},
		currentPhase: "initializing",
		spinner:      s,
		viewport:     vp,
	}
}

func (m Model) Init() tea.Cmd {
	return tea.Batch(
		m.spinner.Tick,
		tickCmd(),
	)
}

func tickCmd() tea.Cmd {
	return tea.Tick(time.Second, func(t time.Time) tea.Msg {
		return TickMsg(t)
	})
}

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":
			m.quitting = true
			return m, tea.Quit
		case "s":
			// Emergency stop
			m.quitting = true
			return m, tea.Quit
		}

	case tea.WindowSizeMsg:
		m.width = msg.Width
		m.height = msg.Height
		m.viewport.Width = msg.Width - 4
		m.viewport.Height = msg.Height - 20

	case EventMsg:
		event := pipeline.CampaignEvent(msg)
		m.events = append(m.events, event)
		m.handleEvent(event)
		m.updateViewport()

	case TickMsg:
		cmds = append(cmds, tickCmd())

	case spinner.TickMsg:
		var cmd tea.Cmd
		m.spinner, cmd = m.spinner.Update(msg)
		cmds = append(cmds, cmd)
	}

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

func (m *Model) handleEvent(event pipeline.CampaignEvent) {
	// Update agent status based on event
	switch event.EventType {
	case pipeline.EventStateChange:
		if strings.Contains(event.Detail, "recon") {
			m.setPhase("Recon")
			m.agents["recon"] = AgentStatus{Name: "Recon Agent", Status: "active", Detail: "Scanning..."}
		} else if strings.Contains(event.Detail, "classif") {
			m.setPhase("Classify")
			m.agents["recon"] = AgentStatus{Name: "Recon Agent", Status: "complete", Detail: "Done"}
			m.agents["classifier"] = AgentStatus{Name: "Classifier", Status: "active", Detail: "Classifying..."}
		} else if strings.Contains(event.Detail, "plan") {
			m.setPhase("Plan")
			m.agents["classifier"] = AgentStatus{Name: "Classifier", Status: "complete", Detail: "Done"}
			m.agents["exploit"] = AgentStatus{Name: "Exploit Agent", Status: "active", Detail: "Building chains..."}
		} else if strings.Contains(event.Detail, "execut") {
			m.setPhase("Execute")
		} else if strings.Contains(event.Detail, "report") {
			m.setPhase("Report")
			m.agents["exploit"] = AgentStatus{Name: "Exploit Agent", Status: "complete", Detail: "Done"}
			m.agents["report"] = AgentStatus{Name: "Report Agent", Status: "active", Detail: "Writing report..."}
		} else if strings.Contains(event.Detail, "complete") {
			m.agents["report"] = AgentStatus{Name: "Report Agent", Status: "complete", Detail: "Done"}
		}

	case pipeline.EventFindingDiscovered:
		// Parse severity from detail
		detail := event.Detail
		var sev pipeline.Severity = pipeline.SeverityMedium
		if strings.Contains(detail, "CRITICAL") {
			sev = pipeline.SeverityCritical
		} else if strings.Contains(detail, "HIGH") {
			sev = pipeline.SeverityHigh
		} else if strings.Contains(detail, "LOW") {
			sev = pipeline.SeverityLow
		}
		m.severityMap[sev]++
		m.findings = append(m.findings, FindingDisplay{Severity: sev, Title: detail})

	case pipeline.EventThought:
		if event.AgentName != "" {
			if a, ok := m.agents[event.AgentName]; ok {
				a.Detail = truncateStr(event.Detail, 50)
				m.agents[event.AgentName] = a
			}
		}

	case pipeline.EventToolCall:
		if event.AgentName != "" {
			if a, ok := m.agents[event.AgentName]; ok {
				a.Status = "active"
				a.Detail = truncateStr(event.Detail, 50)
				m.agents[event.AgentName] = a
			}
		}
	}
}

func (m *Model) setPhase(name string) {
	for i := range m.phases {
		if m.phases[i].Name == name {
			m.phases[i].Status = "active"
			m.currentPhase = name
		} else if m.phases[i].Status == "active" {
			m.phases[i].Status = "done"
		}
	}
}

func (m *Model) updateViewport() {
	var lines []string
	for _, e := range m.events {
		ts := e.Timestamp.Format("15:04:05")
		prefix := dimStyle.Render(ts)

		switch e.EventType {
		case pipeline.EventThought:
			lines = append(lines, fmt.Sprintf("%s [think] %s", prefix, e.Detail))
		case pipeline.EventToolCall:
			lines = append(lines, fmt.Sprintf("%s   [>>]  %s", prefix, e.Detail))
		case pipeline.EventToolResult:
			lines = append(lines, fmt.Sprintf("%s   [<<]  %s", prefix, e.Detail))
		case pipeline.EventFindingDiscovered:
			lines = append(lines, fmt.Sprintf("%s    [!]  %s", prefix, e.Detail))
		case pipeline.EventError:
			lines = append(lines, fmt.Sprintf("%s  [ERR]  %s", prefix, e.Detail))
		case pipeline.EventMilestone:
			lines = append(lines, fmt.Sprintf("%s [DONE]  %s", prefix, e.Detail))
		default:
			lines = append(lines, fmt.Sprintf("%s   [%s] %s", prefix, e.EventType, e.Detail))
		}
	}
	m.viewport.SetContent(strings.Join(lines, "\n"))
	m.viewport.GotoBottom()
}

func (m Model) View() string {
	if m.quitting {
		return ""
	}

	var b strings.Builder

	// Header
	elapsed := time.Since(m.startTime).Round(time.Second)
	header := fmt.Sprintf(" %s  %s  %s  %s",
		titleStyle.Render("PENTEST SWARM"),
		lipgloss.NewStyle().Bold(true).Render(m.target),
		dimStyle.Render(m.objective),
		dimStyle.Render(fmt.Sprintf("%s", elapsed)),
	)
	b.WriteString(header + "\n")

	// Phase bar
	var phases []string
	for _, p := range m.phases {
		switch p.Status {
		case "done":
			phases = append(phases, phaseDone.Render(p.Name+" ✓"))
		case "active":
			phases = append(phases, phaseActive.Render(m.spinner.View()+" "+p.Name))
		default:
			phases = append(phases, phasePending.Render(p.Name))
		}
	}
	b.WriteString(" " + strings.Join(phases, " → ") + "\n")
	b.WriteString(dimStyle.Render(strings.Repeat("─", maxInt(m.width, 60))) + "\n")

	// Two columns: agents (left) + findings (right)
	agentCol := m.renderAgents()
	findingCol := m.renderFindings()

	// Simple side-by-side
	agentLines := strings.Split(agentCol, "\n")
	findingLines := strings.Split(findingCol, "\n")
	maxLines := maxInt(len(agentLines), len(findingLines))

	for i := 0; i < maxLines; i++ {
		left := ""
		right := ""
		if i < len(agentLines) {
			left = agentLines[i]
		}
		if i < len(findingLines) {
			right = findingLines[i]
		}
		b.WriteString(fmt.Sprintf("%-45s %s\n", left, right))
	}

	b.WriteString(dimStyle.Render(strings.Repeat("─", maxInt(m.width, 60))) + "\n")

	// Event log
	b.WriteString(dimStyle.Render(" Event Log") + "\n")
	// Show last 8 events
	start := 0
	if len(m.events) > 8 {
		start = len(m.events) - 8
	}
	for _, e := range m.events[start:] {
		ts := dimStyle.Render(e.Timestamp.Format("15:04:05"))
		b.WriteString(fmt.Sprintf(" %s %s\n", ts, truncateStr(e.Detail, 70)))
	}

	// Footer
	b.WriteString("\n")
	b.WriteString(footerStyle.Render(" q:quit  s:stop  ↑↓:scroll"))

	return b.String()
}

func (m Model) renderAgents() string {
	var b strings.Builder
	b.WriteString(dimStyle.Render(" Agents") + "\n")

	order := []string{"orchestrator", "recon", "classifier", "exploit", "report"}
	for _, name := range order {
		a := m.agents[name]
		style := agentIdleStyle
		statusIcon := "○"

		switch a.Status {
		case "active":
			style = agentActiveStyle
			statusIcon = "●"
		case "complete":
			statusIcon = "✓"
		case "error":
			statusIcon = "✗"
		}

		content := fmt.Sprintf(" %s %s\n %s", statusIcon, a.Name, dimStyle.Render(truncateStr(a.Detail, 35)))
		b.WriteString(style.Render(content) + "\n")
	}

	return b.String()
}

func (m Model) renderFindings() string {
	var b strings.Builder
	b.WriteString(dimStyle.Render(" Findings") + "\n")

	// Severity bars
	c := m.severityMap[pipeline.SeverityCritical]
	h := m.severityMap[pipeline.SeverityHigh]
	med := m.severityMap[pipeline.SeverityMedium]
	l := m.severityMap[pipeline.SeverityLow]

	b.WriteString(fmt.Sprintf(" %s %d  %s %d  %s %d  %s %d\n",
		findingCritical.Render("CRIT"), c,
		findingHigh.Render("HIGH"), h,
		findingMedium.Render("MED"), med,
		findingLow.Render("LOW"), l,
	))
	b.WriteString("\n")

	// Last 5 findings
	start := 0
	if len(m.findings) > 5 {
		start = len(m.findings) - 5
	}
	for _, f := range m.findings[start:] {
		var style lipgloss.Style
		switch f.Severity {
		case pipeline.SeverityCritical:
			style = findingCritical
		case pipeline.SeverityHigh:
			style = findingHigh
		case pipeline.SeverityMedium:
			style = findingMedium
		default:
			style = findingLow
		}
		b.WriteString(" " + style.Render("●") + " " + truncateStr(f.Title, 30) + "\n")
	}

	if len(m.findings) == 0 {
		b.WriteString(dimStyle.Render(" No findings yet") + "\n")
	}

	return b.String()
}

func truncateStr(s string, max int) string {
	if len(s) <= max {
		return s
	}
	return s[:max-3] + "..."
}

func maxInt(a, b int) int {
	if a > b {
		return a
	}
	return b
}
