package main

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/json"
	"fmt"
	"strings"
	"io"
	"net"
	"os"
	"time"

	"github.com/gotk3/gotk3/glib"
	"github.com/gotk3/gotk3/gtk"
	"gopkg.in/yaml.v2"
	"github.com/yang3yen/xxtea-go/xxtea"
)

// Copyright (c) 2025 Robert Strutts, License MIT

// Global Vars:
var (
	encryptor 	Encryptor
	config  	Config
)

type Config struct {
	Window struct {
		Title string `yaml:"title"`
	} `yaml:"window"`
	Connection struct {
		Protocol string `yaml:"protocol"`
		IP       string `yaml:"ip"`
		Port     int    `yaml:"port"`
		Encryption struct {
			Type string `yaml:"type"`
			Key  string `yaml:"key"`
		} `yaml:"encryption"`
	} `yaml:"connection"`
	User struct {
		Username string `yaml:"username"`
		Timezone string `yaml:"timezone"`
	} `yaml:"user"`
}

type Message struct {
	Username string
	Text     string
	Time     string
}

type Encryptor interface {
	Encrypt([]byte) ([]byte, error)
	Decrypt([]byte) ([]byte, error)
}

type AESEncryptor struct {
	key []byte
}

type XXTEAEncryptor struct {
	key []byte
}

func (a *AESEncryptor) Encrypt(plaintext []byte) ([]byte, error) {
	block, err := aes.NewCipher(a.key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	nonce := make([]byte, gcm.NonceSize())
	if _, err = io.ReadFull(rand.Reader, nonce); err != nil {
		return nil, err
	}
	return gcm.Seal(nonce, nonce, plaintext, nil), nil
}

func (a *AESEncryptor) Decrypt(ciphertext []byte) ([]byte, error) {
	block, err := aes.NewCipher(a.key)
	if err != nil {
		return nil, err
	}
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}
	nonceSize := gcm.NonceSize()
	if len(ciphertext) < nonceSize {
		return nil, fmt.Errorf("ciphertext too short")
	}
	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]
	return gcm.Open(nil, nonce, ciphertext, nil)
}

func (x *XXTEAEncryptor) Encrypt(plaintext []byte) ([]byte, error) {
	result, err := xxtea.Encrypt(plaintext, x.key, true, 0)
	if err != nil {
		return nil, err
	}
	return result, nil
}

func (x *XXTEAEncryptor) Decrypt(ciphertext []byte) ([]byte, error) {
	result, err := xxtea.Decrypt(ciphertext, x.key, true, 0)
	if err != nil {
		return nil, err
	}
	return result, nil
}

func loadConfig(filename string) (*Config, error) {
	var key	[]byte
	
	data, err := os.ReadFile(filename)
	if err != nil {
		return nil, fmt.Errorf("error reading config file: %v", err)
	}

	err = yaml.Unmarshal(data, &config)
	if err != nil {
		return nil, fmt.Errorf("error parsing config file: %v", err)
	}

	key = []byte(config.Connection.Encryption.Key)
	
	// Initialize the appropriate encryptor
	switch config.Connection.Encryption.Type {
	case "aes":
		// AES requires exactly 32 bytes key
		if len(key) != 32 {
			return nil, fmt.Errorf("AES key must be exactly 32 bytes")
		}
		encryptor = &AESEncryptor{key: key}
	case "xxtea":
		// XXTEA Recommended 16 bytes key
		if len(key) < 16 {
			return nil, fmt.Errorf("XXTEA key should be at least 16 bytes")
		}
		key := key[:16] // Take the first 16 bytes
		encryptor = &XXTEAEncryptor{key: key}
	default:
		return nil, fmt.Errorf("unsupported encryption type: %s", config.Connection.Encryption.Type)
	}

	return &config, nil
}

func setupTreeView(treeView *gtk.TreeView) *gtk.ListStore {
	renderer, _ := gtk.CellRendererTextNew()

	columns := []struct {
		Title string
		ID    int
	}{
		{"Username", 0},
		{"Message", 1},
		{"Time", 2},
	}

	for _, col := range columns {
		column, _ := gtk.TreeViewColumnNew()
		column.SetTitle(col.Title)
		column.PackStart(renderer, true)
		column.AddAttribute(renderer, "text", col.ID)
		treeView.AppendColumn(column)
	}

	listStore, _ := gtk.ListStoreNew(glib.TYPE_STRING, glib.TYPE_STRING, glib.TYPE_STRING)
	treeView.SetModel(listStore)

	return listStore
}

