package main

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

	"apiSign/internal/config"
	"apiSign/internal/crypto"
	"apiSign/internal/proto"
	"apiSign/internal/replay"
)

func main() {
	cfgPath := flag.String("config", "./demo/agent.json", "agent config path")
	verbose := flag.Bool("verbose", false, "log verbose request/crypto details (encrypted envelope, headers, plaintext preview)")
	flag.Parse()

	cfg, err := config.ReadAgentConfig(*cfgPath)
	if err != nil {
		log.Fatal(err)
	}

	agentSignPriv, err := crypto.ParseEd25519PrivateKey(cfg.SigningKey.PrivateB64)
	if err != nil {
		log.Fatal(err)
	}
	agentEncPriv, err := crypto.ParseX25519PrivateKey(cfg.DecryptKey.PrivateB64)
	if err != nil {
		log.Fatal(err)
	}
	kongSignPub, err := crypto.ParseEd25519PublicKey(cfg.KongSignPub.PublicB64)
	if err != nil {
		log.Fatal(err)
	}
	kongEncPub, err := crypto.ParseX25519PublicKey(cfg.KongEncPub.PublicB64)
	if err != nil {
		log.Fatal(err)
	}

	store := replay.NewMemoryStore()

	mux := http.NewServeMux()
	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write([]byte("ok"))
	})

	mux.HandleFunc("/v1/execute", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}

		if *verbose {
			log.Printf("agent recv %s %s from=%s", r.Method, r.URL.RequestURI(), r.RemoteAddr)
			log.Printf("agent recv headers proto=%q senderKid=%q recipientKid=%q ts=%q exp=%q nonce=%q origCT=%q sig=%q",
				r.Header.Get(proto.HeaderProto),
				r.Header.Get(proto.HeaderSenderKID),
				r.Header.Get(proto.HeaderRecipientKID),
				r.Header.Get(proto.HeaderTimestampMs),
				r.Header.Get(proto.HeaderExpiryMs),
				r.Header.Get(proto.HeaderNonce),
				r.Header.Get(proto.HeaderOrigCT),
				previewString(r.Header.Get(proto.HeaderSignature), 32),
			)
		}

		if r.ContentLength > cfg.Limits.MaxBodyBytes {
			http.Error(w, "body too large", http.StatusRequestEntityTooLarge)
			return
		}
		r.Body = http.MaxBytesReader(w, r.Body, cfg.Limits.MaxBodyBytes)
		body, err := io.ReadAll(r.Body)
		if err != nil {
			http.Error(w, "read body failed", http.StatusBadRequest)
			return
		}
		if *verbose {
			log.Printf("agent recv raw encrypted envelope bytes=%d body=%s", len(body), previewBytes(body, 4096))
		}

		// Verify signature
		bodyHash := crypto.B64Encode(crypto.SHA256(body))
		sigInput := proto.BuildSigInputFromRequest(r, bodyHash)
		if *verbose {
			log.Printf("agent computed bodySha256=%s", bodyHash)
			log.Printf("agent sig canonical=%q", sigInput.CanonicalString())
			log.Printf("agent aead aad=%q", sigInput.AADString())
		}
		sigB64 := r.Header.Get(proto.HeaderSignature)
		sig, err := crypto.B64Decode(sigB64)
		if err != nil {
			http.Error(w, "bad signature encoding", http.StatusBadRequest)
			return
		}
		if !crypto.VerifyEd25519(kongSignPub, []byte(sigInput.CanonicalString()), sig) {
			http.Error(w, "signature verify failed", http.StatusUnauthorized)
			return
		}

		nowMs := time.Now().UnixMilli()
		tsMs, expMs, err := parseTimeHeaders(r)
		if err != nil {
			http.Error(w, err.Error(), http.StatusBadRequest)
			return
		}
		if !withinSkew(nowMs, tsMs, cfg.Limits.MaxClockSkewMs) {
			http.Error(w, "timestamp skew too large", http.StatusUnauthorized)
			return
		}
		if nowMs > expMs {
			http.Error(w, "request expired", http.StatusUnauthorized)
			return
		}

		if cfg.Limits.AntiReplayEnabled {
			nonce := r.Header.Get(proto.HeaderNonce)
			if nonce == "" {
				http.Error(w, "missing nonce", http.StatusBadRequest)
				return
			}
			key := r.Header.Get(proto.HeaderSenderKID) + ":" + nonce
			if store.SeenBefore(key, expMs) {
				http.Error(w, "replay detected", http.StatusUnauthorized)
				return
			}
		}

		var env proto.Envelope
		if err := json.Unmarshal(body, &env); err != nil {
			http.Error(w, "bad envelope", http.StatusBadRequest)
			return
		}
		epK, err := crypto.B64Decode(env.EPK)
		if err != nil {
			http.Error(w, "bad epk", http.StatusBadRequest)
			return
		}
		salt, err := crypto.B64Decode(env.Salt)
		if err != nil {
			http.Error(w, "bad salt", http.StatusBadRequest)
			return
		}
		iv, err := crypto.B64Decode(env.IV)
		if err != nil {
			http.Error(w, "bad iv", http.StatusBadRequest)
			return
		}
		ct, err := crypto.B64Decode(env.CT)
		if err != nil {
			http.Error(w, "bad ct", http.StatusBadRequest)
			return
		}

		aad := []byte(sigInput.AADString())
		pt, err := crypto.OpenFromX25519(agentEncPriv, epK, salt, iv, ct, aad)
		if err != nil {
			http.Error(w, "decrypt failed", http.StatusUnauthorized)
			return
		}

		origCT := r.Header.Get(proto.HeaderOrigCT)
		log.Printf("agent decrypted payload content-type=%q bytes=%d", origCT, len(pt))
		if *verbose {
			log.Printf("agent decrypted plaintext (preview)=%s", previewBytes(pt, 4096))
		}
		log.Printf("agent mock exec input: %s", safeString(pt, 4096))

		// Mock execute result
		respObj := map[string]any{
			"ok":           true,
			"received_ct":  origCT,
			"received_len": len(pt),
			"echo":         safeString(pt, 1024),
			"ts_ms":        time.Now().UnixMilli(),
		}
		respPlain, _ := json.Marshal(respObj)

		// Seal response to Kong
		respSigInput := proto.SigInput{
			Proto:        cfg.ProtoHeader,
			Method:       http.MethodPost,
			PathQuery:    r.URL.RequestURI(),
			SenderKID:    cfg.SigningKey.KID,
			RecipientKID: cfg.KongEncPub.KID,
			TimestampMs:  strconv.FormatInt(time.Now().UnixMilli(), 10),
			ExpiryMs:     strconv.FormatInt(time.Now().UnixMilli()+cfg.Limits.DefaultTTLms, 10),
			Nonce:        r.Header.Get(proto.HeaderNonce),
			OrigCT:       "application/json",
		}
		seal, err := crypto.SealToX25519(kongEncPub, respPlain, []byte(respSigInput.AADString()))
		if err != nil {
			http.Error(w, "encrypt failed", http.StatusInternalServerError)
			return
		}
		respEnv := proto.Envelope{
			V:    1,
			EPK:  crypto.B64Encode(seal.EphemeralPublic),
			Salt: crypto.B64Encode(seal.Salt),
			IV:   crypto.B64Encode(seal.IV),
			CT:   crypto.B64Encode(seal.Ciphertext),
		}
		respBody, _ := json.Marshal(respEnv)
		if *verbose {
			log.Printf("agent resp encrypted envelope bytes=%d body=%s", len(respBody), previewBytes(respBody, 4096))
		}
		respBodyHash := crypto.B64Encode(crypto.SHA256(respBody))
		respSigInput.BodySHA256 = respBodyHash
		sigBytes := crypto.SignEd25519(agentSignPriv, []byte(respSigInput.CanonicalString()))

		w.Header().Set(proto.HeaderProto, cfg.ProtoHeader)
		w.Header().Set(proto.HeaderSenderKID, cfg.SigningKey.KID)
		w.Header().Set(proto.HeaderRecipientKID, cfg.KongEncPub.KID)
		w.Header().Set(proto.HeaderTimestampMs, respSigInput.TimestampMs)
		w.Header().Set(proto.HeaderExpiryMs, respSigInput.ExpiryMs)
		if cfg.Limits.AntiReplayEnabled {
			w.Header().Set(proto.HeaderNonce, respSigInput.Nonce)
		}
		w.Header().Set(proto.HeaderOrigCT, "application/json")
		w.Header().Set(proto.HeaderUpstreamStatus, "200")
		w.Header().Set(proto.HeaderUpstreamCT, "application/json")
		w.Header().Set(proto.HeaderSignature, crypto.B64Encode(sigBytes))
		w.Header().Set("Content-Type", "application/json")
		w.WriteHeader(http.StatusOK)
		_, _ = w.Write(respBody)
	})

	log.Printf("db-agent listening on %s", cfg.ListenAddr)
	log.Fatal(http.ListenAndServe(cfg.ListenAddr, mux))
}

