package main

import (
	"bytes"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"sync"
	"time"
)

type authType int

const (
	authBearer authType = iota
	authHeader
	authBasic
	authBasicPAT // Basic auth with "x-access-token:<token>" for git HTTP
)

// fallback is an alternate credential used when the route's primary TokenKey
// is not configured in secrets.json.
type fallback struct {
	TokenKey   string
	HeaderName string
}

// hostParam names a {placeholder} in a route's Upstream that is filled from a
// leading path segment (in HostParams order). Valid keeps a caller-supplied
// segment from widening the upstream to a different host.
type hostParam struct {
	Name  string
	Valid *regexp.Regexp
}

type route struct {
	Prefix     string
	Upstream   string
	TokenKey   string
	AuthType   authType
	HeaderName string
	Fallback   *fallback
	Timeout    time.Duration // zero means defaultTimeout
	HostParams []hostParam   // fill {name} placeholders in Upstream from leading path segments
}

const (
	defaultTimeout = 30 * time.Second
	// LLM generation runs far longer than a datasource query, and the client
	// timeout covers the whole streamed response, not just the headers.
	llmTimeout = 10 * time.Minute
)

// awsRegion matches an AWS region label (us-west-2, ap-southeast-1,
// us-gov-east-1). It is anchored and allows no dots, so a caller-supplied
// segment can never point the upstream at a different host: without it,
// /bedrock/s3 would resolve to bedrock-runtime.s3.amazonaws.com, an S3 bucket
// someone else may own, and the Bedrock key would be sent there.
var awsRegion = regexp.MustCompile(`^[a-z]{2}(-[a-z]+)+-\d{1,2}$`)

// o11yLabel matches an o11y region or vendor label (us-west-2, us-central1,
// aws, gcp, alicloud). Like awsRegion it starts with a letter and forbids
// dots, so a segment can never re-anchor the host outside
// www.ds.<region>.<vendor>.observability.tidbcloud.com.
var o11yLabel = regexp.MustCompile(`^[a-z][a-z0-9-]{0,31}$`)

// templated reports whether the route's upstream host is completed from
// leading path segments.
func (rt *route) templated() bool { return len(rt.HostParams) > 0 }

// fillHost peels one leading path segment per HostParam, validates each, and
// substitutes it into the upstream template -- returning the completed
// upstream and the remaining path to forward (which keeps any query string).
// When ok is false, bad holds the rejected segment verbatim for the error
// message only -- never build an upstream from it.
func (rt *route) fillHost(path string) (upstream, rest, bad string, ok bool) {
	upstream, rest = rt.Upstream, path
	for _, hp := range rt.HostParams {
		seg, tail, found := strings.Cut(strings.TrimPrefix(rest, "/"), "/")
		if !hp.Valid.MatchString(seg) {
			return "", "", seg, false
		}
		if found {
			rest = "/" + tail
		} else {
			rest = "/"
		}
		upstream = strings.Replace(upstream, "{"+hp.Name+"}", seg, 1)
	}
	return upstream, rest, "", true
}

// hostParamNames renders a route's host params as "<region>/<vendor>" for
// error messages.
func hostParamNames(rt *route) string {
	parts := make([]string, len(rt.HostParams))
	for i, hp := range rt.HostParams {
		parts[i] = "<" + hp.Name + ">"
	}
	return strings.Join(parts, "/")
}

var devopsFallback = &fallback{TokenKey: "DEVOPS_API_KEY", HeaderName: "X-API-KEY"}

