package aastro

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"math"
	"net"
	"net/http"
	"regexp"
	"slices"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"github.com/go-chi/chi/v5"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/trace"
	"go.uber.org/zap"

	"github.com/voidrunner3074/aastro/internal/circuitbreaker"
	"github.com/voidrunner3074/aastro/internal/metric"
)

const methodQuery = "QUERY"

type upstream interface {
	name() string
	call(ctx context.Context, original *http.Request, originalBody []byte) *upstreamResponse
}

type upstreamResponse struct {
	status  int
	headers http.Header
	body    []byte
	err     *upstreamError
}

type httpUpstream struct {
	cfg            upstreamConfig
	state          upstreamState
	circuitBreaker *circuitbreaker.CircuitBreaker
	metrics        *metric.Metrics
	log            *zap.Logger
	client         *http.Client
	streamClient   *http.Client
}

type upstreamConfig struct {
	id     string
	name   string
	hosts  []string
	path   string
	method string

	timeout        time.Duration
	forwardHeaders []string
	forwardQueries []string
	forwardParams  []string
	trustedProxies []*net.IPNet

	lbMode lbMode
	policy upstreamPolicy
}

type upstreamState struct {
	currentHostIdx    uint64
	activeConnections []int64
}

func (u *httpUpstream) name() string { return u.cfg.name }

// call executes the request with retries and circuit breaker protection,
// then validates the final response against the upstream's own policy.
func (u *httpUpstream) call(ctx context.Context, original *http.Request, originalBody []byte) *upstreamResponse {
	log := u.log.With(
		zap.String("upstream", u.cfg.name),
		zap.String("request_id", requestIDFromContext(original.Context())),
	)

	resp := u.callWithRetry(ctx, original, originalBody, log)

	u.updateCircuitBreaker(resp, log)

	// Policy is applied after the circuit breaker update intentionally:
	// a misconfigured requireBody should not cause the breaker to open.
	u.applyPolicy(ctx, resp)

	return resp
}

func (u *httpUpstream) callWithRetry(ctx context.Context, original *http.Request, originalBody []byte, log *zap.Logger) *upstreamResponse {
	retry := u.cfg.policy.retry

	// Same fallback as newRequest/proxy: an unset upstream method means
	// "use the flow's own method". shouldRetry must judge idempotency
	// against that effective method, not the raw (possibly empty) config
	// value - otherwise retries never fire for the common case of an
	// upstream that doesn't override the method.
	method := u.cfg.method
	if method == "" {
		method = original.Method
	}

	var resp *upstreamResponse

	for attempt := 0; attempt <= retry.maxRetries; attempt++ {
		if attempt > 0 {
			log.Debug("retrying upstream call",
				zap.Int("attempt", attempt),
				zap.Int("max_retries", retry.maxRetries),
				zap.Int("previous_status", resp.status),
			)

			u.metrics.IncUpstreamRetriesTotal(routeFromContext(ctx), u.cfg.name)
		}

		if err := ctx.Err(); err != nil {
			return &upstreamResponse{err: &upstreamError{kind: upstreamCanceled, err: err}}
		}

		if u.circuitBreaker != nil && !u.circuitBreaker.Allow() {
			log.Error("circuit breaker open, request denied")
			u.metrics.SetCircuitBreakerState(u.cfg.name, float64(u.circuitBreaker.State()))

			return &upstreamResponse{
				err: &upstreamError{
					kind: upstreamCircuitOpen,
					err:  errors.New("upstream circuit breaker is open"),
				},
			}
		}

		resp = u.doCall(ctx, original, originalBody, log)

		if attempt == retry.maxRetries || !u.shouldRetry(method, resp, retry) {
			break
		}

		if retry.backoffDelay > 0 {
			select {
			case <-time.After(retry.backoffDelay):
			case <-ctx.Done():
				return &upstreamResponse{err: &upstreamError{kind: upstreamCanceled, err: ctx.Err()}}
			}
		}
	}

	return resp
}

func (u *httpUpstream) shouldRetry(method string, resp *upstreamResponse, retry retryPolicy) bool {
	if !isIdempotent(method) {
		return false
	}

	if resp.err == nil {
		return slices.Contains(retry.retryOnStatuses, resp.status)
	}

	switch resp.err.kind.props().retry {
	case retryAlways:
		return true
	case retryByStatus:
		return slices.Contains(retry.retryOnStatuses, resp.status)
	case retryNever:
		return false
	default:
		return false
	}
}

