package internal

import (
	"crypto/rand"
	"database/sql"
	"encoding/base64"
	"fmt"
	"io"
	"strings"
	"time"

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

var (
	msgStyle     = lipgloss.NewStyle().Foreground(lipgloss.Color("252")).PaddingLeft(1)
	senderStyle  = lipgloss.NewStyle().Foreground(lipgloss.Color("205")).Bold(true)
	errStyle     = lipgloss.NewStyle().Foreground(lipgloss.Color("196")).Bold(true)
	sysStyle     = lipgloss.NewStyle().Foreground(lipgloss.Color("117")).Italic(true)
	roomStyle    = lipgloss.NewStyle().Foreground(lipgloss.Color("86")).Bold(true)
	dmStyle      = lipgloss.NewStyle().Foreground(lipgloss.Color("141")).Bold(true)
	noteStyle    = lipgloss.NewStyle().Foreground(lipgloss.Color("228")).Bold(true)
	commandStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("39")).Italic(true)
	timeStyle    = lipgloss.NewStyle().Foreground(lipgloss.Color("246"))
	lockedStyle  = lipgloss.NewStyle().Foreground(lipgloss.Color("240")).Italic(true)
)

type chatMsg struct {
	Timestamp time.Time
	Sender    string
	Content   string
}

type model struct {
	sshSession        ssh.Session
	pubkey            string
	username          string
	authorized        bool
	needsRegistration bool

	state            int
	passphrase       string
	identityKey      *IdentityKeyPair
	encryptedPrivKey []byte
	privKeySalt      []byte
	privKeyNonce     []byte

	currentRoomID    string
	currentRoomEpoch int
	roomKeyCache     map[int][]byte

	rooms []struct {
		ID      string
		Name    string
		Creator string
		IsDM    bool
	}

	err error

	input      textinput.Model
	viewport   viewport.Model
	messages   []chatMsg
	updateChan chan bool
}

func initialModel(s ssh.Session) model {
	pk := s.PublicKey()
	var pubkeyStr string
	var auth bool
	var user string
	var needsRegistration bool

	if pk == nil {
		auth = false
		user = ""
		pubkeyStr = ""
	} else {
		pubkeyStr = fmt.Sprintf("%s %s", pk.Type(), base64.StdEncoding.EncodeToString(pk.Marshal()))

		username, err := getUsername(pubkeyStr)
		if err == sql.ErrNoRows {
			needsRegistration = true
			auth = true
			user = ""
		} else {
			auth = true
			needsRegistration = false
			user = username
		}
	}

	ti := textinput.New()
	if needsRegistration {
		ti.Placeholder = "Choose a username..."
	} else {
		ti.Placeholder = "Enter Passphrase..."
		ti.EchoMode = textinput.EchoPassword
	}
	ti.Focus()
	ti.CharLimit = 156
	ti.Width = 50

	vp := viewport.New(80, 20)

	return model{
		sshSession:        s,
		pubkey:            pubkeyStr,
		username:          user,
		authorized:        auth,
		needsRegistration: needsRegistration,
		state:             0,
		input:             ti,
		viewport:          vp,
		updateChan:        updates.subscribe(),
		roomKeyCache:      make(map[int][]byte),
	}
}

func (m model) Init() tea.Cmd {
	return tea.Batch(textinput.Blink, m.waitForUpdate())
}