func parseTimeHeaders(r *http.Request) (tsMs int64, expMs int64, err error) {
	tsStr := r.Header.Get(proto.HeaderTimestampMs)
	expStr := r.Header.Get(proto.HeaderExpiryMs)
	if tsStr == "" || expStr == "" {
		return 0, 0, fmt.Errorf("missing timestamp/expiry")
	}
	tsMs, err = strconv.ParseInt(tsStr, 10, 64)
	if err != nil {
		return 0, 0, fmt.Errorf("bad timestamp")
	}
	expMs, err = strconv.ParseInt(expStr, 10, 64)
	if err != nil {
		return 0, 0, fmt.Errorf("bad expiry")
	}
	return tsMs, expMs, nil
}

func withinSkew(nowMs, tsMs, maxSkewMs int64) bool {
	d := nowMs - tsMs
	if d < 0 {
		d = -d
	}
	return d <= maxSkewMs
}

func safeString(b []byte, max int) string {
	s := strings.ReplaceAll(string(b), "\n", "\\n")
	if len(s) > max {
		return s[:max] + "..."
	}
	return s
}

func previewBytes(b []byte, max int) string {
	if len(b) <= max {
		return safeString(b, max)
	}
	return safeString(b[:max], max) + fmt.Sprintf("...(truncated, total=%d)", len(b))
}

func previewString(s string, max int) string {
	if len(s) <= max {
		return s
	}
	return s[:max] + "..."
}