func isIdempotent(method string) bool {
	switch method {
	case http.MethodGet, http.MethodHead, http.MethodOptions, http.MethodTrace, http.MethodPut, http.MethodDelete, methodQuery:
		return true
	default:
		return false
	}
}

func (u *httpUpstream) updateCircuitBreaker(resp *upstreamResponse, log *zap.Logger) {
	if u.circuitBreaker == nil {
		return
	}

	switch breakerOutcomeFor(resp.err) {
	case breakerFailure:
		log.Error("upstream request failed, recording circuit breaker failure")
		u.circuitBreaker.OnFailure()
	case breakerSuccess:
		u.circuitBreaker.OnSuccess()
	case breakerNoSignal:
		// The request never reached the upstream - denied by the breaker itself,
		// canceled by the client, or never built. Recording success here would
		// let an open breaker reset itself with the very request it rejected.
	}

	u.metrics.SetCircuitBreakerState(u.cfg.name, float64(u.circuitBreaker.State()))
}

// applyPolicy validates the response against the upstream's own policy rules.
// Violations are merged into resp.err so the original error kind is preserved.
func (u *httpUpstream) applyPolicy(ctx context.Context, resp *upstreamResponse) {
	if resp.err != nil && !isStatusFailure(resp.err.kind) {
		return
	}

	var errs []error

	if u.cfg.policy.requireBody && len(resp.body) == 0 {
		errs = append(errs, errors.New("empty body not allowed by upstream policy"))
	}

	if len(errs) == 0 {
		return
	}

	u.metrics.IncUpstreamErrorsTotal(routeFromContext(ctx), u.cfg.name, "policy_violation")

	combined := errors.Join(errs...)

	switch {
	case resp.err == nil:
		resp.err = &upstreamError{kind: upstreamPolicyViolation, err: combined}
	case resp.err.kind == upstreamClientError, resp.err.kind == upstreamRedirect:
		// A policy violation on top of an otherwise-propagatable client error/redirect
		// means the response isn't worth forwarding - downgrade to a policy violation (502).
		resp.err = &upstreamError{
			kind: upstreamPolicyViolation,
			err:  errors.Join(resp.err.err, combined),
		}
	default:
		// 5xx keeps bad_status: it must keep feeding the circuit breaker
		resp.err.err = errors.Join(resp.err.err, combined)
	}
}

func (u *httpUpstream) doCall(ctx context.Context, original *http.Request, originalBody []byte, log *zap.Logger) *upstreamResponse {
	ctx, cancel := context.WithTimeout(ctx, u.cfg.timeout)
	defer cancel()

	selectedHost := u.selectHost(log)

	if u.cfg.lbMode == lbModeLeastConns {
		atomic.AddInt64(&u.state.activeConnections[selectedHost], 1)
		defer atomic.AddInt64(&u.state.activeConnections[selectedHost], -1)
	}

	req, err := u.newRequest(ctx, original, originalBody, u.cfg.hosts[selectedHost], log)
	if err != nil {
		return &upstreamResponse{err: &upstreamError{kind: upstreamInternal, err: err}}
	}

	span := trace.SpanFromContext(ctx)
	span.SetAttributes(
		attribute.String("http.method", req.Method),
		attribute.String("http.url", req.URL.String()),
		attribute.String("server.address", req.URL.Host),
		attribute.String("aastro.upstream.host", u.cfg.hosts[selectedHost]),
	)

	otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(req.Header))

	httpResp, err := u.client.Do(req)
	if err != nil {
		log.Error("upstream request failed", zap.Error(err))
		span.RecordError(err)
		span.SetStatus(codes.Error, "upstream request failed")

		return &upstreamResponse{err: &upstreamError{kind: u.classifyDoError(err), err: err}}
	}
	defer httpResp.Body.Close()

	span.SetAttributes(attribute.Int("http.status_code", httpResp.StatusCode))

	body, uerr := u.readBody(ctx, httpResp.Body, log)
	if uerr != nil {
		span.RecordError(uerr.Unwrap())
		span.SetStatus(codes.Error, uerr.Error())

		return &upstreamResponse{status: httpResp.StatusCode, err: uerr}
	}

	resp := &upstreamResponse{
		status:  httpResp.StatusCode,
		headers: u.filterHeaders(httpResp.Header),
		body:    body,
	}

	if kind, failed := u.classifyStatus(httpResp.StatusCode); failed {
		if kind == upstreamBadStatus {
			log.Error("upstream returned server error", zap.Int("status_code", httpResp.StatusCode))
			span.SetStatus(codes.Error, http.StatusText(httpResp.StatusCode))
		} else {
			log.Debug(
				"upstream returned non-success status",
				zap.Int("status_code", httpResp.StatusCode),
				zap.String("kind", string(kind)),
			)
		}

		resp.err = &upstreamError{
			kind: kind,
			err:  fmt.Errorf("upstream returned %d", httpResp.StatusCode),
		}
	}

	return resp
}

