// mock-edge-device is a local HTTPS edge-device stub for manual provisioning tests.
package main

import (
	"crypto/ed25519"
	"crypto/rand"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/base64"
	"encoding/hex"
	"encoding/json"
	"encoding/pem"
	"flag"
	"fmt"
	"log"
	"math/big"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"time"

	mqtt "github.com/eclipse/paho.mqtt.golang"
)

const (
	deviceID    = "MOCK-EDGE-001"
	pairingCode = "246810"
)

type mqttConfig struct {
	Host     string `json:"host"`
	Port     int    `json:"port"`
	TLS      bool   `json:"tls"`
	ClientID string `json:"client_id"`
}

type routeConfig struct {
	ParkingLotID uint   `json:"parking_lot_id"`
	BoothID      uint   `json:"booth_id"`
	ChannelID    uint   `json:"channel_id"`
	Direction    string `json:"direction"`
}

type provisionRequest struct {
	Schema     string      `json:"schema"`
	RequestID  string      `json:"request_id"`
	DeviceID   string      `json:"device_id"`
	DeviceCode string      `json:"device_code"`
	MQTT       mqttConfig  `json:"mqtt"`
	Route      routeConfig `json:"route"`
}

type state struct {
	mu          sync.RWMutex
	provisioned bool
	paired      bool
	config      provisionRequest
	client      mqtt.Client
	certPEM     []byte
	keyPEM      []byte
	certFP      string
	publicKeyFP string
	signingKey  ed25519.PrivateKey
	dataPath    string
}

func main() {
	listen := flag.String("listen", "0.0.0.0:18443", "HTTPS listen address")
	dataDir := flag.String("data-dir", ".run/mock-edge-device", "directory for certificate and provision data")
	mqttFallbackPlain := flag.Bool("mqtt-fallback-plain", false, "test-only: use plain MQTT even when provision says tls=true")
	flag.Parse()

	if err := os.MkdirAll(*dataDir, 0700); err != nil {
		log.Fatal(err)
	}
	certPEM, keyPEM, certFP, err := loadOrCreateCertificate(*dataDir)
	if err != nil {
		log.Fatal(err)
	}
	publicKey, signingKey, err := loadOrCreateSigningKey(*dataDir)
	if err != nil {
		log.Fatal(err)
	}
	publicHash := sha256.Sum256(publicKey)

	s := &state{certPEM: certPEM, keyPEM: keyPEM, certFP: certFP,
		publicKeyFP: "SHA256:" + strings.ToUpper(hex.EncodeToString(publicHash[:])),
		signingKey:  signingKey,
		dataPath:    filepath.Join(*dataDir, "provision.json")}
	if raw, err := os.ReadFile(s.dataPath); err == nil {
		_ = json.Unmarshal(raw, &s.config)
		s.provisioned = s.config.DeviceCode != ""
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/api/v1/identity", s.identity)
	mux.HandleFunc("/api/v1/pair", s.pair)
	mux.HandleFunc("/api/v1/provision", func(w http.ResponseWriter, r *http.Request) {
		s.provision(w, r, *mqttFallbackPlain)
	})
	mux.HandleFunc("/health", func(w http.ResponseWriter, _ *http.Request) {
		writeJSON(w, http.StatusOK, map[string]string{"status": "ok"})
	})

	server := &http.Server{Addr: *listen, Handler: mux, ReadHeaderTimeout: 5 * time.Second, ReadTimeout: 10 * time.Second, WriteTimeout: 10 * time.Second}
	cert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		log.Fatal(err)
	}
	server.TLSConfig = &tls.Config{MinVersion: tls.VersionTLS12, Certificates: []tls.Certificate{cert}}
	log.Printf("mock edge device %s listening on https://%s", deviceID, *listen)
	log.Printf("pairing code: %s; certificate fingerprint: %s", pairingCode, certFP)
	log.Printf("identity endpoint: https://<this-host>:%s/api/v1/identity", portOf(*listen))
	go s.serveDiscovery()
	if err := server.ListenAndServeTLS("", ""); err != nil && err != http.ErrServerClosed {
		log.Fatal(err)
	}
}