var routes = []route{
	{Prefix: "/devops/prod", Upstream: "https://prod.devops.pingcap.com", TokenKey: "DEVOPS_API_TOKEN", AuthType: authHeader, HeaderName: "X-API-TOKEN", Fallback: devopsFallback},
	{Prefix: "/devops/staging", Upstream: "https://staging.devops.pingcap.com", TokenKey: "DEVOPS_API_TOKEN", AuthType: authHeader, HeaderName: "X-API-TOKEN", Fallback: devopsFallback},
	// Both o11y query hosts authenticate with the auth0 token; the plain
	// O11Y_TOKEN is only for the direct-mode data-export flow and is never
	// injected here (it is handed out via /o11y/credentials instead).
	//
	// ds host carries the cluster's region and vendor, filled from the two path
	// segments after the prefix (/o11y/ds/<region>/<vendor>/...). o11yLabel keeps
	// them from widening the host past *.observability.tidbcloud.com.
	{Prefix: "/o11y/ds", Upstream: "https://www.ds.{region}.{vendor}.observability.tidbcloud.com", TokenKey: "O11Y_AUTH0_TOKEN", AuthType: authBearer, HostParams: []hostParam{{"region", o11yLabel}, {"vendor", o11yLabel}}},
	// gs is always us-west-2/aws, so its host stays fixed.
	{Prefix: "/o11y/gs", Upstream: "http://www.gs.us-west-2.aws.observability.tidbcloud.com", TokenKey: "O11Y_AUTH0_TOKEN", AuthType: authBearer},
	{Prefix: "/clinic/prod", Upstream: "https://clinic.pingcap.com", TokenKey: "CLINIC_API_KEY", AuthType: authBearer},
	{Prefix: "/clinic/staging", Upstream: "https://staging-clinic.pingcap.com", TokenKey: "CLINIC_API_KEY", AuthType: authBearer},
	{Prefix: "/jira", Upstream: "https://tidb.atlassian.net", TokenKey: "JIRA_BASIC_AUTH_RAW", AuthType: authBasic},
	{Prefix: "/github", Upstream: "https://github.com", TokenKey: "GITHUB_PAT", AuthType: authBasicPAT},
	// DashScope: one route covers both the OpenAI-compatible (/compatible-mode/v1)
	// and the native (/api/v1) paths, since callers keep their own suffix.
	{Prefix: "/dashscope", Upstream: "https://dashscope-intl.aliyuncs.com", TokenKey: "DASHSCOPE_API_KEY", AuthType: authBearer, Timeout: llmTimeout},
	// Bedrock: the region lives in the hostname, so {region} is filled from the
	// path segment after the prefix (/bedrock/us-west-2/... ). awsRegion keeps
	// that segment from widening the host beyond bedrock-runtime.*.amazonaws.com.
	{Prefix: "/bedrock", Upstream: "https://bedrock-runtime.{region}.amazonaws.com", TokenKey: "AWS_BEARER_TOKEN_BEDROCK", AuthType: authBearer, Timeout: llmTimeout, HostParams: []hostParam{{"region", awsRegion}}},
}

// Clinic POST endpoints that are semantically read-only queries.
// Paths are matched after stripping the route prefix (e.g. /clinic/prod).
var clinicPostAllowlist = []string{
	"/data-proxy/query",
	"/data-proxy/batch-query",
}

var jiraPostAllowlist = []string{
	"/rest/api/3/search/jql",
}

// git smart HTTP: POST /git-upload-pack is read-only (clone/fetch),
// POST /git-receive-pack is write (push) and stays blocked.
var gitPostAllowSuffix = "/git-upload-pack"

// Clinic POST paths with wildcard segments: /data-proxy/api/v1/orgs/{orgID}/clusters/{clusterID}/...
// We match the static segments and allow any org/cluster IDs in between.
var clinicPostAllowPatterns = []struct{ prefix, suffix string }{
	{"/data-proxy/api/v1/orgs/", "/sql-queries"},
	{"/data-proxy/api/v1/orgs/", "/sql-query-batches"},
	{"/data-proxy/api/v1/orgs/", "/metrics/query"},
	{"/data-proxy/api/v1/orgs/", "/metrics/query-range"},
}

