package provider

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"time"

	"aigw/internal/config"
	"aigw/internal/domain"
	"aigw/internal/telemetry"
)

type Result struct {
	Response *http.Response
	Route    domain.Route
	Attempts int
}

type Forwarder struct {
	client  *http.Client
	metrics *telemetry.Metrics
}

func New(cfg config.UpstreamHTTPConfig, metrics *telemetry.Metrics) *Forwarder {
	transport := &http.Transport{
		Proxy:                 http.ProxyFromEnvironment,
		DialContext:           (&net.Dialer{Timeout: 5 * time.Second, KeepAlive: 30 * time.Second}).DialContext,
		ForceAttemptHTTP2:     true,
		MaxIdleConns:          cfg.MaxIdleConnections,
		MaxIdleConnsPerHost:   cfg.MaxIdleConnectionsPerHost,
		IdleConnTimeout:       time.Duration(cfg.IdleConnectionTimeoutSecs) * time.Second,
		TLSHandshakeTimeout:   10 * time.Second,
		ResponseHeaderTimeout: time.Duration(cfg.ResponseHeaderTimeoutSecs) * time.Second,
		ExpectContinueTimeout: time.Second,
	}
	return &Forwarder{client: &http.Client{Transport: transport}, metrics: metrics}
}

func (f *Forwarder) Forward(ctx context.Context, protocol domain.Protocol, requestID string, originalBody []byte, sourceHeaders http.Header, routes []domain.Route) (Result, error) {
	var lastErr error
	for i, route := range routes {
		if err := ctx.Err(); err != nil {
			return Result{Attempts: i}, err
		}
		body, err := rewriteModel(originalBody, route.UpstreamModel)
		if err != nil {
			return Result{Attempts: i}, err
		}
		request, err := http.NewRequestWithContext(ctx, http.MethodPost, endpointURL(route.Provider.BaseURL, protocol), bytes.NewReader(body))
		if err != nil {
			return Result{Attempts: i}, fmt.Errorf("build upstream request: %w", err)
		}
		setHeaders(request.Header, sourceHeaders, route.Provider, protocol, requestID)
		f.metrics.UpstreamAttempt()
		response, err := f.client.Do(request)
		if err != nil {
			lastErr = err
			continue
		}
		attempts := i + 1
		if retryableStatus(response.StatusCode) && attempts < len(routes) {
			_, _ = io.CopyN(io.Discard, response.Body, 8<<10)
			_ = response.Body.Close()
			lastErr = fmt.Errorf("upstream %s returned %d", route.Provider.ID, response.StatusCode)
			continue
		}
		return Result{Response: response, Route: route, Attempts: attempts}, nil
	}
	if lastErr == nil {
		lastErr = errors.New("all upstream routes failed")
	}
	return Result{Attempts: len(routes)}, lastErr
}

func rewriteModel(body []byte, upstreamModel string) ([]byte, error) {
	var object map[string]json.RawMessage
	if err := json.Unmarshal(body, &object); err != nil {
		return nil, fmt.Errorf("decode request body: %w", err)
	}
	encoded, _ := json.Marshal(upstreamModel)
	object["model"] = encoded
	result, err := json.Marshal(object)
	if err != nil {
		return nil, fmt.Errorf("encode upstream request: %w", err)
	}
	return result, nil
}

func endpointURL(baseURL string, protocol domain.Protocol) string {
	baseURL = strings.TrimRight(baseURL, "/")
	if protocol == domain.ProtocolAnthropic {
		return baseURL + "/messages"
	}
	return baseURL + "/chat/completions"
}

func setHeaders(target, source http.Header, provider domain.Provider, protocol domain.Protocol, requestID string) {
	target.Set("Content-Type", "application/json")
	target.Set("Accept", source.Get("Accept"))
	if target.Get("Accept") == "" {
		target.Set("Accept", "application/json")
	}
	target.Set("User-Agent", "aigw/0.1")
	target.Set("X-Request-ID", requestID)
	if protocol == domain.ProtocolAnthropic {
		target.Set("x-api-key", provider.APIKey)
		version := source.Get("anthropic-version")
		if version == "" {
			version = "2023-06-01"
		}
		target.Set("anthropic-version", version)
		if beta := source.Get("anthropic-beta"); beta != "" {
			target.Set("anthropic-beta", beta)
		}
		return
	}
	target.Set("Authorization", "Bearer "+provider.APIKey)
}

func retryableStatus(status int) bool {
	switch status {
	case http.StatusTooManyRequests, http.StatusBadGateway, http.StatusServiceUnavailable, http.StatusGatewayTimeout:
		return true
	default:
		return false
	}
}