func (m model) waitForUpdate() tea.Cmd {
	return func() tea.Msg {
		<-m.updateChan
		return struct{}{}
	}
}

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

	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		m.viewport.Width = msg.Width
		m.viewport.Height = msg.Height - 4
		return m, tea.ClearScreen

	case tea.KeyMsg:
		// Clear any displayed errors on user interaction
		if m.err != nil {
			m.err = nil
		}

		switch msg.Type {
		case tea.KeyCtrlC, tea.KeyEsc:
			return m, tea.Quit
		case tea.KeyEnter:
			if !m.authorized {
				return m, tea.Quit
			}

			if m.needsRegistration {
				username := m.input.Value()
				if len(username) > 0 {
					identityKey, err := generateIdentityKeyPair()
					if err != nil {
						m.err = err
						return m, nil
					}

					prekey, prekeySignature, err := generateSignedPrekey()
					if err != nil {
						m.err = err
						return m, nil
					}

					if err := authorizeUser(m.pubkey, username, identityKey.PublicKey[:], prekey[:], prekeySignature); err != nil {
						m.err = err
					} else {
						m.username = username
						m.identityKey = identityKey
						m.needsRegistration = false

						m.input.Reset()
						m.input.Placeholder = "Enter Passphrase..."
						m.input.EchoMode = textinput.EchoPassword
						m.viewport.SetContent("")
						return m, tea.ClearScreen
					}
				}
				return m, nil
			}

			if m.state == 0 {
				pass := m.input.Value()
				if len(pass) > 0 {
					if m.identityKey != nil {
						ciphertext, nonce, salt, err := encryptUserKey(m.identityKey.PrivateKey, pass)
						if err != nil {
							m.err = err
							return m, nil
						}
						m.encryptedPrivKey = ciphertext
						m.privKeyNonce = nonce
						m.privKeySalt = salt
					}

					m.passphrase = pass
					m.state = 1

					m.input.Reset()
					m.input.EchoMode = textinput.EchoNormal
					m.input.Placeholder = "Loading rooms..."
					m.viewport.SetContent("")

					m.loadRooms()
					return m, tea.ClearScreen
				}
			} else if m.state == 1 {
				val := m.input.Value()
				if len(val) > 0 {
					m.handleRoomListInput(val)
					m.input.Reset()
					if m.state == 2 {
						return m, tea.ClearScreen
					}
				}
			} else if m.state == 2 {
				val := m.input.Value()
				if len(val) > 0 {
					if val == "/back" {
						m.exitChat()
						return m, tea.ClearScreen
					}

					if val == "/leave" {
						creator, _ := getRoomCreator(m.currentRoomID)
						if creator == m.username {
							m.saveMessage("System", "Error: Owners cannot leave. Use /delete to destroy room.")
							m.input.Reset()
							return m, nil
						}
						leaveRoom(m.currentRoomID, m.username)
						m.rotateRoomKey(m.currentRoomID)
						m.exitChat()
						return m, tea.ClearScreen
					}

					if val == "/delete" {
						creator, _ := getRoomCreator(m.currentRoomID)
						if creator == m.username {
							deleteRoom(m.currentRoomID)
							m.exitChat()
							return m, tea.ClearScreen
						} else {
							m.saveMessage("System", "Error: Only the creator can delete this room.")
							m.input.Reset()
							return m, nil
						}
					}
					m.handleChatInput(val)
					m.input.Reset()
				}
			}
		}

	case struct{}:
		if m.state == 2 {
			m.loadMessages()
			return m, m.waitForUpdate()
		}
	}

	m.input, cmd = m.input.Update(msg)
	return m, cmd
}

func (m *model) exitChat() {
	m.state = 1
	m.currentRoomID = ""
	m.currentRoomEpoch = 0
	m.messages = nil
	m.roomKeyCache = make(map[int][]byte)
	m.viewport.SetContent("")
	m.loadRooms()
	m.input.Reset()
}

func (m *model) loadRooms() {
	rooms, err := listUserRooms(m.username)
	if err != nil {
		m.err = err
		return
	}
	m.rooms = rooms
	m.viewport.SetContent("")
	m.input.SetValue("")
	m.input.Placeholder = "Enter room # or: /join <name>, /dm <username>"
}

func (m *model) handleRoomListInput(text string) {
	if strings.HasPrefix(text, "/join ") {
		roomName := strings.TrimPrefix(text, "/join ")
		m.handleJoinRoom(roomName)
		return
	}
	if strings.HasPrefix(text, "/new ") {
		roomName := strings.TrimPrefix(text, "/new ")
		m.handleJoinRoom(roomName)
		return
	}

	if strings.HasPrefix(text, "/dm ") {
		targetUser := strings.TrimPrefix(text, "/dm ")
		m.handleSelectUserForDM(targetUser)
		return
	}

	if idx := parseInt(text); idx > 0 && idx <= len(m.rooms) {
		room := m.rooms[idx-1]
		m.enterRoom(room.ID, room.Name)
	}
}

func (m *model) handleJoinRoom(roomName string) {
	existingID, err := getRoomIDByName(roomName)
	if err == nil && existingID != "" {
		if err := joinRoomMember(existingID, m.username); err != nil {
			m.err = err
			return
		}
		m.rotateRoomKey(existingID)
		m.enterRoom(existingID, roomName)
		return
	}

	roomID := generateRoomID()
	if err := createRoom(roomID, roomName, m.username, false); err != nil {
		m.err = err
		return
	}
	if err := joinRoomMember(roomID, m.username); err != nil {
		m.err = err
		return
	}
	m.rotateRoomKey(roomID)
	m.enterRoom(roomID, roomName)
}