func scrollToBottom(scrolledWindow *gtk.ScrolledWindow, treeView *gtk.TreeView) {
    // Get the vertical adjustment
    adjustment := scrolledWindow.GetVAdjustment()
    
    // Get the last path in the tree view
    model, err := treeView.GetModel()
    if err != nil {
        fmt.Printf("Error getting model: %v\n", err)
        return
    }
    
    listStore, ok := model.(*gtk.ListStore)
    if !ok {
        fmt.Printf("Error: model is not a ListStore\n")
        return
    }

    // Get first iterator
    iter, valid := listStore.GetIterFirst()
    if !valid {
        return
    }

    var lastPath *gtk.TreePath
    for {
        path, err := listStore.GetPath(iter)
        if err != nil {
            break
        }
        lastPath = path
        
        if !listStore.IterNext(iter) {
            break
        }
    }
    
    if lastPath != nil {
        // First scroll to make sure the last row is in view
        treeView.ScrollToCell(lastPath, nil, true, 1.0, 1.0)
        
        // Process events to ensure scroll takes effect
        for gtk.EventsPending() {
            gtk.MainIterationDo(false)
        }

        // Force the adjustment to the maximum value
        upper := adjustment.GetUpper()
        pageSize := adjustment.GetPageSize()
        
        // Set value to maximum possible
        adjustment.SetValue(upper - pageSize)
        
        // Force another GUI update
        for gtk.EventsPending() {
            gtk.MainIterationDo(false)
        }
        
        // Set the value one more time to ensure it sticks
        adjustment.SetValue(upper - pageSize)
    }
}

func wrapText(text string, lineWidth int) string {
	words := strings.Fields(text) // Split the text into words
	var result strings.Builder
	var line string

	for _, word := range words {
		if len(line)+len(word)+1 > lineWidth {
			result.WriteString(line + "\n")
			line = word
		} else {
			if line != "" {
				line += " "
			}
			line += word
		}
	}

	if line != "" {
		result.WriteString(line)
	}

	return result.String()
}

func say(conn net.Conn, text string, username string, timezone *time.Location) {
	msg := Message{
		Username: username,
		Text:     text,
		Time:     time.Now().In(timezone).Format("2006-01-02 03:04:05 PM"),
	}

	jsonMsg, _ := json.Marshal(msg)
	encrypted, err := encryptor.Encrypt(jsonMsg)
	if err != nil {
		fmt.Printf("Encryption error: %v\n", err)
		return
	}
	conn.Write(encrypted)
}

func showErrorDialog(parent *gtk.Window, message string) {
	// Create a new message dialog
	dialog := gtk.MessageDialogNew(
		parent,                      // Parent window
		gtk.DIALOG_MODAL,            // Make the dialog modal
		gtk.MESSAGE_ERROR,           // Message type (error)
		gtk.BUTTONS_OK,              // Buttons to show ("OK")
		message,                     // Message text
	)

	// Run the dialog and wait for the user to close it
	dialog.Run()
	dialog.Destroy() // Clean up the dialog
}

func tryReconnect(conn *net.Conn, window *gtk.Window, addr string, config *Config) error {
    var newConn net.Conn
    var err error

    if config.Connection.Protocol == "tcp" {
        newConn, err = net.Dial("tcp", addr)
    } else {
        udpAddr, resolveErr := net.ResolveUDPAddr("udp", addr)
        if resolveErr != nil {
            return fmt.Errorf("failed to resolve UDP address: %v", resolveErr)
        }
        newConn, err = net.DialUDP("udp", nil, udpAddr)
    }

    if err != nil {
		showErrorDialog(window, "Failed to reconnect. Exiting...")
        gtk.MainQuit()
        return err
    }

    // Update the caller's connection variable
    *conn = newConn
    return nil
}