func (u *httpUpstream) readBody(ctx context.Context, body io.ReadCloser, log *zap.Logger) ([]byte, *upstreamError) {
	reader := io.Reader(body)

	if u.cfg.policy.maxResponseBodySize > 0 {
		log.Debug("applying response body size limit policy", zap.Int64("limit", u.cfg.policy.maxResponseBodySize))

		// Read one byte past the limit: LimitReader signals exhaustion with io.EOF,
		// which ReadAll reports as success. The extra byte is the only way to tell
		// "body ended" from "body was truncated" - see the length check below
		reader = io.LimitReader(reader, u.cfg.policy.maxResponseBodySize+1)
	}

	data, err := io.ReadAll(reader)
	if err != nil {
		if ctx.Err() != nil {
			return nil, &upstreamError{kind: upstreamCanceled, err: ctx.Err()}
		}

		return nil, &upstreamError{kind: upstreamReadError, err: err}
	}

	if u.cfg.policy.maxResponseBodySize > 0 && int64(len(data)) > u.cfg.policy.maxResponseBodySize {
		return nil, &upstreamError{
			kind: upstreamBodyTooLarge,
			err:  fmt.Errorf("response size exceeds limit of %d bytes", u.cfg.policy.maxResponseBodySize),
		}
	}

	return data, nil
}

func (u *httpUpstream) filterHeaders(headers http.Header) http.Header {
	if len(u.cfg.policy.headerBlacklist) == 0 {
		return headers.Clone()
	}

	filtered := make(http.Header, len(headers))
	for header, values := range headers {
		if _, blocked := u.cfg.policy.headerBlacklist[header]; !blocked {
			filtered[header] = values
		}
	}

	return filtered
}

func (u *httpUpstream) classifyStatus(status int) (upstreamErrorKind, bool) {
	switch {
	case status >= http.StatusInternalServerError:
		return upstreamBadStatus, true
	case status >= http.StatusBadRequest:
		return upstreamClientError, true
	case status >= http.StatusMultipleChoices:
		return upstreamRedirect, true
	default:
		return "", false
	}
}

func (u *httpUpstream) classifyDoError(err error) upstreamErrorKind {
	switch {
	case errors.Is(err, context.DeadlineExceeded):
		return upstreamTimeout
	case errors.Is(err, context.Canceled):
		return upstreamCanceled
	default:
		return upstreamConnection
	}
}

func (u *httpUpstream) newRequest(ctx context.Context, original *http.Request, originalBody []byte, targetHost string, log *zap.Logger) (*http.Request, error) {
	path := expandPathParams(u.cfg.path, original)
	path = strings.TrimPrefix(path, "/")

	var hostPath string
	if strings.HasSuffix(targetHost, "/") {
		hostPath = targetHost + path
	} else {
		hostPath = targetHost + "/" + path
	}

	method := u.cfg.method
	if method == "" {
		method = original.Method
	}

	if method != http.MethodPost && method != http.MethodPut && method != http.MethodPatch && method != methodQuery {
		originalBody = nil
	}

	target, err := http.NewRequestWithContext(ctx, method, hostPath, bytes.NewReader(originalBody))
	if err != nil {
		return nil, err
	}

	u.resolveQueries(target, original)
	u.resolveHeaders(target, original, log)

	return target, nil
}

var pathParamRegexp = regexp.MustCompile(`\{([^}]+)\}`)

func expandPathParams(path string, req *http.Request) string {
	return pathParamRegexp.ReplaceAllStringFunc(path, func(match string) string {
		name := match[1 : len(match)-1]

		if value := chi.URLParam(req, name); value != "" {
			return value
		}

		return match
	})
}