// LLM inference endpoints. These POSTs don't mutate any datasource, so they
// pass the same read-only guard as the query endpoints above.
var llmPostAllowlist = []string{
	"/compatible-mode/v1/chat/completions",             // DashScope, OpenAI-compatible
	"/compatible-mode/v1/embeddings",                   // DashScope, OpenAI-compatible
	"/api/v1/services/aigc/text-generation/generation", // DashScope, native
	"/openai/v1/chat/completions",                      // Bedrock, OpenAI-compatible
	"/apps/anthropic/v1/messages",                      // DashScope, Anthropic-compatible
	"/apps/anthropic/v1/messages/count_tokens",         // DashScope, Anthropic-compatible
}

// Bedrock native invoke: /model/{modelId}/invoke[-with-response-stream].
// modelId may be an inference-profile ARN, so match on the static segments.
// Note the two suffixes are disjoint: "/invoke-with-response-stream" does not
// end with "/invoke".
var bedrockPostAllowPatterns = []struct{ prefix, suffix string }{
	{"/model/", "/invoke"},
	{"/model/", "/invoke-with-response-stream"},
}

func isPostAllowed(upstreamPath string) bool {
	for _, p := range clinicPostAllowlist {
		if upstreamPath == p {
			return true
		}
	}
	for _, pat := range clinicPostAllowPatterns {
		if strings.HasPrefix(upstreamPath, pat.prefix) && strings.HasSuffix(upstreamPath, pat.suffix) {
			return true
		}
	}
	for _, p := range jiraPostAllowlist {
		if upstreamPath == p {
			return true
		}
	}
	for _, p := range llmPostAllowlist {
		if upstreamPath == p {
			return true
		}
	}
	for _, pat := range bedrockPostAllowPatterns {
		if strings.HasPrefix(upstreamPath, pat.prefix) && strings.HasSuffix(upstreamPath, pat.suffix) {
			return true
		}
	}
	if strings.HasSuffix(upstreamPath, gitPostAllowSuffix) {
		return true
	}
	return false
}

func isReadOnly(method, upstreamPath string) bool {
	if method == http.MethodGet || method == http.MethodHead {
		return true
	}
	if method == http.MethodPost {
		return isPostAllowed(upstreamPath)
	}
	return false
}

func matchRoute(path string) *route {
	for i := range routes {
		r := &routes[i]
		if path == r.Prefix || strings.HasPrefix(path, r.Prefix+"/") {
			return r
		}
	}
	return nil
}

// --- secrets store ---
// Reads secrets.json (named "secrets" instead of "config" so AI agents don't try to parse it).

type secrets struct {
	DevopsAPIToken string `json:"DEVOPS_API_TOKEN"`
	DevopsAPIKey   string `json:"DEVOPS_API_KEY"`
	ClinicAPIKey   string `json:"CLINIC_API_KEY"`
	JiraBasicAuth  string `json:"JIRA_BASIC_AUTH_RAW"`
	GitHubPAT      string `json:"GITHUB_PAT"`
	DashScopeKey   string `json:"DASHSCOPE_API_KEY"`
	BedrockToken   string `json:"AWS_BEARER_TOKEN_BEDROCK"`
	O11yEnabled    bool   `json:"O11Y_ENABLED"`
	// O11yExposeCreds opens GET /o11y/credentials, which hands out the
	// credentials for the direct-only o11y flows the proxy deliberately does not
	// proxy: O11Y_TOKEN (data export) and the S3 STS credentials. Off by default:
	// it is the one route that returns credentials rather than injecting them, so
	// enable it only when a caller genuinely needs direct export/S3 access. The
	// auth0 token (used for the proxied queries) and the PAT are never handed out.
	O11yExposeCreds bool `json:"O11Y_EXPOSE_CREDENTIALS"`
}

type secretsStore struct {
	mu   sync.RWMutex
	data secrets
	path string
}

func newSecretsStore(path string) *secretsStore {
	ss := &secretsStore{path: path}
	ss.load()
	go ss.watch()
	return ss
}

func (ss *secretsStore) load() {
	data, err := os.ReadFile(ss.path)
	if err != nil {
		slog.Warn("cannot read secrets", "path", ss.path, "err", err)
		return
	}
	var s secrets
	if err := json.Unmarshal(data, &s); err != nil {
		slog.Warn("cannot parse secrets", "path", ss.path, "err", err)
		return
	}
	ss.mu.Lock()
	ss.data = s
	ss.mu.Unlock()
	slog.Info("secrets loaded", "path", ss.path)
}

