package main

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"io"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"
)

type Email struct {
	ID        string    `json:"id"`
	From      string    `json:"from"`
	To        string    `json:"to"`
	Subject   string    `json:"subject"`
	Body      string    `json:"body"`
	Timestamp time.Time `json:"timestamp"`
	Domain    string    `json:"domain"`
}

// User struct for account management
type User struct {
	Username string `json:"username"`
	Password string `json:"password"` // hashed
}

var (
	mailboxDir = "mailboxes"
	mu         sync.Mutex
	usersFile  = "users.json"
)

// Load users from file
func loadUsers() (map[string]string, error) {
	users := make(map[string]string)
	data, err := ioutil.ReadFile(usersFile)
	if err != nil {
		if os.IsNotExist(err) {
			return users, nil
		}
		return nil, err
	}
	var userList []User
	if err := json.Unmarshal(data, &userList); err != nil {
		return nil, err
	}
	for _, u := range userList {
		users[u.Username] = u.Password
	}
	return users, nil
}

// Save users to file
func saveUsers(users map[string]string) error {
	var userList []User
	for username, password := range users {
		userList = append(userList, User{Username: username, Password: password})
	}
	data, err := json.MarshalIndent(userList, "", "  ")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(usersFile, data, 0644)
}

// Hash password
func hashPassword(password string) string {
	h := sha256.Sum256([]byte(password))
	return base64.StdEncoding.EncodeToString(h[:])
}

func saveEmail(email Email) error {
	mu.Lock()
	defer mu.Unlock()
	userDir := filepath.Join(mailboxDir, email.To)
	if err := os.MkdirAll(userDir, 0755); err != nil {
		return err
	}
	filePath := filepath.Join(userDir, email.ID+".json")
	f, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer f.Close()
	return json.NewEncoder(f).Encode(email)
}

func createUserHandler(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
		return
	}
	var req struct {
		Username string `json:"username"`
		Password string `json:"password"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		http.Error(w, "Invalid JSON", http.StatusBadRequest)
		return
	}
	if req.Username == "" || req.Password == "" {
		http.Error(w, "Username and password required", http.StatusBadRequest)
		return
	}
	mu.Lock()
	defer mu.Unlock()
	users, err := loadUsers()
	if err != nil {
		http.Error(w, "Server error", http.StatusInternalServerError)
		return
	}
	if _, exists := users[req.Username]; exists {
		http.Error(w, "User already exists", http.StatusConflict)
		return
	}
	hashed := hashPassword(req.Password)
	users[req.Username] = hashed
	if err := saveUsers(users); err != nil {
		http.Error(w, "Server error", http.StatusInternalServerError)
		return
	}
	w.WriteHeader(http.StatusCreated)
	json.NewEncoder(w).Encode(map[string]string{"status": "user created"})
}

// Authenticate user from Basic Auth
func authenticate(r *http.Request) (string, bool) {
	username, password, ok := r.BasicAuth()
	if !ok {
		return "", false
	}
	users, err := loadUsers()
	if err != nil {
		return "", false
	}
	hashed := hashPassword(password)
	if users[username] == hashed {
		return username, true
	}
	return "", false
}

func receiveEmailHandler(w http.ResponseWriter, r *http.Request) {
	_, ok := authenticate(r)
	if !ok {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}
	var email Email
	if err := json.NewDecoder(r.Body).Decode(&email); err != nil {
		http.Error(w, "Invalid JSON", http.StatusBadRequest)
		return
	}
	if email.Domain == "brandon.ad" {
		email.ID = fmt.Sprintf("%d", time.Now().UnixNano())
		email.Timestamp = time.Now()
		if err := saveEmail(email); err != nil {
			http.Error(w, "Failed to save email", http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusCreated)
		json.NewEncoder(w).Encode(map[string]string{"status": "ok", "id": email.ID})
		return
	}
	// Forward to another domain
	forwardURL := fmt.Sprintf("http://%s/email", email.Domain)
	jsonData, err := json.Marshal(email)
	if err != nil {
		http.Error(w, "Failed to marshal email", http.StatusInternalServerError)
		return
	}
	resp, err := http.Post(forwardURL, "application/json", bytes.NewReader(jsonData))
	if err != nil {
		http.Error(w, "Failed to forward email", http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

func listMailboxHandler(w http.ResponseWriter, r *http.Request) {
	username, ok := authenticate(r)
	if !ok {
		http.Error(w, "Unauthorized", http.StatusUnauthorized)
		return
	}
	user := r.URL.Query().Get("user")
	if user == "" {
		http.Error(w, "Missing user", http.StatusBadRequest)
		return
	}
	if user != username {
		http.Error(w, "Forbidden", http.StatusForbidden)
		return
	}
	userDir := filepath.Join(mailboxDir, user)
	files, err := os.ReadDir(userDir)
	if err != nil {
		http.Error(w, "Mailbox not found", http.StatusNotFound)
		return
	}
	var emails []Email
	for _, file := range files {
		f, err := os.Open(filepath.Join(userDir, file.Name()))
		if err != nil {
			continue
		}
		var email Email
		if err := json.NewDecoder(f).Decode(&email); err == nil {
			emails = append(emails, email)
		}
		f.Close()
	}
	json.NewEncoder(w).Encode(emails)
}

// CORS middleware
func withCORS(handler http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Access-Control-Allow-Origin", "*")
		w.Header().Set("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
		w.Header().Set("Access-Control-Allow-Headers", "Content-Type, Authorization")
		if r.Method == http.MethodOptions {
			w.WriteHeader(http.StatusOK)
			return
		}
		handler(w, r)
	}
}

func main() {
	http.HandleFunc("/users", withCORS(createUserHandler))
	http.HandleFunc("/email", withCORS(receiveEmailHandler))
	http.HandleFunc("/mailbox", withCORS(listMailboxHandler))
	log.Println("Email server running on :8080")
	log.Fatal(http.ListenAndServe(":8080", nil))
}