func (u *httpUpstream) selectHost(log *zap.Logger) int64 {
	if len(u.cfg.hosts) == 1 {
		return 0
	}

	var selected int64

	switch u.cfg.lbMode {
	case lbModeRoundRobin:
		idx := atomic.AddUint64(&u.state.currentHostIdx, 1)

		// safe: idx % len is bounded by len(hosts) which is well within int64 range
		selected = int64(idx % uint64(len(u.cfg.hosts))) //nolint:gosec // bounded by slice length
	case lbModeLeastConns:
		var minConns int64 = math.MaxInt64

		for i := range u.cfg.hosts {
			if c := atomic.LoadInt64(&u.state.activeConnections[i]); c < minConns {
				minConns = c
				selected = int64(i)
			}
		}
	}

	log.Debug("host selected",
		zap.String("host", u.cfg.hosts[selected]),
		zap.String("upstream", u.cfg.name),
	)

	return selected
}

func (u *httpUpstream) resolveQueries(target, original *http.Request) {
	targetQ := target.URL.Query()
	originalQ := original.URL.Query()

	for _, pattern := range u.cfg.forwardQueries {
		if pattern == "*" {
			target.URL.RawQuery = originalQ.Encode()
			return
		}

		if v := originalQ.Get(pattern); v != "" {
			targetQ.Add(pattern, v)
		}
	}

	for _, param := range u.cfg.forwardParams {
		if param == "*" {
			if rctx := chi.RouteContext(original.Context()); rctx != nil {
				for i, key := range rctx.URLParams.Keys {
					targetQ.Set(key, rctx.URLParams.Values[i])
				}
			}

			target.URL.RawQuery = targetQ.Encode()

			return
		}

		if v := chi.URLParam(original, param); v != "" {
			targetQ.Set(param, v)
		}
	}

	target.URL.RawQuery = targetQ.Encode()
}

func (u *httpUpstream) resolveHeaders(target, original *http.Request, log *zap.Logger) {
	for _, pattern := range u.cfg.forwardHeaders {
		if pattern == "*" {
			target.Header = original.Header.Clone()
			break
		}

		if strings.HasSuffix(pattern, "*") {
			u.forwardHeadersByPrefix(target, original.Header, strings.TrimSuffix(pattern, "*"))
			continue
		}

		if v := original.Header.Get(pattern); v != "" {
			target.Header.Add(pattern, v)
		}
	}

	target.Header.Set("Content-Type", original.Header.Get("Content-Type"))

	clientIP := clientIPFromContext(original.Context())
	if clientIP == "" {
		clientIP = extractClientIP(original, u.cfg.trustedProxies)
	}

	parsedClientIP := net.ParseIP(clientIP)
	if parsedClientIP == nil {
		log.Warn("cannot parse client IP, skipping forward headers", zap.String("client_ip", clientIP))
		return
	}

	isTLS := original.TLS != nil
	port := u.resolvePort(u.effectiveHost(original), isTLS)
	proto := "http"

	if isTLS {
		proto = "https"
	}

	if !u.isTrustedProxy(parsedClientIP) {
		u.setUntrustedForwardingHeaders(original, target, clientIP, proto, port)
	} else {
		u.appendTrustedForwardingHeaders(original, target, clientIP, proto, port)
	}
}

func (u *httpUpstream) forwardHeadersByPrefix(target *http.Request, src http.Header, prefix string) {
	for name, values := range src {
		if strings.HasPrefix(name, prefix) {
			for _, v := range values {
				target.Header.Add(name, v)
			}
		}
	}
}

func (u *httpUpstream) setUntrustedForwardingHeaders(original, target *http.Request, clientIP, proto, port string) {
	host := u.effectiveHost(original)

	target.Header.Set("X-Forwarded-For", clientIP)
	target.Header.Set("X-Forwarded-Proto", proto)
	target.Header.Set("X-Forwarded-Host", host)
	target.Header.Set("X-Forwarded-Port", port)
	target.Header.Set("Forwarded", fmt.Sprintf("for=%s; proto=%s; host=%s", clientIP, proto, host))
}