func main() {
	// Load configuration
	config, err := loadConfig("chat_client.yaml")
	if err != nil {
		fmt.Println("Error loading config:", err)
		os.Exit(1)
	}

	// Load timezone
	timezone, err := time.LoadLocation(config.User.Timezone)
	if err != nil {
		fmt.Println("Error loading timezone:", err)
		timezone = time.Local // Fallback to local time if timezone loading fails
	}

	gtk.Init(nil)

	builder, _ := gtk.BuilderNew()
	builder.AddFromString(`
	<interface>
		<object class="GtkWindow" id="window">
			<property name="title">` + config.Window.Title + `</property>
			<property name="default-width">600</property>
			<property name="default-height">400</property>
			<child>
				<object class="GtkBox" id="box">
					<property name="orientation">vertical</property>
					<child>
						<object class="GtkScrolledWindow" id="scrolled">
							<property name="vexpand">true</property>
							<property name="hexpand">true</property>
							<child>
								<object class="GtkTreeView" id="treeview">
									<property name="vexpand">true</property>
									<property name="hexpand">true</property>
								</object>
							</child>
						</object>
					</child>
					<child>
						<object class="GtkEntry" id="entry">
							<signal name="activate" handler="send_message"/>
						</object>
					</child>
				</object>
			</child>
		</object>
	</interface>
	`)

	windowObj, _ := builder.GetObject("window")
	window := windowObj.(*gtk.Window)
	window.Connect("destroy", gtk.MainQuit)

	treeViewObj, _ := builder.GetObject("treeview")
	treeView := treeViewObj.(*gtk.TreeView)
	listStore := setupTreeView(treeView)

	scrolledObj, _ := builder.GetObject("scrolled")
	scrolledWindow := scrolledObj.(*gtk.ScrolledWindow)

	entryObj, _ := builder.GetObject("entry")
	entry := entryObj.(*gtk.Entry)

	// Setup connection
	addr := fmt.Sprintf("%s:%d", config.Connection.IP, config.Connection.Port)
	var conn net.Conn
	var connErr error

	if config.Connection.Protocol == "tcp" {
		conn, connErr = net.Dial("tcp", addr)
	} else {
		udpAddr, _ := net.ResolveUDPAddr("udp", addr)
		conn, connErr = net.DialUDP("udp", nil, udpAddr)
	}

	// Connect the "destroy" signal to quit the application
	window.Connect("destroy", func() {
		fmt.Println("Window closed")
		say(conn, "Good Bye, User(Off-Line)", config.User.Username, timezone)
		gtk.MainQuit()
	})

	if connErr != nil {
		fmt.Printf("Connection error (%s://%s): %v\n", config.Connection.Protocol, addr, connErr)
		showErrorDialog(window, "Unable to connect to the server. Please check the server and try again.")
		os.Exit(1)
	}
	defer conn.Close()
	say(conn, "Welcome, User(On-line)", config.User.Username, timezone)

	builder.ConnectSignals(map[string]interface{}{
		"send_message": func() {
			text, _ := entry.GetText()
			entry.SetText("")
			say(conn, text, config.User.Username, timezone)
		},
	})

	go func() {
		for {
			buf := make([]byte, 1024)
			n, err := conn.Read(buf)
			if err != nil {
				//if err == io.EOF {
				showErrorDialog(window, "Disconnected, try again...")
				err := tryReconnect(&conn, window, addr, config)
				if err != nil {
					fmt.Printf("Reconnection failed: %v\n", err)
					gtk.MainQuit()
				}
				continue
			}
			decrypted, err := encryptor.Decrypt(buf[:n])
			if err != nil {
				continue
			}

			var msg Message
			json.Unmarshal(decrypted, &msg)

			glib.IdleAdd(func() {
				iter := listStore.Append()
				err := listStore.Set(iter,
					[]int{0, 1, 2},
					[]interface{}{msg.Username, wrapText(msg.Text, 45), msg.Time},
				)
				if err != nil {
					fmt.Println("Failed to set row values:", err)
				}
				scrollToBottom(scrolledWindow, treeView)
				
				window.Deiconify() // UnMinimize -- nope, doesn't always do that...
				window.Present()   // Popup on New MSG!
			})
		}
	}()

	window.ShowAll()
	gtk.Main()
}