func (ss *secretsStore) watch() {
	var lastMod time.Time
	for {
		time.Sleep(5 * time.Second)
		info, err := os.Stat(ss.path)
		if err != nil {
			continue
		}
		if info.ModTime().After(lastMod) {
			lastMod = info.ModTime()
			ss.load()
		}
	}
}

func (ss *secretsStore) get(key string) (string, bool) {
	ss.mu.RLock()
	defer ss.mu.RUnlock()
	var v string
	switch key {
	case "DEVOPS_API_TOKEN":
		v = ss.data.DevopsAPIToken
	case "DEVOPS_API_KEY":
		v = ss.data.DevopsAPIKey
	case "CLINIC_API_KEY":
		v = ss.data.ClinicAPIKey
	case "JIRA_BASIC_AUTH_RAW":
		v = ss.data.JiraBasicAuth
	case "GITHUB_PAT":
		v = ss.data.GitHubPAT
	case "DASHSCOPE_API_KEY":
		v = ss.data.DashScopeKey
	case "AWS_BEARER_TOKEN_BEDROCK":
		v = ss.data.BedrockToken
	}
	return v, v != ""
}

func (ss *secretsStore) snapshot() secrets {
	ss.mu.RLock()
	defer ss.mu.RUnlock()
	return ss.data
}

// --- o11y token rotation ---

const (
	o11yAuthURL       = "https://www.ds.us-west-2.aws.observability.tidbcloud.com/client/auth"
	o11yTokenTTL      = 12 * time.Hour
	o11yRefreshBuffer = 1 * time.Hour
)

type o11yAuth struct {
	mu        sync.RWMutex
	token     string
	auth0     string
	awsAK     string
	awsSK     string
	awsST     string
	expiresAt time.Time
	secrets   *secretsStore
}

func newO11yAuth(ss *secretsStore) *o11yAuth {
	oa := &o11yAuth{secrets: ss}
	oa.refresh()
	go oa.loop()
	return oa
}

func (oa *o11yAuth) refresh() {
	pat, ok := oa.secrets.get("GITHUB_PAT")
	if !ok {
		slog.Warn("o11y: GITHUB_PAT not configured, skipping token refresh")
		return
	}

	body, _ := json.Marshal(map[string]string{"github_token": pat, "provider": "aws"})
	resp, err := (&http.Client{Timeout: 10 * time.Second}).Post(
		o11yAuthURL, "application/json", bytes.NewReader(body),
	)
	if err != nil {
		slog.Error("o11y: auth request failed", "err", err)
		return
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		data, _ := io.ReadAll(resp.Body)
		slog.Error("o11y: auth failed", "status", resp.StatusCode, "body", string(data))
		return
	}

	var result struct {
		Token           string `json:"token"`
		Auth0Token      string `json:"auth0_token"`
		AccessKeyID     string `json:"access_key_id"`
		SecretAccessKey string `json:"secret_access_key"`
		SessionToken    string `json:"session_token"`
	}
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		slog.Error("o11y: decode auth response failed", "err", err)
		return
	}
	if result.Token == "" || result.Auth0Token == "" {
		slog.Error("o11y: auth response missing tokens")
		return
	}

	oa.mu.Lock()
	oa.token = result.Token
	oa.auth0 = result.Auth0Token
	// The same /client/auth call returns short-lived AWS STS credentials scoped
	// to the caller's o11y S3 bucket. They are served (never the PAT) to callers
	// that need direct S3 access; see serveCredentials.
	oa.awsAK = result.AccessKeyID
	oa.awsSK = result.SecretAccessKey
	oa.awsST = result.SessionToken
	oa.expiresAt = time.Now().Add(o11yTokenTTL)
	oa.mu.Unlock()

	slog.Info("o11y: tokens refreshed", "expires_at", oa.expiresAt.Format(time.RFC3339))
}

