package main

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"sync"
	"time"
)

const PORT = ":5518"

type PacketID int

const (
	CONNECT_REQUEST PacketID = iota
	CONNECT_RESPONSE
	JOIN_REQUEST
	JOIN_RESPONSE
	CREATE_REQUEST
	CREATE_RESPONSE
	READ_DATA
	SEND_DATA_REQUEST
	SEND_DATA_RESPONSE
	LEAVE_ROOM
)

type Client struct {
	Conn     net.Conn
	ClientID string
	RoomCode string
}

type Room struct {
	Code     string
	Password string
	Clients  []*Client
}

var rooms = map[string]*Room{}
var clients = map[string]*Client{}
var mu sync.Mutex

func readPacket(conn net.Conn) (PacketID, []byte, error) {
	// First, read the length of the packet (4 bytes)
	var length uint32
	if err := binary.Read(conn, binary.BigEndian, &length); err != nil {
		return 0, nil, err
	}

	// Then read the packet ID (1 byte)
	var messageType byte
	if err := binary.Read(conn, binary.BigEndian, &messageType); err != nil {
		return 0, nil, err
	}

	// Read the remaining data (length - 5 bytes, since 4 bytes for length and 1 byte for messageType)
	data := make([]byte, length-5)
	if _, err := io.ReadFull(conn, data); err != nil {
		return 0, nil, err
	}

	return PacketID(int(messageType)), data, nil
}

func writePacket(conn net.Conn, messageType byte, data []byte) error {
	// Calculate the total length of the packet
	// 4 bytes for length, 1 byte for messageType
	length := uint32(5 + len(data))

	// Write the length and the message type
	if err := binary.Write(conn, binary.BigEndian, length); err != nil {
		return err
	}
	if err := binary.Write(conn, binary.BigEndian, messageType); err != nil {
		return err
	}

	// Write the data
	_, err := conn.Write(data)
	return err
}

func main() {
	listener, err := net.Listen("tcp", PORT)
	if err != nil {
		fmt.Println("Error starting server:", err)
		return
	}
	defer listener.Close()

	fmt.Println("Server started on port", PORT)

	for {
		conn, err := listener.Accept()
		if err != nil {
			fmt.Println("Error accepting connection:", err)
			continue
		}
		go handleClient(conn)
	}
}

type ConnectionRequestPacket struct {
	UserID   string
	RoomCode string
	Password string
}

type ConnectionResponsePacket struct {
	Success bool
	Reason  string
}

type CreateRoomRequestPacket struct {
	UserID   string
	Password string
}

type CreateRoomResponsePacket struct {
	Success  bool
	Reason   string
	RoomCode string
}

type JoinRequestPacket struct {
	UserID   string
	Password string
	RoomCode string
}

type JoinRequestResponsePacket struct {
	Success     bool
	Reason      string
	CurrentData string
}

type SendDataRequestPacket struct {
	UserID string
	Data   string
}

type ForwardDataPacket struct {
	Data string
}

type SendDataResponsePacket struct {
	Success bool
	Reason  string
}

func handleClient(conn net.Conn) {
	// defer conn.Close()

	for {

		var packetId PacketID
		packetId, data, err := readPacket(conn)
		if err != nil {
			fmt.Println("Error reading packet (handleclient) :", err)
			handleLeaveRoom(conn)
			return
		}

		switch packetId {
		case CONNECT_REQUEST:
			var packet ConnectionRequestPacket
			if err := json.Unmarshal(data, &packet); err != nil {
				fmt.Println("Error decoding connection request:", err)
				return
			}
			handleConnectRequest(conn, packet)
		case CREATE_REQUEST:
			var packet CreateRoomRequestPacket
			if err := json.Unmarshal(data, &packet); err != nil {
				fmt.Println("Error decoding connection request:", err)
				return
			}
			handleCreateRoomRequest(conn, packet)
		case JOIN_REQUEST:
			var packet JoinRequestPacket
			if err := json.Unmarshal(data, &packet); err != nil {
				fmt.Println("Error decoding connection request:", err)
				return
			}
			handleJoinRequest(conn, packet)
		case SEND_DATA_REQUEST:
			var packet SendDataRequestPacket
			if err := json.Unmarshal(data, &packet); err != nil {
				fmt.Println("Error decoding connection request:", err)
				return
			}
			handleSendData(conn, packet)
		case LEAVE_ROOM:
			handleLeaveRoom(conn)
		}
	}
}

func handleConnectRequest(conn net.Conn, packet ConnectionRequestPacket) {
	fmt.Println("Incoming connection request")
	room := findRoom(packet.RoomCode)
	if room == nil {
		writePacket(conn, byte(CONNECT_RESPONSE), encodeResponsePacket(ConnectionResponsePacket{
			Success: false,
			Reason:  "room does not exist",
		}))
		fmt.Println("Attempted room does not exist")
		return
	}

	if room.Password == "" || packet.Password == room.Password {
		writePacket(conn, byte(CONNECT_RESPONSE), encodeResponsePacket(ConnectionResponsePacket{
			Success: true,
			Reason:  "",
		}))
		fmt.Println("Invalid password")
		return
	}

	writePacket(conn, byte(CONNECT_RESPONSE), encodeResponsePacket(ConnectionResponsePacket{
		Success: false,
		Reason:  "invalid password",
	}))
	fmt.Println("Connected user to room ")
}

