package main

import (
	"context"
	"errors"
	"fmt"
	"net"
	"os"
	"os/signal"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"time"

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

type programInfo struct {
	roomName string
	userID   string
}

// app contains a wish server and manages chat rooms and programs.
type app struct {
	*ssh.Server
	activePrograms map[*tea.Program]programInfo // program -> programInfo
	rooms          map[string][]*tea.Program    // roomName -> list of programs in that room
	mu             sync.Mutex
}

// BroadcastToRoom sends a message to all running programs in a specific room.
func (a *app) BroadcastToRoom(roomName string, msg tea.Msg) {
	a.mu.Lock()
	defer a.mu.Unlock()

	if progsInRoom, ok := a.rooms[roomName]; ok {
		log.Debugf("Broadcasting to room %s (%d programs): %+v", roomName, len(progsInRoom), msg)
		for _, p := range progsInRoom {
			go p.Send(msg)
		}
	} else {
		log.Warnf("Attempted to broadcast to non-existent or empty room: %s", roomName)
	}
}

func newApp() *app {
	a := &app{
		activePrograms: make(map[*tea.Program]programInfo), // Initialize correctly
		rooms:          make(map[string][]*tea.Program),
	}
	s, err := wish.NewServer(
		wish.WithAddress(net.JoinHostPort(host, port)),
		wish.WithHostKeyPath(".ssh/id_ed25519"), // Ensure this key exists (e.g., run ssh-keygen -t ed25519 -f .ssh/id_ed25519 -N "")
		wish.WithMiddleware(
			bubbletea.MiddlewareWithProgramHandler(a.ProgramHandler, termenv.ANSI256),
			activeterm.Middleware(),
			logging.Middleware(),
		),
	)
	if err != nil {
		log.Fatal("Could not initialize server", "error", err)
	}

	a.Server = s
	return a
}

func (a *app) Start() {
	var err error
	done := make(chan os.Signal, 1)
	signal.Notify(done, os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
	log.Info("Starting SSH server", "host", host, "port", port)
	go func() {
		if err = a.ListenAndServe(); err != nil && !errors.Is(err, ssh.ErrServerClosed) {
			log.Error("Could not start server", "error", err)
			done <- syscall.SIGTERM // Signal main to exit if server fails
		}
	}()

	<-done
	log.Info("Stopping SSH server...")
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer func() {
		cancel()
		// Clean up internal state after server shutdown
		a.mu.Lock()
		a.activePrograms = make(map[*tea.Program]programInfo) // Correct type
		a.rooms = make(map[string][]*tea.Program)
		a.mu.Unlock()
		log.Info("Cleared app program stores.")
	}()

	if err := a.Shutdown(ctx); err != nil {
		log.Error("Could not gracefully stop server", "error", err)
	}
	log.Info("SSH server stopped.")
}

func (a *app) ProgramHandler(s ssh.Session) *tea.Program {
	userID := s.User()
	m := initialModel(a, userID) // initialModel returns *model

	p := tea.NewProgram(m, bubbletea.MakeOptions(s)...)

	m.SetProgram(p)         // Give model a reference to its own program
	a.addProgram(p, userID) // Register program with the app, passing userID

	return p
}

func (a *app) addProgram(p *tea.Program, userID string) {
	a.mu.Lock()
	defer a.mu.Unlock()
	a.activePrograms[p] = programInfo{roomName: "", userID: userID} // Store userID
	log.Debugf("Program %p added for user %s (unassigned)", p, userID)
}

func (a *app) AssignProgramToRoom(prog *tea.Program, roomName string) {
	a.mu.Lock()
	defer a.mu.Unlock()

	currentInfo, exists := a.activePrograms[prog]
	if !exists {
		// userID is NOT available here if the program isn't in activePrograms.
		// NO prog.Model() call should be here.
		log.Warnf("Program %p not found in activePrograms, cannot assign to room %s", prog, roomName)
		return
	}
	// Correct: UserID is obtained from the stored currentInfo.
	userID := currentInfo.userID
	currentRoom := currentInfo.roomName

	if currentRoom != "" && currentRoom != roomName {
		var updatedRoomProgs []*tea.Program
		for _, pInRoom := range a.rooms[currentRoom] {
			if pInRoom != prog {
				updatedRoomProgs = append(updatedRoomProgs, pInRoom)
			}
		}
		if len(updatedRoomProgs) > 0 {
			a.rooms[currentRoom] = updatedRoomProgs
		} else {
			delete(a.rooms, currentRoom)
		}
		// Correct: userID is used from the variable. NO prog.Model() here.
		log.Debugf("Program %p (user: %s) moved from room %s", prog, userID, currentRoom)
	}

	if _, ok := a.rooms[roomName]; !ok {
		a.rooms[roomName] = []*tea.Program{}
	}

	foundInNewRoomList := false
	for _, pInList := range a.rooms[roomName] {
		if pInList == prog {
			foundInNewRoomList = true
			break
		}
	}
	if !foundInNewRoomList {
		a.rooms[roomName] = append(a.rooms[roomName], prog)
	}

	a.activePrograms[prog] = programInfo{roomName: roomName, userID: userID}
	// Correct: userID is used from the variable. NO prog.Model() here.
	// If your error is on this line (or similar), ensure 'userID' is not derived from prog.Model().
	log.Infof("Program %p assigned/confirmed in room %s (user: %s)", prog, roomName, userID)
}

func (a *app) RemoveProgram(prog *tea.Program) {
	a.mu.Lock()
	defer a.mu.Unlock()

	info, exists := a.activePrograms[prog]
	if !exists {
		// Correct: NO prog.Model() call here.
		log.Warnf("Attempted to remove unknown or already removed program: %p", prog)
		return
	}
	// Correct: UserID and roomName are obtained from the stored info.
	userID := info.userID
	roomName := info.roomName

	if roomName != "" {
		var updatedRoomProgs []*tea.Program
		for _, pInRoom := range a.rooms[roomName] {
			if pInRoom != prog {
				updatedRoomProgs = append(updatedRoomProgs, pInRoom)
			}
		}
		if len(updatedRoomProgs) > 0 {
			a.rooms[roomName] = updatedRoomProgs
		} else {
			delete(a.rooms, roomName)
		}
		// Correct: userID is used from the variable. NO prog.Model() here.
		// If your error is on this line, ensure 'userID' is not derived from prog.Model().
		log.Infof("Program %p removed from room %s (user: %s)", prog, roomName, userID)
	} else {
		// Correct: userID is used from the variable. NO prog.Model() here.
		log.Debugf("Program %p was not in any room, removing from active list (user: %s).", prog, userID)
	}

	delete(a.activePrograms, prog)
	// Correct: userID is used from the variable. NO prog.Model() here.
	log.Infof("Program %p removed from active list (user: %s).", prog, userID)
}

func main() {
	// For more detailed logs:
	// file, err := os.OpenFile("chat_server.log", os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0666)
	// if err == nil {
	// 	log.Default().SetOutput(file)
	// }
	// log.SetLevel(log.DebugLevel) // Show debug messages

	app := newApp()
	app.Start()
}

type (
	errMsg  error
	chatMsg struct {
		id   string
		text string
		room string // So message carries its intended room
	}
)

type modelPhase int

const (
	phaseRoomInput modelPhase = iota
	phaseChat
)

var roomNameRegex = regexp.MustCompile("^[a-zA-Z0-9]{1,24}$")

func isValidRoomName(name string) bool {
	return roomNameRegex.MatchString(name)
}

type model struct {
	app         *app         // Pointer to the main application state
	program     *tea.Program // Pointer to this model's own tea.Program
	phase       modelPhase
	roomInput   textinput.Model
	currentRoom string
	id          string // User ID from SSH session

	// Chat fields
	viewport    viewport.Model
	messages    []string
	textarea    textarea.Model
	senderStyle lipgloss.Style
	err         error
}

func initialModel(appInstance *app, userID string) *model {
	roomTi := textinput.New()
	roomTi.Placeholder = "e.g., 'dev-chat' or 'gaming'"
	roomTi.Focus()
	roomTi.CharLimit = 24
	roomTi.Width = 40

	ta := textarea.New()
	ta.Placeholder = "Type your message here and press Enter..."
	ta.Prompt = "┃ "
	ta.CharLimit = 280
	ta.SetWidth(60)
	ta.SetHeight(3)
	ta.FocusedStyle.CursorLine = lipgloss.NewStyle()
	ta.ShowLineNumbers = false
	ta.KeyMap.InsertNewline.SetEnabled(false)

	vp := viewport.New(60, 15)
	vp.SetContent("Welcome! Please enter a room name to join or create.")

	return &model{
		app:         appInstance,
		id:          userID,
		phase:       phaseRoomInput,
		roomInput:   roomTi,
		textarea:    ta,
		messages:    []string{},
		viewport:    vp,
		senderStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("12")),
		err:         nil,
	}
}

func (m *model) SetProgram(p *tea.Program) {
	m.program = p
}

func (m *model) Init() tea.Cmd {
	switch m.phase {
	case phaseRoomInput:
		return textinput.Blink
	case phaseChat:
		return textarea.Blink
	default:
		return nil
	}
}

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

	switch m.phase {
	case phaseRoomInput:
		switch msg := msg.(type) {
		case tea.KeyMsg:
			switch msg.Type {
			case tea.KeyCtrlC, tea.KeyEsc:
				if m.program != nil {
					m.app.RemoveProgram(m.program)
				}
				return m, tea.Quit
			case tea.KeyEnter:
				roomName := strings.TrimSpace(m.roomInput.Value())
				if isValidRoomName(roomName) {
					m.currentRoom = roomName
					m.phase = phaseChat
					m.err = nil

					if m.program == nil {
						log.Error("CRITICAL: m.program is nil when trying to assign to room. This should not happen.")
						m.err = errors.New("internal setup error: program reference missing")
						return m, nil
					}
					m.app.AssignProgramToRoom(m.program, m.currentRoom)

					m.textarea.Focus()
					m.roomInput.Blur()

					m.viewport.SetContent(fmt.Sprintf("Welcome to room '%s'! (User: %s)\nMessages will appear below.", m.currentRoom, m.id))
					m.messages = []string{}
					cmds = append(cmds, textarea.Blink)
				} else {
					m.err = errors.New("invalid room name: 1-24 alphanumeric characters required")
					m.roomInput.Focus()
					cmds = append(cmds, textinput.Blink)
				}
				return m, tea.Batch(cmds...)
			}
		case errMsg:
			m.err = msg
			return m, nil
		}

		m.roomInput, cmd = m.roomInput.Update(msg)
		cmds = append(cmds, cmd)

	case phaseChat:
		var tiCmd tea.Cmd
		var vpCmd tea.Cmd

		m.textarea, tiCmd = m.textarea.Update(msg)
		m.viewport, vpCmd = m.viewport.Update(msg)
		cmds = append(cmds, tiCmd, vpCmd)

		switch msg := msg.(type) {
		case tea.KeyMsg:
			switch msg.Type {
			case tea.KeyCtrlC, tea.KeyEsc:
				if m.program != nil {
					m.app.RemoveProgram(m.program)
				}
				return m, tea.Quit
			case tea.KeyEnter:
				val := strings.TrimSpace(m.textarea.Value())
				if val != "" {
					m.app.BroadcastToRoom(m.currentRoom, chatMsg{
						id:   m.id,
						text: val,
						room: m.currentRoom,
					})
					m.textarea.Reset()
					m.viewport.GotoBottom()
				}
			}
		case chatMsg:
			if msg.room == m.currentRoom {
				styledId := m.senderStyle.Render(msg.id)
				if msg.id == m.id {
					styledId = lipgloss.NewStyle().Foreground(lipgloss.Color("10")).Render(msg.id + " (you)")
				}
				m.messages = append(m.messages, styledId+": "+msg.text)
				m.viewport.SetContent(strings.Join(m.messages, "\n"))
				m.viewport.GotoBottom()
			} else {
				log.Warnf("Client in room '%s' (user '%s') incorrectly received msg for room '%s'. Ignored.", m.currentRoom, m.id, msg.room)
			}
		case errMsg:
			m.err = msg
			log.Error("Error in chat model", "user", m.id, "room", m.currentRoom, "error", msg)
			return m, nil
		}
	}
	return m, tea.Batch(cmds...)
}

