package main

import (
	"encoding/json"
	"log"
	"net/http"
	"os"

	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
	"github.com/gorilla/handlers"
)

type mxError struct {
	ErrCode string `json:"errcode"`
	Error   string `json:"error"`
}

var internalErrorBuf = []byte(`{"errcode":"M_UNKNOWN", "error": "Internal error"}`)

type RoomMapping struct {
	RoomID  string   `json:"room_id"`
	Servers []string `json:"servers"`
}

var KnownRooms = map[string]RoomMapping{}

func SendJSON(w http.ResponseWriter, obj any, status int) {
	w.Header().Set("Content-Type", "application/json")
	buf, err := json.Marshal(obj)
	if err != nil {
		log.Printf("Error marshaling JSON response: %v", err)
		w.WriteHeader(500)
		_, _ = w.Write(internalErrorBuf)
		return
	}

	w.WriteHeader(status)
	_, _ = w.Write(buf)
}

func NotFound(w http.ResponseWriter, r *http.Request) {
	SendJSON(w, mxError{"M_UNRECOGNIZED", "Unsupported Endpoint"}, 404)
}

func MethodNotAllowed(w http.ResponseWriter, r *http.Request) {
	SendJSON(w, mxError{"M_UNRECOGNIZED", "Unrecognized Method"}, 405)
}

func Version(w http.ResponseWriter, r *http.Request) {
	SendJSON(w, map[string]any{
		"name":    "Matrix-Vanity",
		"version": "0.0",
	}, 200)
}

func QueryDirectory(w http.ResponseWriter, r *http.Request) {
	q := r.URL.Query()
	roomAlias := q.Get("room_alias")

	mapping, ok := KnownRooms[roomAlias]
	if !ok {
		log.Printf("Room '%s' not found", roomAlias)
		SendJSON(w, mxError{"M_NOT_FOUND", "Room alias not found."}, 404)
		return
	}

	SendJSON(w, mapping, 200)
}

func Index(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(200)
	w.Write([]byte("This is a Matrix-Vanity server"))
}

func main() {
	useProxyIP := os.Getenv("MXV_USE_PROXY_HEADERS") != ""

	configFile := os.Getenv("MXV_ROOMS_FILE")
	if configFile == "" {
		configFile = "/etc/matrix-vanity-rooms.json"
	}
	log.Printf("Loading room list from %s", configFile)

	listenAddr := os.Getenv("MXV_LISTEN")
	if listenAddr == "" {
		listenAddr = ":3333"
	}

	buf, err := os.ReadFile(configFile)
	if err != nil {
		log.Fatalf("Error reading room file: %v", err)
	}

	if err := json.Unmarshal(buf, &KnownRooms); err != nil {
		log.Fatalf("Error parsing room file: %v", err)
	}

	log.Print("Defined Rooms:")
	for alias, info := range KnownRooms {
		log.Printf(" - '%s', ID: '%s', Servers: %v", alias, info.RoomID, info.Servers)
	}
	if len(KnownRooms) == 0 {
		log.Printf("⚠️ No rooms defined. That doesn't seem very useful ")
	}
	log.Print("")

	r := chi.NewRouter()
	if useProxyIP {
		log.Print("X-Forwarded-For header intepretation enabled; it is assumed you are running this service behind a proxy")
		r.Use(handlers.ProxyHeaders)
	}
	r.Use(middleware.Logger)
	r.Use(middleware.Recoverer)

	r.NotFound(NotFound)
	r.MethodNotAllowed(MethodNotAllowed)
	r.Get("/", Index)
	r.Get("/_matrix/federation/v1/version", Version)
	r.Get("/_matrix/federation/v1/query/directory", QueryDirectory)
	log.Printf("Listening on %q", listenAddr)
	http.ListenAndServe(listenAddr, r)
}