func (oa *o11yAuth) loop() {
	for {
		time.Sleep(5 * time.Minute)

		oa.mu.RLock()
		needsRefresh := oa.token == "" || time.Until(oa.expiresAt) < o11yRefreshBuffer
		oa.mu.RUnlock()

		if needsRefresh {
			oa.refresh()
		}
	}
}

func (oa *o11yAuth) get(key string) (string, bool) {
	oa.mu.RLock()
	defer oa.mu.RUnlock()
	switch key {
	case "O11Y_TOKEN":
		return oa.token, oa.token != ""
	case "O11Y_AUTH0_TOKEN":
		return oa.auth0, oa.auth0 != ""
	}
	return "", false
}

// awsCreds returns the STS credentials from the last successful auth, if any.
func (oa *o11yAuth) awsCreds() (ak, sk, st string, ok bool) {
	oa.mu.RLock()
	defer oa.mu.RUnlock()
	if oa.awsAK == "" {
		return "", "", "", false
	}
	return oa.awsAK, oa.awsSK, oa.awsST, true
}

// --- proxy ---

type proxy struct {
	secrets *secretsStore
	o11y    *o11yAuth
}

func matchService(path string) string {
	parts := strings.SplitN(strings.TrimPrefix(path, "/"), "/", 3)
	if len(parts) >= 1 {
		return "/" + parts[0]
	}
	return ""
}

func serviceEnvs(service string) []string {
	var envs []string
	for _, r := range routes {
		if strings.HasPrefix(r.Prefix, service+"/") {
			envs = append(envs, r.Prefix)
		}
	}
	return envs
}

func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path == "/healthz" {
		w.WriteHeader(http.StatusOK)
		return
	}

	if r.URL.Path == "/o11y/credentials" {
		p.serveCredentials(w, r)
		return
	}

	if r.URL.Path == "/" {
		var available []string
		for _, rt := range routes {
			if _, _, ok := p.resolveToken(&rt); ok {
				available = append(available, rt.Prefix)
			}
		}
		writeJSON(w, http.StatusOK, map[string]any{"routes": available})
		return
	}

	rt := matchRoute(r.URL.Path)
	if rt == nil {
		svc := matchService(r.URL.Path)
		if envs := serviceEnvs(svc); len(envs) > 0 {
			writeJSON(w, http.StatusNotFound, map[string]any{
				"error":     fmt.Sprintf("[datasource-proxy] unknown route %q", r.URL.Path),
				"available": envs,
			})
		} else {
			services := make([]string, len(routes))
			for i, r := range routes {
				services[i] = r.Prefix
			}
			writeJSON(w, http.StatusNotFound, map[string]any{
				"error":     fmt.Sprintf("[datasource-proxy] unknown route %q", r.URL.Path),
				"available": services,
			})
		}
		return
	}

	upstreamPath := strings.TrimPrefix(r.URL.Path, rt.Prefix)
	if upstreamPath == "" || upstreamPath[0] != '/' {
		upstreamPath = "/" + upstreamPath
	}

	// Templated host segments (region/vendor) are part of the address, not of
	// the upstream path, so they come off before the allowlist sees the path.
	if rt.templated() {
		_, rest, bad, ok := rt.fillHost(upstreamPath)
		if !ok {
			writeJSON(w, http.StatusNotFound, map[string]any{
				"error": fmt.Sprintf("[datasource-proxy] %q is not a valid host segment; expected %s/%s/...", bad, rt.Prefix, hostParamNames(rt)),
			})
			return
		}
		upstreamPath = rest
	}

	if !isReadOnly(r.Method, upstreamPath) {
		slog.Warn("method not allowed", "method", r.Method, "path", r.URL.Path)
		writeJSON(w, http.StatusMethodNotAllowed, map[string]string{
			"error": "[datasource-proxy] method not allowed for this path",
		})
		return
	}

	p.forward(w, r, rt)
}

