package main

import (
	"compress/gzip"
	"encoding/json"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

type gzipResponseWriter struct {
	http.ResponseWriter
	io.Writer
}

func (g *gzipResponseWriter) Write(data []byte) (int, error) {
	return g.Writer.Write(data)
}

func GzipMiddleware(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if !strings.Contains(r.Header.Get("Accept-Encoding"), "gzip") {
			next.ServeHTTP(w, r)
			return
		}

		gz := gzip.NewWriter(w)
		defer gz.Close()

		w.Header().Set("Content-Encoding", "gzip")
		next.ServeHTTP(&gzipResponseWriter{ResponseWriter: w, Writer: gz}, r)
	})
}

func getClientIP(r *http.Request) string {
	if fwdIP := r.Header.Get("X-Forwarded-For"); fwdIP != "" {
		return strings.Split(fwdIP, ",")[0]
	}

	clientIP := r.RemoteAddr
	if host, _, err := net.SplitHostPort(clientIP); err == nil {
		return host
	}
	return clientIP
}

type Message struct {
	Content        string
	SenderIp       string
	SenderUsername string
	Timestamp      string
}

func NewMessage(username, ip, content string) *Message {
	timestamp := time.Now().Format("2006-01-02 15:04:05")
	return &Message{
		Content:        content,
		SenderIp:       ip,
		SenderUsername: username,
		Timestamp:      timestamp,
	}
}

type Server struct {
	Ip       string
	Port     int
	Messages []Message
	// Add mappings for IPs and usernames
	Connected map[string]time.Time // Map IP -> Last activity time
	Usernames map[string]string    // Map IP -> Username
	mu        sync.Mutex           // For thread safety
}

func NewServer(ip string, port int) *Server {
	return &Server{
		Ip:        ip,
		Port:      port,
		Messages:  make([]Message, 0),
		Connected: make(map[string]time.Time),
		Usernames: make(map[string]string),
		mu:        sync.Mutex{},
	}
}

func (s *Server) AddMessage(username string, userip string, contents string) {
	message := NewMessage(username, userip, contents)
	s.Messages = append(s.Messages, *message)
}

func (s *Server) updateActivity(ip string) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.Connected[ip] = time.Now()
}

func (s *Server) cleanupActivity() {
	s.mu.Lock()
	defer s.mu.Unlock()
	for ip, lastActivity := range s.Connected {
		if time.Since(lastActivity) > 10*time.Second {
			delete(s.Connected, ip)
		}
	}
}

func (s *Server) handlePing(w http.ResponseWriter, r *http.Request) {
	clientIP := getClientIP(r)
	s.updateActivity(clientIP)
	s.cleanupActivity()
	w.WriteHeader(http.StatusOK)
}

func validUsername(username string) bool {
	for _, c := range username {
		if !((c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_') {
			return false
		}
	}
	return true
}

func (s *Server) handleUsername(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPut {
		http.Error(w, `{"error": "Method not allowed"}`, http.StatusMethodNotAllowed)
		return
	}

	w.Header().Set("Content-Type", "application/json")

	clientIP := getClientIP(r)

	var req struct {
		Username string `json:"username"`
	}

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, `{"error": "Invalid JSON"}`, http.StatusBadRequest)
		return
	}

	s.mu.Lock()

	if len(req.Username) > 64 {
		http.Error(w, fmt.Sprintf(`{"error": "Username too long (must be less than 64 characters)"}`), http.StatusRequestEntityTooLarge)
		s.mu.Unlock()
		return
	}

	if !validUsername(req.Username) {
		http.Error(w, fmt.Sprintf(`{"error": "Username must only contain alphanumeric characters and/or underscores"}`), http.StatusBadRequest)
		s.mu.Unlock()
		return
	}

	for _, username := range s.Usernames {
		if username == req.Username {
			s.mu.Unlock()
			http.Error(w, fmt.Sprintf(`{"error": "Username already exists"}`), http.StatusConflict)
			return
		}
	}

	if _, ok := s.Usernames[clientIP]; ok {
		fmt.Println("Name change detected, fixing messages to reflect the change")
		for i, msg := range s.Messages {
			if msg.SenderIp == clientIP {
				s.Messages[i].SenderUsername = req.Username
			}
		}
	}
	s.Usernames[clientIP] = req.Username

	s.mu.Unlock()

	json.NewEncoder(w).Encode(map[string]string{"status": "Username registered"})
	fmt.Println(clientIP, "->", req.Username)
}

