package main

import (
	"context"
	"errors"
	"io"
	"net/http"
	"strings"
	"time"
)

var hopByHop = map[string]bool{
	"connection": true, "keep-alive": true, "proxy-authenticate": true, "proxy-authorization": true,
	"proxy-connection": true, "te": true, "trailer": true, "trailers": true,
	"transfer-encoding": true, "upgrade": true, "content-length": true, "host": true,
}

var identityHeaders = map[string]bool{
	"cookie": true, "cookie2": true, "authorization": true, "set-cookie": true, "set-cookie2": true,
}

type FetchResult struct {
	Status    int
	Headers   map[string]string
	Body      []byte
	Truncated bool
	Elapsed   time.Duration
}

type Fetcher struct {
	policy *Policy
	client *http.Client
}

func NewFetcher(policy *Policy) *Fetcher {
	transport := &http.Transport{
		Proxy:                  nil,
		DialContext:            policy.DialContext,
		ForceAttemptHTTP2:      true,
		TLSHandshakeTimeout:    10 * time.Second,
		ResponseHeaderTimeout:  20 * time.Second,
		MaxResponseHeaderBytes: 64 * 1024,
		MaxIdleConns:           4,
		MaxIdleConnsPerHost:    2,
		IdleConnTimeout:        30 * time.Second,
	}
	client := &http.Client{
		Transport: transport,
		Jar:       nil,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			if len(via) > MaxRedirects {
				return refuse(RefuseRedirect, "too many redirects")
			}
			_, err := policy.CheckRedirect(req.URL.String(), req.Method)
			return err
		},
	}
	return &Fetcher{policy: policy, client: client}
}

func requestHeaders(source map[string]any) http.Header {
	header := http.Header{}
	count := 0
	for name, raw := range source {
		value, ok := raw.(string)
		lower := strings.ToLower(strings.TrimSpace(name))
		if !ok || lower == "" || hopByHop[lower] || identityHeaders[lower] || len(value) > MaxHeaderValueBytes {
			continue
		}
		if strings.ContainsAny(value, "\r\n") || strings.ContainsAny(lower, "\r\n: ") {
			continue
		}
		count++
		if count > MaxRequestHeaders {
			break
		}
		header.Set(lower, value)
	}
	return header
}

func responseHeaders(source http.Header) map[string]string {
	out := map[string]string{}
	for name, values := range source {
		lower := strings.ToLower(name)
		if hopByHop[lower] || identityHeaders[lower] || len(out) >= MaxResponseHeaders {
			continue
		}
		joined := strings.Join(values, ", ")
		if len(joined) > MaxHeaderValueBytes {
			joined = joined[:MaxHeaderValueBytes]
		}
		out[lower] = joined
	}
	return out
}

func ErrorCode(err error) string {
	var policyErr *PolicyError
	if errors.As(err, &policyErr) {
		return policyErr.Code
	}
	if errors.Is(err, context.DeadlineExceeded) {
		return RefuseTimeout
	}
	return FailFetch
}

func clampTimeout(seconds int) time.Duration {
	if seconds < 1 {
		seconds = DefaultJobTimeout
	}
	if seconds > MaxJobTimeout {
		seconds = MaxJobTimeout
	}
	return time.Duration(seconds) * time.Second
}

func (f *Fetcher) Do(ctx context.Context, frame Frame, maxBody int, budget int64) (*FetchResult, error) {
	target, err := f.policy.Check(frame.URL, frame.Method, "fetch")
	if err != nil {
		return nil, err
	}
	ctx, cancel := context.WithTimeout(ctx, clampTimeout(frame.Timeout))
	defer cancel()
	started := time.Now()
	request, err := http.NewRequestWithContext(ctx, target.Method, target.URL, nil)
	if err != nil {
		return nil, refuse(RefuseBadURL, "request cannot be built")
	}
	request.Header = requestHeaders(frame.Headers)
	response, err := f.client.Do(request)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	allowed := int64(maxBody)
	if budget < allowed {
		allowed = budget
	}
	body, err := io.ReadAll(io.LimitReader(response.Body, allowed+1))
	if err != nil {
		return nil, err
	}
	truncated := int64(len(body)) > allowed
	if truncated {
		body = body[:allowed]
	}
	return &FetchResult{
		Status:    response.StatusCode,
		Headers:   responseHeaders(response.Header),
		Body:      body,
		Truncated: truncated,
		Elapsed:   time.Since(started),
	}, nil
}