func (m *model) rotateRoomKey(roomID string) {
	newKey := make([]byte, 32)
	io.ReadFull(rand.Reader, newKey)

	newEpoch, err := incrementRoomEpoch(roomID)
	if err != nil {
		m.err = err
		return
	}

	members, err := getRoomMembers(roomID)
	if err != nil {
		m.err = err
		return
	}

	for _, user := range members {
		idKey, err := getMemberIdentityKey(user)
		if err != nil {
			continue
		}
		encrypted, err := EncryptKeyForUser(idKey, newKey)
		if err != nil {
			continue
		}
		storeUserRoomKey(user, roomID, newEpoch, encrypted)
	}
}

func (m *model) handleSelectUserForDM(username string) {
	if username == "/cancel" {
		m.state = 1
		m.input.Placeholder = "Enter room # or: /join <name>, /dm <username>"
		return
	}
	if username == m.username {
		m.createNoteToSelf()
		return
	}

	// CHECK IF USER EXISTS
	identityKey, prekey, prekeySignature, err := getUserKeys(username)
	if err != nil {
		if err == sql.ErrNoRows {
			m.err = fmt.Errorf("user '%s' not found", username)
		} else {
			m.err = err
		}
		return
	}

	if len(identityKey) == 0 {
		m.err = fmt.Errorf("user '%s' has no keys set up", username)
		return
	}

	// Pre-calc shared secret for DM
	var theirIdentityKey, theirPrekey [32]byte
	copy(theirIdentityKey[:], identityKey)
	copy(theirPrekey[:], prekey)

	ephemeralKey, err := generateIdentityKeyPair()
	if err != nil {
		m.err = err
		return
	}

	dh1, _ := performDH(m.identityKey.PrivateKey, theirPrekey)
	dh2, _ := performDH(ephemeralKey.PrivateKey, theirIdentityKey)
	dh3, _ := performDH(ephemeralKey.PrivateKey, theirPrekey)

	_ = deriveSharedSecret(dh1, dh2, dh3)
	_ = prekeySignature

	roomID := generateRoomID()
	roomName := fmt.Sprintf("DM: %s <-> %s", m.username, username)

	if err := createRoom(roomID, roomName, m.username, true); err != nil {
		m.err = err
		return
	}
	joinRoomMember(roomID, m.username)
	joinRoomMember(roomID, username)

	m.rotateRoomKey(roomID)
	m.enterRoom(roomID, roomName)
}

func (m *model) createNoteToSelf() {
	roomID := generateRoomID()
	roomName := "[Note to Self]"
	if err := createRoom(roomID, roomName, m.username, true); err != nil {
		m.err = err
		return
	}
	joinRoomMember(roomID, m.username)
	m.rotateRoomKey(roomID)
	m.enterRoom(roomID, roomName)
}

func (m *model) enterRoom(roomID, roomName string) {
	epoch, _ := getRoomCurrentEpoch(roomID)
	creator, _ := getRoomCreator(roomID)

	m.currentRoomID = roomID
	m.currentRoomEpoch = epoch
	m.state = 2

	if creator == m.username {
		m.input.Placeholder = fmt.Sprintf("[%s] /back to menu, /delete to destroy", roomName)
	} else {
		m.input.Placeholder = fmt.Sprintf("[%s] /back to menu, /leave to quit", roomName)
	}
	m.loadMessages()
}

func (m *model) handleChatInput(text string) {
	m.saveMessage(m.username, text)
}

func parseInt(s string) int {
	var num int
	fmt.Sscanf(s, "%d", &num)
	return num
}

func generateRoomID() string {
	b := make([]byte, 16)
	io.ReadFull(rand.Reader, b)
	return base64.URLEncoding.EncodeToString(b)
}

func (m *model) getEpochKey(epoch int) []byte {
	if key, ok := m.roomKeyCache[epoch]; ok {
		return key
	}
	encKey, err := getUserRoomKey(m.username, m.currentRoomID, epoch)
	if err != nil || encKey == nil {
		return nil
	}
	key, err := DecryptKeyForUser(m.identityKey.PrivateKey, encKey)
	if err != nil {
		return nil
	}
	m.roomKeyCache[epoch] = key
	return key
}