type discoveryRequest struct {
	Schema    string `json:"schema"`
	RequestID string `json:"request_id"`
	IssuedAt  int64  `json:"issued_at"`
	ExpiresAt int64  `json:"expires_at"`
	SystemID  string `json:"system_id"`
	Nonce     string `json:"nonce"`
}

func (s *state) serveDiscovery() {
	conn, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4zero, Port: 31001})
	if err != nil {
		log.Printf("UDP discovery disabled: %v", err)
		return
	}
	defer conn.Close()
	log.Printf("UDP discovery listening on 0.0.0.0:31001")
	buffer := make([]byte, 4*1024)
	for {
		n, remote, err := conn.ReadFromUDP(buffer)
		if err != nil {
			log.Printf("UDP discovery read failed: %v", err)
			return
		}
		if n == len(buffer) {
			continue
		}
		var req discoveryRequest
		if json.Unmarshal(buffer[:n], &req) != nil || req.Schema != "discovery.request.v1" || strings.TrimSpace(req.RequestID) == "" || strings.TrimSpace(req.Nonce) == "" || req.ExpiresAt <= time.Now().Unix() || req.ExpiresAt > time.Now().Add(2*time.Minute).Unix() {
			continue
		}
		ip := discoverySourceIP(remote.IP)
		if ip == nil {
			continue
		}
		configURL := fmt.Sprintf("https://%s:18443/api/v1/provision", ip.String())
		packet := map[string]string{
			"schema": "discovery.response.v1", "request_id": req.RequestID, "device_id": deviceID,
			"device_code": s.currentDeviceCode(), "device_type": "lpr-gate", "device_model": "MOCK-LPR-GATE",
			"firmware_version": "mock-1.0.0", "ip": ip.String(), "config_url": configURL,
			"public_key_fingerprint": s.publicKeyFP, "signing_public_key": base64.StdEncoding.EncodeToString(s.signingKey.Public().(ed25519.PublicKey)), "nonce": req.Nonce,
		}
		payload := strings.Join([]string{"discovery.response.v1", req.RequestID, deviceID, ip.String(), configURL, req.Nonce}, "\n")
		packet["signature"] = base64.StdEncoding.EncodeToString(ed25519.Sign(s.signingKey, []byte(payload)))
		body, _ := json.Marshal(packet)
		if _, err := conn.WriteToUDP(body, remote); err != nil {
			log.Printf("UDP discovery response failed: %v", err)
		}
	}
}

func (s *state) currentDeviceCode() string {
	s.mu.RLock()
	defer s.mu.RUnlock()
	return s.config.DeviceCode
}

func discoverySourceIP(remote net.IP) net.IP {
	conn, err := net.DialUDP("udp4", nil, &net.UDPAddr{IP: remote, Port: 31001})
	if err != nil {
		return nil
	}
	defer conn.Close()
	local, ok := conn.LocalAddr().(*net.UDPAddr)
	if !ok || local.IP == nil || local.IP.IsUnspecified() || local.IP.IsLoopback() {
		return nil
	}
	return local.IP.To4()
}

func (s *state) identity(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
		return
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	writeJSON(w, http.StatusOK, map[string]interface{}{
		"device_id": deviceID, "device_code": s.config.DeviceCode, "device_type": "lpr-gate",
		"device_model": "MOCK-LPR-GATE", "firmware_version": "mock-1.0.0",
		"certificate_fingerprint": s.certFP, "public_key_fingerprint": s.publicKeyFP,
		"pairing_required": !s.paired,
	})
}

func (s *state) pair(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
		return
	}
	var req struct {
		PairingCode string `json:"pairing_code"`
	}
	if json.NewDecoder(http.MaxBytesReader(w, r.Body, 4096)).Decode(&req) != nil || req.PairingCode != pairingCode {
		writeJSON(w, http.StatusForbidden, map[string]string{"error": "invalid pairing code"})
		return
	}
	s.mu.Lock()
	s.paired = true
	s.mu.Unlock()
	writeJSON(w, http.StatusOK, map[string]interface{}{"result": "paired", "device_id": deviceID, "expires_at": time.Now().Add(24 * time.Hour).Unix()})
}

