package main

import (
	"embed"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"time"

	"zoraxy-firewall/zoraxy_plugin"
)

//go:embed web/*
var webFS embed.FS

func main() {
	loadConfig()
	loadHostMap()

	spec := &zoraxy_plugin.IntroSpect{
		ID:            "zoraxy-firewall",
		Name:          "Firewall",
		Author:        "Zoraxy Community",
		AuthorContact: "",
		Description:   "Transparent WAF proxy — forwards traffic to firewall for inline inspection.",
		URL:           "",
		Type:          zoraxy_plugin.PluginType_Router,
		VersionMajor:  1,
		VersionMinor:  0,
		VersionPatch:  0,
		UIPath:        "/ui",

		StaticCapturePaths:   []zoraxy_plugin.StaticCaptureRule{{CapturePath: "/"}},
		StaticCaptureIngress: "/inspect",
	}

	config, err := zoraxy_plugin.ServeAndRecvSpec(spec)
	if err != nil {
		log.Fatalf("failed to receive config: %v", err)
	}

	startHealthCheck()

	// Return port listener — receives traffic from WAF, proxies to Zoraxy.
	cfg := getConfig()
	if cfg.ReturnPort > 0 {
		go func() {
			rmux := http.NewServeMux()
			rmux.HandleFunc("/", handleReturn)
			addr := fmt.Sprintf(":%d", cfg.ReturnPort)
			log.Printf("Firewall: return listener on %s", addr)
			if err := http.ListenAndServe(addr, rmux); err != nil {
				log.Printf("Firewall: return listener error: %v", err)
			}
		}()
	}

	mux := http.NewServeMux()
	mux.Handle("/inspect/", http.HandlerFunc(handleInspect))
	mux.HandleFunc("/inspect", handleInspect)

	uiRouter := zoraxy_plugin.NewPluginEmbedUIRouter(spec.ID, &webFS, "web", spec.UIPath)
	uiRouter.HandleFunc("/api/config", handleGetConfig, mux)
	uiRouter.HandleFunc("/api/config/save", handleSaveConfig, mux)
	uiRouter.HandleFunc("/api/stats", handleGetStats, mux)

	uiRouter.RegisterTerminateHandler(func() {
		log.Println("firewall shutting down")
	}, mux)
	uiRouter.AttachHandlerToMux(mux)

	addr := fmt.Sprintf("127.0.0.1:%d", config.Port)
	log.Printf("firewall v%d.%d.%d listening on %s", spec.VersionMajor, spec.VersionMinor, spec.VersionPatch, addr)

	if err := http.ListenAndServe(addr, mux); err != nil {
		log.Fatalf("server error: %v", err)
	}
}

func handleGetConfig(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, getConfig())
}

func handleSaveConfig(w http.ResponseWriter, r *http.Request) {
	if err := r.ParseForm(); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid form"})
		return
	}

	cfgMu.Lock()
	cfg.Enabled = r.FormValue("enabled") == "true"
	if u := r.FormValue("waf_url"); u != "" {
		cfg.WAFURL = u
	}
	if v := r.FormValue("return_port"); v != "" {
		fmt.Sscanf(v, "%d", &cfg.ReturnPort)
	}
	if v := r.FormValue("health_interval"); v != "" {
		fmt.Sscanf(v, "%d", &cfg.HealthInterval)
	}
	if v := r.FormValue("timeout_ms"); v != "" {
		fmt.Sscanf(v, "%d", &cfg.TimeoutMs)
	}
	cfgMu.Unlock()

	if err := saveConfig(); err != nil {
		log.Printf("ERROR saving config: %v", err)
		writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "save failed"})
		return
	}
	writeJSON(w, http.StatusOK, map[string]string{"status": "saved"})
}

func handleGetStats(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, metricsSnapshot())
}

func handleReturn(w http.ResponseWriter, r *http.Request) {
	backend := resolveBackend(r.Host)
	if !strings.HasPrefix(backend, "http") {
		backend = "http://" + backend
	}
	log.Printf("DEBUG return: %s %s -> %s%s (Host: %s)", r.Method, r.URL.Path, backend, r.URL.Path, r.Host)
	proxyReq, err := http.NewRequestWithContext(r.Context(), r.Method, backend+r.URL.Path, r.Body)
	if err != nil {
		http.Error(w, "proxy error", http.StatusInternalServerError)
		return
	}
	proxyReq.Host = r.Host
	for k, vv := range r.Header {
		for _, v := range vv {
			proxyReq.Header.Add(k, v)
		}
	}
	client := &http.Client{
		Timeout: 30 * time.Second,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
	resp, err := client.Do(proxyReq)
	if err != nil {
		http.Error(w, "backend unreachable", http.StatusBadGateway)
		return
	}
	defer resp.Body.Close()
	for k, vv := range resp.Header {
		for _, v := range vv {
			w.Header().Add(k, v)
		}
	}
	w.WriteHeader(resp.StatusCode)
	io.Copy(w, resp.Body)
}

func writeJSON(w http.ResponseWriter, status int, v interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(v)
}