package kb

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"

	"github.com/tidbcloud/AgentTeam/internal/agsbridge"
)

func (s AGSHTTPRepositoryStore) doAGSAPI(ctx context.Context, method string, path string, payload any) (int, []byte, error) {
	const defaultMaxResponseBytes = 1 << 20
	return s.doAGSAPILimited(ctx, method, path, payload, defaultMaxResponseBytes)
}

func (s AGSHTTPRepositoryStore) doAGSAPILimited(ctx context.Context, method string, path string, payload any, maxResponseBytes int64) (int, []byte, error) {
	var body io.Reader
	if payload != nil {
		encoded, err := json.Marshal(payload)
		if err != nil {
			return 0, nil, err
		}
		body = bytes.NewReader(encoded)
	}
	if s.Handler != nil {
		request := httptest.NewRequest(method, path, body).WithContext(agsbridge.ContextWithSiteAdmin(ctx))
		if payload != nil {
			request.Header.Set("Content-Type", "application/json")
		}
		setAGSAuthHeader(request.Header, s.AuthHeader, s.AuthorizationHeader)
		recorder := newBoundedAGSResponseWriter(maxResponseBytes)
		s.Handler.ServeHTTP(recorder, request)
		if recorder.exceeded {
			return 0, nil, fmt.Errorf("AGS response exceeds limit")
		}
		return recorder.statusCode(), recorder.body.Bytes(), nil
	}
	baseURL := strings.TrimRight(strings.TrimSpace(s.BaseURL), "/")
	if baseURL == "" {
		return 0, nil, invalidArgument("AGS base URL is required")
	}
	parsed, err := url.Parse(baseURL)
	if err != nil || parsed.Scheme == "" || parsed.Host == "" {
		return 0, nil, invalidArgument("valid AGS base URL is required")
	}
	apiPath, err := url.Parse(path)
	if err != nil {
		return 0, nil, err
	}
	parsed.Path = strings.TrimRight(parsed.Path, "/") + apiPath.Path
	parsed.RawQuery = apiPath.RawQuery
	parsed.Fragment = ""
	request, err := http.NewRequestWithContext(ctx, method, parsed.String(), body)
	if err != nil {
		return 0, nil, err
	}
	if payload != nil {
		request.Header.Set("Content-Type", "application/json")
	}
	setAGSAuthHeader(request.Header, s.AuthHeader, s.AuthorizationHeader)
	response, err := http.DefaultClient.Do(request)
	if err != nil {
		return 0, nil, err
	}
	defer response.Body.Close()
	responseBody, err := io.ReadAll(io.LimitReader(response.Body, maxResponseBytes+1))
	if err != nil {
		return 0, nil, err
	}
	if int64(len(responseBody)) > maxResponseBytes {
		return 0, nil, fmt.Errorf("AGS response exceeds limit")
	}
	return response.StatusCode, responseBody, nil
}

type boundedAGSResponseWriter struct {
	header   http.Header
	body     bytes.Buffer
	status   int
	maxBytes int64
	exceeded bool
}

func newBoundedAGSResponseWriter(maxBytes int64) *boundedAGSResponseWriter {
	return &boundedAGSResponseWriter{header: make(http.Header), maxBytes: maxBytes}
}

func (w *boundedAGSResponseWriter) Header() http.Header { return w.header }

func (w *boundedAGSResponseWriter) WriteHeader(status int) {
	if w.status == 0 {
		w.status = status
	}
}

func (w *boundedAGSResponseWriter) Write(data []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	remaining := w.maxBytes - int64(w.body.Len())
	if remaining > 0 {
		keep := int64(len(data))
		if keep > remaining {
			keep = remaining
		}
		_, _ = w.body.Write(data[:keep])
	}
	if int64(len(data)) > remaining {
		w.exceeded = true
	}
	return len(data), nil
}

func (w *boundedAGSResponseWriter) statusCode() int {
	if w.status == 0 {
		return http.StatusOK
	}
	return w.status
}