func (s *state) provision(w http.ResponseWriter, r *http.Request, mqttFallbackPlain bool) {
	if r.Method != http.MethodPost {
		writeJSON(w, http.StatusMethodNotAllowed, map[string]string{"error": "method not allowed"})
		return
	}
	var req provisionRequest
	if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 64*1024)).Decode(&req); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid JSON"})
		return
	}
	s.mu.RLock()
	paired := s.paired
	s.mu.RUnlock()
	if !paired {
		writeJSON(w, http.StatusForbidden, map[string]string{"error": "device is not paired"})
		return
	}
	if req.Schema != "provision.request.v1" || req.DeviceID != deviceID || req.DeviceCode == "" || req.MQTT.Host == "" || req.MQTT.Port < 1 || req.MQTT.Port > 65535 || req.Route.ParkingLotID == 0 || req.Route.BoothID == 0 || req.Route.ChannelID == 0 {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid provision fields"})
		return
	}
	raw, _ := json.MarshalIndent(req, "", "  ")
	if err := atomicWrite(s.dataPath, raw); err != nil {
		writeJSON(w, http.StatusInternalServerError, map[string]string{"error": err.Error()})
		return
	}
	s.mu.Lock()
	s.config, s.provisioned = req, true
	old := s.client
	s.client = nil
	s.mu.Unlock()
	if old != nil {
		old.Disconnect(250)
	}
	go s.connectMQTT(req, mqttFallbackPlain)
	writeJSON(w, http.StatusOK, map[string]string{"result": "accepted", "request_id": req.RequestID, "message": "configuration saved; mqtt reconnecting"})
}

func (s *state) connectMQTT(req provisionRequest, mqttFallbackPlain bool) {
	server := fmt.Sprintf("tcp://%s:%d", req.MQTT.Host, req.MQTT.Port)
	useTLS := req.MQTT.TLS && !mqttFallbackPlain
	if useTLS {
		log.Printf("MQTT test stub received tls=true; embedded broker is plain TCP, connection may fail")
	}
	opts := mqtt.NewClientOptions().AddBroker(server).SetClientID(req.MQTT.ClientID).
		SetKeepAlive(30*time.Second).SetAutoReconnect(true).
		SetWill(topic(req.Route, "gate/lwt"), fmt.Sprintf(`{"schema":"gate.lwt.v1","device_code":%q}`, req.DeviceCode), 1, true)
	if useTLS {
		opts.SetTLSConfig(&tls.Config{MinVersion: tls.VersionTLS12, InsecureSkipVerify: true})
	}
	opts.SetOnConnectHandler(func(client mqtt.Client) {
		client.Subscribe(topic(req.Route, "gate/cmd"), 1, func(c mqtt.Client, msg mqtt.Message) { s.handleCommand(c, msg, req) })
		client.Publish(topic(req.Route, "gate/lwt"), 1, true, []byte{})
		payload, _ := json.Marshal(map[string]interface{}{"schema": "gate.state.v1", "device_code": req.DeviceCode, "state": "closed", "ts": time.Now().Unix()})
		client.Publish(topic(req.Route, "gate/state"), 1, true, payload)
		log.Printf("MQTT online; published gate.state.v1 to %s", topic(req.Route, "gate/state"))
	})
	client := mqtt.NewClient(opts)
	s.mu.Lock()
	s.client = client
	s.mu.Unlock()
	if token := client.Connect(); token.Wait() && token.Error() != nil {
		log.Printf("MQTT connect failed: %v", token.Error())
	}
}

func (s *state) handleCommand(client mqtt.Client, msg mqtt.Message, req provisionRequest) {
	var command struct{ Schema, CmdID, DeviceCode, Action string }
	if json.Unmarshal(msg.Payload(), &command) != nil || command.Schema != "gate.cmd.v1" || command.DeviceCode != req.DeviceCode || command.CmdID == "" {
		return
	}
	payload, _ := json.Marshal(map[string]string{"schema": "gate.ack.v1", "cmd_id": command.CmdID, "device_code": req.DeviceCode, "result": "success", "error": ""})
	client.Publish(topic(req.Route, "gate/cmd/ack"), 1, false, payload)
}