// resolveToken returns the credential to send and the header name it belongs
// under (the fallback credential may use a different header, e.g. devops
// X-API-KEY vs X-API-TOKEN).
// serveCredentials hands out the credentials for the direct-only o11y flows the
// proxy does not proxy -- O11Y_TOKEN (data export) and the S3 STS credentials --
// so a caller can run them without ever holding the GitHub PAT. The auth0 token
// (used for the proxied queries) and the PAT never leave the process.
func (p *proxy) serveCredentials(w http.ResponseWriter, r *http.Request) {
	// Off by default: when the operator hasn't opted in, the endpoint reads as
	// absent so it neither serves creds nor advertises the capability.
	if !p.secrets.snapshot().O11yExposeCreds {
		writeJSON(w, http.StatusNotFound, map[string]string{
			"error": "[datasource-proxy] /o11y/credentials is disabled; set O11Y_EXPOSE_CREDENTIALS=true in secrets.json",
		})
		return
	}
	if r.Method != http.MethodGet {
		writeJSON(w, http.StatusMethodNotAllowed, map[string]string{
			"error": "[datasource-proxy] method not allowed for this path",
		})
		return
	}
	if p.o11y == nil {
		writeJSON(w, http.StatusServiceUnavailable, map[string]string{
			"error": "[datasource-proxy] o11y rotation disabled; set O11Y_ENABLED=true in secrets.json",
		})
		return
	}
	token, ok := p.o11y.get("O11Y_TOKEN")
	if !ok {
		writeJSON(w, http.StatusServiceUnavailable, map[string]string{
			"error": "[datasource-proxy] o11y credentials not available yet",
		})
		return
	}
	out := map[string]string{"o11y_token": token}
	if ak, sk, st, ok := p.o11y.awsCreds(); ok {
		out["access_key_id"] = ak
		out["secret_access_key"] = sk
		out["session_token"] = st
	}
	writeJSON(w, http.StatusOK, out)
}

func (p *proxy) resolveToken(rt *route) (token string, headerName string, ok bool) {
	if p.o11y != nil {
		if token, ok := p.o11y.get(rt.TokenKey); ok {
			return token, rt.HeaderName, true
		}
	}
	if token, ok := p.secrets.get(rt.TokenKey); ok {
		return token, rt.HeaderName, true
	}
	if fb := rt.Fallback; fb != nil {
		if token, ok := p.secrets.get(fb.TokenKey); ok {
			return token, fb.HeaderName, true
		}
	}
	return "", "", false
}

