package main

import (
	"context"
	"errors"
	"log"
	"math/rand"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/charmbracelet/ssh"
	"github.com/charmbracelet/wish"
	"github.com/charmbracelet/wish/activeterm"
	"github.com/charmbracelet/wish/bubbletea"
	"github.com/charmbracelet/wish/logging"
	"github.com/muesli/termenv"
)

const (
	host = "0.0.0.0"
	port = "2222"
)

func init() {
	rand.Seed(time.Now().UnixNano())
}

func main() {
	s, err := wish.NewServer(
		wish.WithAddress(net.JoinHostPort(host, port)),
		wish.WithHostKeyPath(".ssh/id_ed25519"),
		wish.WithMiddleware(
			bubbletea.Middleware(teaHandler),
			activeterm.Middleware(),
			logging.Middleware(),
		),
	)
	if err != nil {
		log.Fatalln(err)
	}

	done := make(chan os.Signal, 1)
	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)

	log.Printf("🔥 Yule Log SSH Server starting on %s:%s", host, port)
	log.Printf("Connect with: ssh localhost -p %s", port)

	go func() {
		if err = s.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
			log.Fatalln(err)
		}
	}()

	<-done
	log.Println("Shutting down server...")
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()
	if err := s.Shutdown(ctx); err != nil {
		log.Fatalln(err)
	}
}

func teaHandler(s ssh.Session) (tea.Model, []tea.ProgramOption) {
	pty, _, _ := s.Pty()
	renderer := bubbletea.MakeRenderer(s)
	renderer.SetColorProfile(termenv.ANSI256)
	m := newModel(pty.Window.Width, pty.Window.Height, renderer)
	return m, []tea.ProgramOption{tea.WithAltScreen()}
}

type tickMsg time.Time

type model struct {
	width    int
	height   int
	viewport viewport.Model
	buffer   []int
	renderer *lipgloss.Renderer
}

func newModel(width, height int, renderer *lipgloss.Renderer) model {
	vp := viewport.New(width, height)
	size := width * height
	return model{
		width:    width,
		height:   height,
		viewport: vp,
		buffer:   make([]int, size+width+1),
		renderer: renderer,
	}
}

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

func tick() tea.Cmd {
	return tea.Tick(30*time.Millisecond, func(t time.Time) tea.Msg {
		return tickMsg(t)
	})
}

func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	switch msg := msg.(type) {
	case tea.KeyMsg:
		switch msg.String() {
		case "q", "ctrl+c", "esc":
			return m, tea.Quit
		}
	case tea.WindowSizeMsg:
		m.width = msg.Width
		m.height = msg.Height
		m.viewport.Width = msg.Width
		m.viewport.Height = msg.Height
		size := msg.Width * msg.Height
		m.buffer = make([]int, size+msg.Width+1)
	case tickMsg:
		m.updateFire()
		return m, tick()
	}
	return m, nil
}

func (m *model) updateFire() {
	width := m.width
	height := m.height - 1
	if width <= 0 || height <= 0 {
		return
	}

	size := width * height
	if len(m.buffer) < size+width+1 {
		m.buffer = make([]int, size+width+1)
	}

	// Inject heat on bottom row
	for i := 0; i < width/9; i++ {
		idx := rand.Intn(width) + width*(height-1)
		if idx >= 0 && idx < len(m.buffer) {
			m.buffer[idx] = 65
		}
	}

	// Propagate and cool
	for i := 0; i < size; i++ {
		if i+width+1 >= len(m.buffer) {
			continue
		}
		b0 := m.buffer[i]
		b1 := m.buffer[i+1]
		b2 := m.buffer[i+width]
		b3 := m.buffer[i+width+1]
		v := (b0 + b1 + b2 + b3) / 4
		m.buffer[i] = v
	}
}

func (m model) View() string {
	return m.renderFire()
}

var fireColors = []string{"0", "52", "88", "124", "160", "196", "202", "208", "214", "220", "226", "227", "228", "229", "230"}
var fireChars = []rune{' ', '.', ':', '^', '*', 'x', 's', 'S', '#', '$'}

func getFireColorIndex(v int) int {
	if v <= 0 {
		return 0
	} else if v <= 2 {
		return 1
	} else if v <= 4 {
		return 2
	} else if v <= 6 {
		return 3
	} else if v <= 8 {
		return 4
	} else if v <= 10 {
		return 5
	} else if v <= 13 {
		return 6
	} else if v <= 16 {
		return 7
	} else if v <= 20 {
		return 8
	} else if v <= 25 {
		return 9
	} else if v <= 30 {
		return 10
	} else if v <= 40 {
		return 11
	} else if v <= 50 {
		return 12
	} else if v <= 60 {
		return 13
	}
	return 14
}

func (m model) renderFire() string {
	width := m.width
	height := m.height - 1

	if width < 10 || height < 5 {
		return "Terminal too small"
	}

	var output string

	for row := 0; row < height; row++ {
		line := ""
		for col := 0; col < width; col++ {
			i := col + row*width
			v := 0
			if i < len(m.buffer) {
				v = m.buffer[i]
			}

			colorIdx := getFireColorIndex(v)
			style := m.renderer.NewStyle().Foreground(lipgloss.Color(fireColors[colorIdx]))

			chIdx := v
			if chIdx > 9 {
				chIdx = 9
			}
			if chIdx < 0 {
				chIdx = 0
			}

			line += style.Render(string(fireChars[chIdx]))
		}
		output += line + "\n"
	}

	footer := "Press 'q' to exit"
	footerStyle := m.renderer.NewStyle().
		Foreground(lipgloss.Color("240")).
		Width(width).
		Align(lipgloss.Center)
	output += footerStyle.Render(footer)

	return output
}