func getMessageTemplate(file string, body string) string {
	contents, _ := os.ReadFile(file)
	return strings.Replace(string(contents), "{{body}}", body, 1)
}

func (s *Server) handleMessages(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case http.MethodGet:
		w.Header().Set("Content-Type", "text/html")

		var body string
		for _, msg := range s.Messages {
			body += fmt.Sprintf(`<p><span class="username">%s  </span><span class="timestamp">%s</span><br><span class="message">%s</span></p>`,
				msg.SenderUsername, msg.Timestamp, msg.Content)
		}

		w.Write([]byte(getMessageTemplate("messages.html", body)))

	case http.MethodPut:
		w.Header().Set("Content-Type", "application/json")

		// Get client's IP
		clientIP := getClientIP(r)

		// Check if IP is registered
		s.mu.Lock()
		username, exists := s.Usernames[clientIP]
		s.mu.Unlock()

		if !exists {
			errorFmt := fmt.Sprintf(`{"error": "IP %s not registered with username"}`, clientIP)
			http.Error(w, errorFmt, http.StatusUnauthorized)
			return
		}

		var msg struct {
			Message string `json:"message"`
		}

		if err := json.NewDecoder(r.Body).Decode(&msg); err != nil {
			http.Error(w, `{"error": "Invalid JSON"}`, http.StatusBadRequest)
			return
		}

		// Add message using the stored username and IP
		s.AddMessage(username, clientIP, msg.Message)

		json.NewEncoder(w).Encode(map[string]string{
			"status": "Message received",
			"from":   username,
		})
	default:
		http.Error(w, `{"error": "Method not allowed"}`, http.StatusMethodNotAllowed)
		return
	}
}

func (s *Server) handleUsernameStatus(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, `{"error": "Method not allowed"}`, http.StatusMethodNotAllowed)
		return
	}

	w.Header().Set("Content-Type", "application/json")
	clientIP := getClientIP(r)

	s.mu.Lock()
	username, exists := s.Usernames[clientIP]
	s.mu.Unlock()

	json.NewEncoder(w).Encode(map[string]interface{}{
		"hasUsername": exists,
		"username":    username,
	})
}

func (s *Server) handleUsers(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		http.Error(w, `{"error": "Method not allowed"}`, http.StatusMethodNotAllowed)
		return
	}

	w.Header().Set("Content-Type", "application/json")

	clientIP := getClientIP(r)
	s.updateActivity(clientIP)
	s.cleanupActivity()
	s.mu.Lock()
	var users []string
	for ip := range s.Connected {
		users = append(users, s.Usernames[ip])
	}
	s.mu.Unlock()

	json.NewEncoder(w).Encode(map[string]interface{}{
		"users": users,
	})
}

func (s *Server) handleRoot(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.NotFound(w, r)
		return
	}
	w.Header().Set("Content-Type", "text/html")
	w.Write(readFile("root.html"))
}

func readFile(filepath string) []byte {
	contents, _ := os.ReadFile(filepath)
	return contents
}

func (s *Server) handleJs(w http.ResponseWriter, r *http.Request) {
	_ = r
	w.Header().Set("Content-Type", "application/javascript")
	w.Write(readFile("root.js"))
}

func (s *Server) handleCss(w http.ResponseWriter, r *http.Request) {
	_ = r
	w.Header().Set("Content-Type", "text/css")
	w.Write(readFile("root.css"))
}

func (s *Server) Run() {
	handler := http.NewServeMux()
	handler.HandleFunc("/ping", s.handlePing)
	handler.HandleFunc("/username", s.handleUsername)
	handler.HandleFunc("/messages", s.handleMessages)
	handler.HandleFunc("/username/status", s.handleUsernameStatus)
	handler.HandleFunc("/users", s.handleUsers)
	handler.HandleFunc("/", s.handleRoot)
	handler.HandleFunc("/root.js", s.handleJs)
	handler.HandleFunc("/root.css", s.handleCss)
	fmt.Printf("Server starting on %s:%d\n", s.Ip, s.Port)
	if err := http.ListenAndServe(fmt.Sprintf("%s:%d", s.Ip, s.Port), GzipMiddleware(handler)); err != nil {
		fmt.Printf("Server error: %v\n", err)
	}
}

func main() {
	ip := os.Args[1]
	port, _ := strconv.Atoi(os.Args[2])
	// ip := "localhost"
	// port := 8080
	server := NewServer(ip, port)
	server.Run()
}