func (p *proxy) forward(w http.ResponseWriter, r *http.Request, rt *route) {
	forwardPath := r.URL.RequestURI()
	forwardPath = strings.TrimPrefix(forwardPath, rt.Prefix)
	if forwardPath == "" || forwardPath[0] != '/' {
		forwardPath = "/" + forwardPath
	}

	upstreamURL := rt.Upstream
	if rt.templated() {
		// ServeHTTP already validated this; repeat rather than trust a caller
		// of forward to have done it.
		filled, rest, _, ok := rt.fillHost(forwardPath)
		if !ok {
			writeJSON(w, http.StatusNotFound, map[string]string{
				"error": "[datasource-proxy] missing or malformed host segment",
			})
			return
		}
		upstreamURL = filled
		forwardPath = rest
	}

	upstream, _ := url.Parse(upstreamURL)

	target := upstream.Scheme + "://" + upstream.Host + strings.TrimSuffix(upstream.Path, "/") + forwardPath

	var body io.Reader
	if r.Method == http.MethodPost {
		body = r.Body
	}

	proxyReq, err := http.NewRequestWithContext(r.Context(), r.Method, target, body)
	if err != nil {
		writeJSON(w, http.StatusBadGateway, map[string]string{"error": err.Error()})
		return
	}

	for k, vv := range r.Header {
		switch strings.ToLower(k) {
		case "host", "authorization", "x-api-token", "x-api-key":
			continue
		}
		for _, v := range vv {
			proxyReq.Header.Add(k, v)
		}
	}

	token, headerName, ok := p.resolveToken(rt)
	if !ok {
		slog.Warn("token not configured", "token_key", rt.TokenKey)
		writeJSON(w, http.StatusBadGateway, map[string]string{
			"error": fmt.Sprintf("[datasource-proxy] token %q not configured in secrets.json", rt.TokenKey),
		})
		return
	}
	switch rt.AuthType {
	case authBearer:
		proxyReq.Header.Set("Authorization", "Bearer "+token)
	case authHeader:
		proxyReq.Header.Set(headerName, token)
	case authBasic:
		proxyReq.Header.Set("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte(token)))
	case authBasicPAT:
		proxyReq.Header.Set("Authorization", "Basic "+base64.StdEncoding.EncodeToString([]byte("x-access-token:"+token)))
	}

	timeout := rt.Timeout
	if timeout == 0 {
		timeout = defaultTimeout
	}

	start := time.Now()
	client := &http.Client{
		Timeout: timeout,
		CheckRedirect: func(*http.Request, []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}

	resp, err := client.Do(proxyReq)
	duration := time.Since(start)
	if err != nil {
		slog.Error("upstream request failed", "method", r.Method, "path", r.URL.RequestURI(), "upstream", target, "duration", duration, "err", err)
		writeJSON(w, http.StatusBadGateway, map[string]string{"error": "[datasource-proxy] upstream error: " + err.Error()})
		return
	}
	defer resp.Body.Close()

	slog.Info("request", "method", r.Method, "path", r.URL.RequestURI(), "upstream", target, "status", resp.StatusCode, "duration", duration)

	for k, vv := range resp.Header {
		for _, v := range vv {
			w.Header().Add(k, v)
		}
	}
	w.WriteHeader(resp.StatusCode)
	copyStream(w, resp.Body)
}

// copyStream forwards the body, flushing after every chunk. Without the flush
// a streamed response (SSE from the OpenAI-compatible endpoints, AWS
// event-stream from Bedrock invoke-with-response-stream) sits in the
// ResponseWriter buffer and reaches the client as one late blob.
func copyStream(w http.ResponseWriter, body io.Reader) {
	rc := http.NewResponseController(w)
	buf := make([]byte, 32*1024)
	for {
		n, readErr := body.Read(buf)
		if n > 0 {
			if _, err := w.Write(buf[:n]); err != nil {
				return
			}
			rc.Flush() // fails only if the writer can't flush; nothing to do about it
		}
		if readErr != nil {
			return
		}
	}
}

func writeJSON(w http.ResponseWriter, status int, v any) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(status)
	json.NewEncoder(w).Encode(v)
}

func main() {
	var (
		listenAddr string
		tokensPath string
	)
	flag.StringVar(&listenAddr, "listen", "127.0.0.1:8990", "listen address")
	flag.StringVar(&tokensPath, "secrets", "", "path to secrets.json")
	flag.Parse()

	if tokensPath == "" {
		tokensPath = os.Getenv("DS_PROXY_SECRETS")
	}
	if tokensPath == "" {
		exe, _ := os.Executable()
		dir := filepath.Dir(exe)
		if dir == "" || dir == "." {
			dir, _ = os.Getwd()
		}
		tokensPath = filepath.Join(dir, "secrets.json")
	}

	ss := newSecretsStore(tokensPath)

	var o11y *o11yAuth
	if ss.snapshot().O11yEnabled {
		o11y = newO11yAuth(ss)
	} else {
		slog.Info("o11y token rotation disabled (set O11Y_ENABLED=true in secrets.json to enable)")
	}
	p := &proxy{secrets: ss, o11y: o11y}

	slog.Info("server starting", "addr", listenAddr, "secrets", tokensPath)
	for _, r := range routes {
		slog.Info("route registered", "prefix", r.Prefix, "upstream", r.Upstream, "token_key", r.TokenKey)
	}

	if err := http.ListenAndServe(listenAddr, p); err != nil { // nosemgrep: go.lang.security.audit.net.use-tls.use-tls
		slog.Error("server error", "err", err)
		os.Exit(1)
	}
}