func (m *model) View() string {
	var viewBuilder strings.Builder

	titleStyle := lipgloss.NewStyle().Bold(true).Foreground(lipgloss.Color("63"))
	errorStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("9")).Bold(true)

	switch m.phase {
	case phaseRoomInput:
		viewBuilder.WriteString(titleStyle.Render("chatter.sh") + "\n\n")
		viewBuilder.WriteString("Please enter a room name to join or create.\n")
		viewBuilder.WriteString("Rules: 1-24 characters, letters (a-z, A-Z) and numbers (0-9) only.\n\n")
		viewBuilder.WriteString(m.roomInput.View())
		viewBuilder.WriteString("\n\n(Ctrl+C or Esc to quit)")
	case phaseChat:
		header := fmt.Sprintf("Room: %s / User: %s", m.currentRoom, m.id)
		viewBuilder.WriteString(titleStyle.Render(header) + "\n")
		viewBuilder.WriteString(lipgloss.NewStyle().BorderStyle(lipgloss.NormalBorder()).BorderBottom(true).Width(m.viewport.Width).Render(m.viewport.View()) + "\n\n")
		viewBuilder.WriteString(m.textarea.View())
		viewBuilder.WriteString("\n\n(Ctrl+C or Esc to quit)")
	default:
		viewBuilder.WriteString("Unknown application state. This is a bug.")
	}

	if m.err != nil {
		viewBuilder.WriteString("\n\n" + errorStyle.Render("Error: "+m.err.Error()))
	}

	return viewBuilder.String() + "\n"
}