package main

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"time"
)

// InspectionRequest is sent to the inspection endpoint.
type InspectionRequest struct {
	Method        string              `json:"method"`
	URL           string              `json:"url"`
	Host          string              `json:"host"`
	RemoteAddr    string              `json:"remote_addr"`
	ContentType   string              `json:"content_type"`
	ContentLength int64               `json:"content_length"`
	Headers       map[string][]string `json:"headers"`
	BodySample    []byte              `json:"body_sample"` // first N bytes
}

// InspectionResponse is returned by the inspection endpoint.
type InspectionResponse struct {
	Verdict string `json:"verdict"` // "allow" or "block"
	Reason  string `json:"reason,omitempty"`
	Score   int    `json:"score,omitempty"`
}

const maxBodySample = 4096 // only send first 4KB to inspection endpoint

// inspect sends a synchronous inspection request to the inspection endpoint.
// Returns nil if allowed, or an error describing the block reason.
func inspect(r *http.Request) (*InspectionResponse, error) {
	c := getConfig()

	// Sample the first N bytes of the body.
	var bodySample []byte
	if r.Body != nil {
		bodySample = make([]byte, maxBodySample)
		n, _ := io.ReadFull(r.Body, bodySample)
		bodySample = bodySample[:n]
		// Restore body for downstream use.
		r.Body = io.NopCloser(io.MultiReader(
			bytes.NewReader(bodySample),
			r.Body,
		))
	}

	inspReq := InspectionRequest{
		Method:        r.Method,
		URL:           r.URL.String(),
		Host:          r.Host,
		RemoteAddr:    r.RemoteAddr,
		ContentType:   r.Header.Get("Content-Type"),
		ContentLength: r.ContentLength,
		Headers:       map[string][]string(r.Header),
		BodySample:    bodySample,
	}

	body, err := json.Marshal(inspReq)
	if err != nil {
		return nil, fmt.Errorf("marshal inspection request: %w", err)
	}

	timeout := time.Duration(c.TimeoutMs) * time.Millisecond
	client := &http.Client{Timeout: timeout}

	resp, err := client.Post(c.NodeURL+"/inspect", "application/json", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("inspection node unreachable: %w", err)
	}
	defer resp.Body.Close()

	var inspResp InspectionResponse
	if err := json.NewDecoder(resp.Body).Decode(&inspResp); err != nil {
		return nil, fmt.Errorf("decode inspection response: %w", err)
	}

	log.Printf("WAF: %s %s → %s (score=%d reason=%s)",
		r.Method, r.URL.Path, inspResp.Verdict, inspResp.Score, inspResp.Reason)

	return &inspResp, nil
}