func (m *model) saveMessage(sender, text string) {
	m.currentRoomEpoch, _ = getRoomCurrentEpoch(m.currentRoomID)
	key := m.getEpochKey(m.currentRoomEpoch)
	if key == nil {
		m.err = fmt.Errorf("no key for current epoch")
		return
	}

	ct, nonce, err := encryptMsg(text, key)
	if err != nil {
		m.err = err
		return
	}

	err = saveMessage(m.currentRoomID, sender, m.currentRoomEpoch, ct, nonce)
	if err != nil {
		m.err = err
		return
	}

	go updates.broadcast()
}

func (m *model) loadMessages() {
	if m.currentRoomID == "" {
		return
	}

	dbMessages, err := loadMessages(m.currentRoomID)
	if err != nil {
		m.err = err
		return
	}

	var msgs []chatMsg
	for _, dbMsg := range dbMessages {
		key := m.getEpochKey(dbMsg.Epoch)
		var plain string
		if key == nil {
			plain = "[Unreadable: Old History or Key Rotated]"
		} else {
			p, err := decryptMsg(dbMsg.Ciphertext, dbMsg.Nonce, key)
			if err != nil {
				plain = "[Decryption Failed]"
			} else {
				plain = p
			}
		}

		msgs = append(msgs, chatMsg{
			Timestamp: dbMsg.Timestamp,
			Sender:    dbMsg.Sender,
			Content:   plain,
		})
	}
	m.messages = msgs
	m.updateViewport()
}

func (m *model) updateViewport() {
	var b strings.Builder
	for _, msg := range m.messages {
		dateStr := msg.Timestamp.Format("Jan 2 15:04")

		if strings.Contains(msg.Content, "[Unreadable") {
			b.WriteString(fmt.Sprintf("%s %s\n",
				timeStyle.Render(dateStr),
				lockedStyle.Render(msg.Content),
			))
		} else {
			b.WriteString(fmt.Sprintf("%s %s: %s\n",
				timeStyle.Render(dateStr),
				senderStyle.Render(msg.Sender),
				msgStyle.Render(msg.Content),
			))
		}
	}
	m.viewport.SetContent(b.String())
	m.viewport.GotoBottom()
}

func (m model) View() string {
	if !m.authorized {
		return fmt.Sprintf(`
%s
ACCESS DENIED
No public key authentication provided.

(Press Ctrl+C to quit)
`, errStyle.Render("SECURITY ALERT"))
	}

	if m.needsRegistration {
		errStr := ""
		if m.err != nil {
			errStr = "\n" + errStyle.Render(m.err.Error()) + "\n"
		}
		return fmt.Sprintf(`
%s
Welcome! You're a new user.
Please choose a username to register.
%s
%s
`, sysStyle.Render("SECURE TUI CHAT - REGISTRATION"), errStr, m.input.View())
	}

	if m.state == 0 {
		errStr := ""
		if m.err != nil {
			errStr = "\n" + errStyle.Render(m.err.Error()) + "\n"
		}
		return fmt.Sprintf(`
%s
Welcome back, %s.
This environment is encrypted at rest.
Please enter your passphrase to unlock your keys.
%s
%s
`, sysStyle.Render("SECURE TUI CHAT"), senderStyle.Render(m.username), errStr, m.input.View())
	}

	if m.state == 1 {
		var b strings.Builder
		b.WriteString(sysStyle.Render("=== YOUR ROOMS ===") + "\n\n")

		if len(m.rooms) == 0 {
			b.WriteString(sysStyle.Render("No rooms yet.") + "\n")
			b.WriteString(commandStyle.Render("  /join <name>") + " - Join/Create room\n")
			b.WriteString(commandStyle.Render("  /dm <user>") + " - Start a DM\n")
		} else {
			for i, room := range m.rooms {
				var style lipgloss.Style
				var prefix string
				if strings.Contains(room.Name, "Note to Self") {
					style = noteStyle
					prefix = "[*]"
				} else if room.IsDM {
					style = dmStyle
					prefix = "[DM]"
				} else {
					style = roomStyle
					prefix = "[#]"
				}
				b.WriteString(fmt.Sprintf("%s %s %s\n",
					senderStyle.Render(fmt.Sprintf("%d.", i+1)),
					prefix,
					style.Render(room.Name)))
			}
			b.WriteString("\n" + sysStyle.Render("Commands: ") + commandStyle.Render("/join /dm") + "\n")
		}

		if m.err != nil {
			b.WriteString("\n" + errStyle.Render(m.err.Error()))
		}
		b.WriteString("\n" + m.input.View())
		return b.String()
	}

	if m.state == 2 {
		return fmt.Sprintf("%s\n%s", m.viewport.View(), m.input.View())
	}

	return m.input.View()
}