func (u *httpUpstream) appendTrustedForwardingHeaders(original, target *http.Request, clientIP, proto, port string) {
	if xff := original.Header.Get("X-Forwarded-For"); xff != "" {
		target.Header.Set("X-Forwarded-For", xff+", "+clientIP)
	} else {
		target.Header.Set("X-Forwarded-For", clientIP)
	}

	if p := original.Header.Get("X-Forwarded-Proto"); p == "http" || p == "https" {
		proto = p
	}
	target.Header.Set("X-Forwarded-Proto", proto)

	host := u.effectiveHost(original)
	if h := original.Header.Get("X-Forwarded-Host"); h != "" {
		host = h
	}
	target.Header.Set("X-Forwarded-Host", host)

	if p := original.Header.Get("X-Forwarded-Port"); isValidPort(p) {
		target.Header.Set("X-Forwarded-Port", p)
	} else {
		target.Header.Set("X-Forwarded-Port", port)
	}

	hop := fmt.Sprintf("for=%s; proto=%s; host=%s", clientIP, proto, host)
	if fwd := original.Header.Get("Forwarded"); fwd != "" {
		target.Header.Set("Forwarded", fwd+", "+hop)
	} else {
		target.Header.Set("Forwarded", hop)
	}
}

func (u *httpUpstream) isTrustedProxy(ip net.IP) bool {
	for _, cidr := range u.cfg.trustedProxies {
		if cidr.Contains(ip) {
			return true
		}
	}

	return false
}

func (u *httpUpstream) resolvePort(host string, isTLS bool) string {
	_, port, err := net.SplitHostPort(host)
	if err != nil {
		if isTLS {
			return "443"
		}

		return "80"
	}

	return port
}

func (u *httpUpstream) effectiveHost(req *http.Request) string {
	if req.Host != "" {
		return req.Host
	}

	return req.URL.Host
}

func isValidPort(p string) bool {
	n, err := strconv.Atoi(p)
	return err == nil && n >= 1 && n <= 65535
}

// proxy implements proxyCapable for streaming flows.
func (u *httpUpstream) proxy(ctx context.Context, w http.ResponseWriter, original *http.Request) error {
	selectedHost := u.selectHost(u.log)
	host := u.cfg.hosts[selectedHost]

	path := expandPathParams(u.cfg.path, original)
	path = strings.TrimPrefix(path, "/")

	hostPath := host
	if !strings.HasSuffix(host, "/") {
		hostPath += "/"
	}
	hostPath += path

	method := u.cfg.method
	if method == "" {
		method = original.Method
	}

	req, err := http.NewRequestWithContext(ctx, method, hostPath, original.Body)
	if err != nil {
		return fmt.Errorf("build upstream request: %w", err)
	}

	req.ContentLength = original.ContentLength
	req.TransferEncoding = original.TransferEncoding

	u.resolveQueries(req, original)
	u.resolveHeaders(req, original, u.log)

	span := trace.SpanFromContext(ctx)
	span.SetAttributes(
		attribute.String("http.method", req.Method),
		attribute.String("http.url", req.URL.String()),
		attribute.String("server.address", req.URL.Host),
		attribute.String("aastro.upstream.host", host),
	)

	otel.GetTextMapPropagator().Inject(ctx, propagation.HeaderCarrier(req.Header))

	resp, err := u.streamClient.Do(req) // #nosec G704
	if err != nil {
		span.RecordError(err)
		span.SetStatus(codes.Error, "streaming upstream call failed")

		return fmt.Errorf("upstream call: %w", err)
	}
	defer resp.Body.Close()

	span.SetAttributes(attribute.Int("http.status_code", resp.StatusCode))
	if resp.StatusCode >= http.StatusInternalServerError {
		span.SetStatus(codes.Error, http.StatusText(resp.StatusCode))
	}

	for k, vv := range resp.Header {
		if _, skip := hopByHopHeaders[k]; skip {
			continue
		}

		for _, v := range vv {
			w.Header().Add(k, v)
		}
	}

	w.Header().Del("Content-Length")
	w.WriteHeader(resp.StatusCode)

	if err = streamCopy(w, resp.Body); err != nil {
		if errors.Is(err, context.Canceled) {
			span.SetAttributes(attribute.String("aastro.streaming.end_reason", "client_disconnect"))
			return nil
		}

		span.RecordError(err)
		span.SetStatus(codes.Error, "stream body failed")

		return fmt.Errorf("stream body: %w", err)
	}

	return nil
}