func handleCreateRoomRequest(conn net.Conn, packet CreateRoomRequestPacket) {
	if !isValidPassword(packet.Password) {
		writePacket(conn, byte(CREATE_RESPONSE), encodeResponsePacket(CreateRoomResponsePacket{
			Success:  false,
			Reason:   "invalid password",
			RoomCode: "",
		}))
		return
	}
	code := generateRandomRoomCode()

	room := &Room{
		Code:     code,
		Password: packet.Password,
		Clients:  []*Client{},
	}

	mu.Lock()
	rooms[code] = room
	mu.Unlock()

	writePacket(conn, byte(CREATE_RESPONSE), encodeResponsePacket(CreateRoomResponsePacket{
		Success:  true,
		Reason:   "",
		RoomCode: code,
	}))
}

func handleJoinRequest(conn net.Conn, packet JoinRequestPacket) {
	room := findRoom(packet.RoomCode)
	if room == nil {
		writePacket(conn, byte(JOIN_RESPONSE), encodeResponsePacket(JoinRequestResponsePacket{
			Success: false,
			Reason:  "room not found",
		}))
		return
	}

	if !isValidPassword(packet.Password) {
		writePacket(conn, byte(JOIN_RESPONSE), encodeResponsePacket(JoinRequestResponsePacket{
			Success: false,
			Reason:  "invalid password",
		}))
		return
	}

	if room.Password == "" || room.Password == packet.Password {
		client := &Client{
			Conn:     conn,
			ClientID: packet.UserID,
			RoomCode: packet.RoomCode,
		}

		mu.Lock()
		room.Clients = append(room.Clients, client)
		clients[packet.UserID] = client
		mu.Unlock()

		writePacket(conn, byte(JOIN_RESPONSE), encodeResponsePacket(JoinRequestResponsePacket{
			Success:     true,
			Reason:      "",
			CurrentData: "{}",
			// TODO: Send current data
		}))
		return
	}

	writePacket(conn, byte(JOIN_RESPONSE), encodeResponsePacket(JoinRequestResponsePacket{
		Success: false,
		Reason:  "invalid password",
	}))
	return
}

func handleSendData(conn net.Conn, packet SendDataRequestPacket) {
	fmt.Println("Incoming send data request")
	mu.Lock()
	client, exists := clients[packet.UserID]
	mu.Unlock()

	if !exists || client.RoomCode == "" {
		writePacket(conn, byte(SEND_DATA_RESPONSE), encodeResponsePacket(SendDataResponsePacket{
			Success: false,
			Reason:  "client not in a room",
		}))
		return
	}

	room := findRoom(client.RoomCode)
	if room == nil {
		writePacket(conn, byte(SEND_DATA_RESPONSE), encodeResponsePacket(SendDataResponsePacket{
			Success: false,
			Reason:  "room not found",
		}))
		return
	}

	dataPacket := encodeResponsePacket(ForwardDataPacket{
		Data: packet.Data,
	})

	mu.Lock()
	for _, roomClient := range room.Clients {
		// if roomClient.ClientID != packet.UserID {
		// 	// Send data to all other clients except the sender
		// 	writePacket(roomClient.Conn, byte(READ_DATA), dataPacket)
		// }
		writePacket(roomClient.Conn, byte(READ_DATA), dataPacket)
	}
	mu.Unlock()

	writePacket(conn, byte(SEND_DATA_RESPONSE), encodeResponsePacket(SendDataResponsePacket{
		Success: true,
		Reason:  "",
	}))
}

func handleLeaveRoom(conn net.Conn) {
	mu.Lock()
	defer mu.Unlock()

	var leavingClient *Client
	for _, client := range clients {
		if client.Conn == conn {
			leavingClient = client
			break
		}
	}
	if leavingClient == nil {
		fmt.Println("Client not found")
		return
	}

	room := rooms[leavingClient.RoomCode]
	if room == nil {
		fmt.Println("Room not found")
		return
	}

	for i, client := range room.Clients {
		if client == leavingClient {
			room.Clients = append(room.Clients[:i], room.Clients[i+1:]...)
			break
		}
	}

	// Delete the room if its empty
	if len(room.Clients) == 0 {
		delete(rooms, room.Code)
		fmt.Println("Room", room.Code, "deleted")
	}

	delete(clients, leavingClient.ClientID)

	done := make(chan struct{})
	go func() {
		leavingClient.Conn.Close()
		close(done)
	}()

	select {
	case <-done:
	case <-time.After(5 * time.Second):
		log.Fatalln("Timeout while closing leaving client's connection")
	}

	fmt.Println("Client", leavingClient.ClientID, "left the room and connection closed")
}

func findRoom(code string) *Room {
	mu.Lock()
	defer mu.Unlock()
	return rooms[code]
}

func encodeResponsePacket(packet interface{}) []byte {
	data, err := json.Marshal(packet)
	if err != nil {
		fmt.Println("Error encoding response packet:", err)
		return nil
	}
	return data
}

func generateRandomRoomCode() string {
	validChars := "abcdefghijklmnopqrstuvwxyz0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZ"
	code := make([]byte, 4)
	for i := range code {
		code[i] = validChars[rand.Intn(len(validChars))]
	}
	return string(code)
}

func isValidPassword(password string) bool {
	validChars := "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789?!()[]{}<>-_"
	charMap := make(map[rune]bool)
	for _, char := range validChars {
		charMap[char] = true
	}

	for _, char := range password {
		if !charMap[char] {
			return false
		}
	}
	return true
}