package main

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

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

func main() {
	cfgPath := flag.String("config", "./demo/kong.json", "kong config path")
	dumpHTTP := flag.Bool("dump-http", false, "log outbound Kong->Agent request in raw HTTP format (BurpSuite-like)")
	flag.Parse()

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

	kongSignPriv, err := crypto.ParseEd25519PrivateKey(cfg.SigningKey.PrivateB64)
	if err != nil {
		log.Fatal(err)
	}
	kongDecPriv, err := crypto.ParseX25519PrivateKey(cfg.DecryptKey.PrivateB64)
	if err != nil {
		log.Fatal(err)
	}
	agentEncPub, err := crypto.ParseX25519PublicKey(cfg.AgentEncPub.PublicB64)
	if err != nil {
		log.Fatal(err)
	}
	agentSignPub, err := crypto.ParseEd25519PublicKey(cfg.AgentSignPub.PublicB64)
	if err != nil {
		log.Fatal(err)
	}

	client := &http.Client{Timeout: 10 * time.Second}

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

	mux.HandleFunc("/proxy", func(w http.ResponseWriter, r *http.Request) {
		if r.Method != http.MethodPost {
			http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
			return
		}
		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)
		origBody, err := io.ReadAll(r.Body)
		if err != nil {
			http.Error(w, "read body failed", http.StatusBadRequest)
			return
		}
		origCT := r.Header.Get("Content-Type")
		if origCT == "" {
			origCT = "application/octet-stream"
		}

		tsMs := time.Now().UnixMilli()
		expMs := tsMs + cfg.Limits.DefaultTTLms
		nonce := ""
		if cfg.Limits.AntiReplayEnabled {
			b, err := crypto.RandBytes(16)
			if err != nil {
				http.Error(w, "nonce gen failed", http.StatusInternalServerError)
				return
			}
			nonce = crypto.B64Encode(b)
		}

		// Build signature input for outbound request
		outSigInput := proto.SigInput{
			Proto:        cfg.ProtoHeader,
			Method:       http.MethodPost,
			PathQuery:    "/v1/execute",
			SenderKID:    cfg.SigningKey.KID,
			RecipientKID: cfg.AgentEncPub.KID,
			TimestampMs:  strconv.FormatInt(tsMs, 10),
			ExpiryMs:     strconv.FormatInt(expMs, 10),
			Nonce:        nonce,
			OrigCT:       origCT,
		}
		aad := []byte(outSigInput.AADString())

		seal, err := crypto.SealToX25519(agentEncPub, origBody, aad)
		if err != nil {
			http.Error(w, "encrypt failed", http.StatusInternalServerError)
			return
		}
		env := proto.Envelope{
			V:    1,
			EPK:  crypto.B64Encode(seal.EphemeralPublic),
			Salt: crypto.B64Encode(seal.Salt),
			IV:   crypto.B64Encode(seal.IV),
			CT:   crypto.B64Encode(seal.Ciphertext),
		}
		bodyBytes, _ := json.Marshal(env)
		outSigInput.BodySHA256 = crypto.B64Encode(crypto.SHA256(bodyBytes))
		sig := crypto.SignEd25519(kongSignPriv, []byte(outSigInput.CanonicalString()))

		outReq, err := http.NewRequest(http.MethodPost, cfg.AgentURL, bytes.NewReader(bodyBytes))
		if err != nil {
			http.Error(w, "build request failed", http.StatusInternalServerError)
			return
		}
		// Ensure we can dump the request (which consumes Body) without breaking the actual send.
		outReq.GetBody = func() (io.ReadCloser, error) {
			return io.NopCloser(bytes.NewReader(bodyBytes)), nil
		}
		outReq.Header.Set("Content-Type", "application/json")
		outReq.Header.Set(proto.HeaderProto, cfg.ProtoHeader)
		outReq.Header.Set(proto.HeaderSenderKID, cfg.SigningKey.KID)
		outReq.Header.Set(proto.HeaderRecipientKID, cfg.AgentEncPub.KID)
		outReq.Header.Set(proto.HeaderTimestampMs, outSigInput.TimestampMs)
		outReq.Header.Set(proto.HeaderExpiryMs, outSigInput.ExpiryMs)
		if cfg.Limits.AntiReplayEnabled {
			outReq.Header.Set(proto.HeaderNonce, nonce)
		}
		outReq.Header.Set(proto.HeaderOrigCT, origCT)
		outReq.Header.Set(proto.HeaderSignature, crypto.B64Encode(sig))

		// Demo visibility: log the outbound encrypted "packet" (headers + encrypted envelope body).
		// This intentionally does NOT log the plaintext payload.
		log.Printf("kong->agent request")
		log.Printf("  url: %s", cfg.AgentURL)
		log.Printf("  headers:")
		log.Printf("    %s: %q", proto.HeaderProto, outReq.Header.Get(proto.HeaderProto))
		log.Printf("    %s: %q", proto.HeaderSenderKID, outReq.Header.Get(proto.HeaderSenderKID))
		log.Printf("    %s: %q", proto.HeaderRecipientKID, outReq.Header.Get(proto.HeaderRecipientKID))
		log.Printf("    %s: %q", proto.HeaderTimestampMs, outReq.Header.Get(proto.HeaderTimestampMs))
		log.Printf("    %s: %q", proto.HeaderExpiryMs, outReq.Header.Get(proto.HeaderExpiryMs))
		if cfg.Limits.AntiReplayEnabled {
			log.Printf("    %s: %q", proto.HeaderNonce, outReq.Header.Get(proto.HeaderNonce))
		}
		log.Printf("    %s: %q", proto.HeaderOrigCT, outReq.Header.Get(proto.HeaderOrigCT))
		log.Printf("    %s: %q", proto.HeaderSignature, previewString(outReq.Header.Get(proto.HeaderSignature), 32))
		log.Printf("  body: encrypted envelope (%d bytes)", len(bodyBytes))
		log.Printf("%s", indentLines(prettyJSON(bodyBytes), "    "))

		if *dumpHTTP {
			// DumpRequestOut will read Body; restore it afterwards.
			dump, err := httputil.DumpRequestOut(outReq, true)
			if err != nil {
				log.Printf("kong->agent raw http dump failed: %v", err)
			} else {
				log.Printf("kong->agent raw http request (encrypted body)\n%s", string(dump))
			}
			if outReq.GetBody != nil {
				if rc, err := outReq.GetBody(); err == nil {
					outReq.Body = rc
				}
			}
		}

		resp, err := client.Do(outReq)
		if err != nil {
			http.Error(w, "agent request failed", http.StatusBadGateway)
			return
		}
		defer resp.Body.Close()

		respBody, err := io.ReadAll(resp.Body)
		if err != nil {
			http.Error(w, "read agent resp failed", http.StatusBadGateway)
			return
		}

		// Verify agent signature over its response
		respBodyHash := crypto.B64Encode(crypto.SHA256(respBody))
		respSigInput := proto.BuildSigInputFromRequest(resp.Request, respBodyHash)
		// BuildSigInputFromRequest expects an http.Request; we must reconstruct from agent response headers.
		respSigInput.Proto = resp.Header.Get(proto.HeaderProto)
		respSigInput.Method = http.MethodPost
		respSigInput.PathQuery = "/v1/execute"
		respSigInput.SenderKID = resp.Header.Get(proto.HeaderSenderKID)
		respSigInput.RecipientKID = resp.Header.Get(proto.HeaderRecipientKID)
		respSigInput.TimestampMs = resp.Header.Get(proto.HeaderTimestampMs)
		respSigInput.ExpiryMs = resp.Header.Get(proto.HeaderExpiryMs)
		respSigInput.Nonce = resp.Header.Get(proto.HeaderNonce)
		respSigInput.OrigCT = resp.Header.Get(proto.HeaderOrigCT)
		respSigInput.BodySHA256 = respBodyHash

		agentSigB64 := resp.Header.Get(proto.HeaderSignature)
		agentSig, err := crypto.B64Decode(agentSigB64)
		if err != nil {
			http.Error(w, "bad agent signature encoding", http.StatusBadGateway)
			return
		}
		if !crypto.VerifyEd25519(agentSignPub, []byte(respSigInput.CanonicalString()), agentSig) {
			http.Error(w, "agent signature verify failed", http.StatusBadGateway)
			return
		}

		var respEnv proto.Envelope
		if err := json.Unmarshal(respBody, &respEnv); err != nil {
			http.Error(w, "bad agent envelope", http.StatusBadGateway)
			return
		}
		epK, err := crypto.B64Decode(respEnv.EPK)
		if err != nil {
			http.Error(w, "bad agent epk", http.StatusBadGateway)
			return
		}
		salt, err := crypto.B64Decode(respEnv.Salt)
		if err != nil {
			http.Error(w, "bad agent salt", http.StatusBadGateway)
			return
		}
		iv, err := crypto.B64Decode(respEnv.IV)
		if err != nil {
			http.Error(w, "bad agent iv", http.StatusBadGateway)
			return
		}
		ct, err := crypto.B64Decode(respEnv.CT)
		if err != nil {
			http.Error(w, "bad agent ct", http.StatusBadGateway)
			return
		}
		aadResp := []byte(respSigInput.AADString())
		pt, err := crypto.OpenFromX25519(kongDecPriv, epK, salt, iv, ct, aadResp)
		if err != nil {
			http.Error(w, "agent resp decrypt failed", http.StatusBadGateway)
			return
		}

		upStatusStr := resp.Header.Get(proto.HeaderUpstreamStatus)
		upCT := resp.Header.Get(proto.HeaderUpstreamCT)
		upStatus := http.StatusOK
		if upStatusStr != "" {
			if v, err := strconv.Atoi(upStatusStr); err == nil {
				upStatus = v
			}
		}
		if upCT == "" {
			upCT = "application/octet-stream"
		}
		w.Header().Set("Content-Type", upCT)
		w.WriteHeader(upStatus)
		_, _ = w.Write(pt)
	})

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

func prettyJSON(b []byte) string {
	var out bytes.Buffer
	if err := json.Indent(&out, b, "", "  "); err != nil {
		return string(b)
	}
	return out.String()
}

func indentLines(s, prefix string) string {
	lines := strings.Split(s, "\n")
	for i := range lines {
		if lines[i] != "" {
			lines[i] = prefix + lines[i]
		}
	}
	return strings.Join(lines, "\n")
}

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