func topic(route routeConfig, suffix string) string {
	return fmt.Sprintf("parking/lot/%d/booth/%d/channel/%d/%s", route.ParkingLotID, route.BoothID, route.ChannelID, suffix)
}

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

func atomicWrite(path string, data []byte) error {
	tmp := path + ".tmp"
	if err := os.WriteFile(tmp, data, 0600); err != nil {
		return err
	}
	return os.Rename(tmp, path)
}

func portOf(addr string) string {
	_, port, err := net.SplitHostPort(addr)
	if err != nil {
		return "18443"
	}
	return port
}

func loadOrCreateCertificate(dir string) ([]byte, []byte, string, error) {
	certPath, keyPath := filepath.Join(dir, "device.crt.pem"), filepath.Join(dir, "device.key.pem")
	if certPEM, err := os.ReadFile(certPath); err == nil {
		if keyPEM, keyErr := os.ReadFile(keyPath); keyErr == nil {
			fp, fpErr := certificateFingerprint(certPEM)
			return certPEM, keyPEM, fp, fpErr
		}
	}
	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, "", err
	}
	serial, _ := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 120))
	now := time.Now()
	template := &x509.Certificate{SerialNumber: serial, Subject: pkix.Name{CommonName: deviceID}, NotBefore: now.Add(-time.Minute), NotAfter: now.AddDate(10, 0, 0), KeyUsage: x509.KeyUsageDigitalSignature, ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth}, BasicConstraintsValid: true, DNSNames: []string{"localhost"}}
	for _, ip := range localIPs() {
		template.IPAddresses = append(template.IPAddresses, ip)
	}
	der, err := x509.CreateCertificate(rand.Reader, template, template, publicKey, privateKey)
	if err != nil {
		return nil, nil, "", err
	}
	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: der})
	keyBytes, _ := x509.MarshalPKCS8PrivateKey(privateKey)
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: keyBytes})
	if err := os.WriteFile(certPath, certPEM, 0600); err != nil {
		return nil, nil, "", err
	}
	if err := os.WriteFile(keyPath, keyPEM, 0600); err != nil {
		return nil, nil, "", err
	}
	hash := sha256.Sum256(der)
	return certPEM, keyPEM, "SHA256:" + strings.ToUpper(hex.EncodeToString(hash[:])), nil
}

func certificateFingerprint(certPEM []byte) (string, error) {
	block, _ := pem.Decode(certPEM)
	if block == nil {
		return "", fmt.Errorf("invalid certificate PEM")
	}
	hash := sha256.Sum256(block.Bytes)
	return "SHA256:" + strings.ToUpper(hex.EncodeToString(hash[:])), nil
}

func loadOrCreateSigningKey(dir string) (ed25519.PublicKey, ed25519.PrivateKey, error) {
	path := filepath.Join(dir, "signing-key.bin")
	if raw, err := os.ReadFile(path); err == nil && len(raw) == ed25519.PrivateKeySize {
		privateKey := ed25519.PrivateKey(raw)
		return privateKey.Public().(ed25519.PublicKey), privateKey, nil
	}
	publicKey, privateKey, err := ed25519.GenerateKey(rand.Reader)
	if err != nil {
		return nil, nil, err
	}
	if err := os.WriteFile(path, privateKey, 0600); err != nil {
		return nil, nil, err
	}
	return publicKey, privateKey, nil
}

func localIPs() []net.IP {
	var result []net.IP
	interfaces, _ := net.Interfaces()
	for _, iface := range interfaces {
		addrs, _ := iface.Addrs()
		for _, addr := range addrs {
			var ip net.IP
			switch value := addr.(type) {
			case *net.IPNet:
				ip = value.IP
			case *net.IPAddr:
				ip = value.IP
			}
			if ip != nil && ip.To4() != nil && !ip.IsLoopback() {
				result = append(result, ip.To4())
			}
		}
	}
	return result
}