package websearch

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/netip"
	"net/url"
	"strings"
	"time"
)

// Fetch limits. A page is read to be quoted into a context window, so the
// ceiling is set by what is worth sending to the model rather than by what the
// machine could hold.
const (
	// DefaultMaxBytes is how much of a page is read when nothing says otherwise.
	DefaultMaxBytes = 512 << 10
	// MaxBytes is the ceiling on asking for more.
	MaxBytes = 2 << 20
	// FetchTimeout is how long one page may take, start to finish.
	FetchTimeout = 15 * time.Second
)

// Page is a fetched document, reduced to text.
type Page struct {
	// URL is where the content finally came from, after redirects.
	URL         string
	ContentType string
	Text        string
	Bytes       int
	// Truncated reports that the read stopped at the byte ceiling, so the text
	// ends in the middle of the document rather than at the end of it.
	Truncated bool
}

// CheckURL parses and vets a URL before anything dials it.
//
// This is the path jail's twin, and it exists for the same reason: the argument
// is written by the model, not by the user. A fetch tool that will dial
// whatever it is handed is an SSRF hole pointed at everything the machine can
// reach but the internet cannot — 169.254.169.254 holds cloud credentials, and
// localhost holds whatever the user is developing. The check is on resolved
// addresses rather than on the hostname, because a name can point anywhere and
// an attacker who controls a DNS record controls what a name means.
func CheckURL(ctx context.Context, raw string) (*url.URL, error) {
	u, err := url.Parse(strings.TrimSpace(raw))
	if err != nil {
		return nil, fmt.Errorf("not a URL: %v", err)
	}
	switch u.Scheme {
	case "http", "https":
	case "":
		return nil, fmt.Errorf("%s has no scheme — give an absolute http or https URL", raw)
	default:
		return nil, fmt.Errorf("%s: only http and https can be fetched", u.Scheme)
	}
	host := u.Hostname()
	if host == "" {
		return nil, fmt.Errorf("%s has no host", raw)
	}

	addrs, err := net.DefaultResolver.LookupNetIP(ctx, "ip", host)
	if err != nil {
		return nil, fmt.Errorf("cannot resolve %s: %v", host, err)
	}
	for _, a := range addrs {
		if err := checkAddr(a); err != nil {
			return nil, fmt.Errorf("%s resolves to %s: %w", host, a, err)
		}
	}
	return u, nil
}

// checkAddr is the rule itself: an address must be one the public internet
// could also route to. Everything else belongs to this machine or this network
// and is none of the model's business.
func checkAddr(a netip.Addr) error {
	a = a.Unmap()
	switch {
	case a.IsLoopback():
		return fmt.Errorf("a loopback address, which is this machine")
	case a.IsPrivate():
		return fmt.Errorf("a private address, which is this network")
	case a.IsLinkLocalUnicast(), a.IsLinkLocalMulticast():
		return fmt.Errorf("a link-local address (169.254.169.254 and friends hold cloud credentials)")
	case a.IsUnspecified():
		return fmt.Errorf("an unspecified address")
	case a.IsMulticast(), a.IsInterfaceLocalMulticast():
		return fmt.Errorf("a multicast address")
	case isUniqueLocal(a):
		return fmt.Errorf("a unique-local address, which is this network")
	case !a.IsGlobalUnicast():
		return fmt.Errorf("not a routable address")
	}
	return nil
}

// isUniqueLocal covers fc00::/7, IPv6's answer to 10.0.0.0/8. netip's IsPrivate
// already says yes to it, but the two are spelled out separately so that a
// change in either does not silently widen the other.
func isUniqueLocal(a netip.Addr) bool {
	return a.Is6() && a.As16()[0]&0xfe == 0xfc
}

// SafeClient is an HTTP client that cannot be talked out of the jail.
//
// The check in CheckURL happens before the request; these two happen during it.
// CheckRedirect re-runs the whole check on every hop, because a public URL that
// 302s to http://127.0.0.1:8080 would otherwise walk straight through the front
// door. DialContext re-checks the address actually being connected to, which
// closes the window between resolving a name and dialling it — the same name
// can answer differently twice.
func SafeClient(timeout time.Duration) *http.Client {
	dialer := &net.Dialer{Timeout: 10 * time.Second}
	transport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			host, _, err := net.SplitHostPort(addr)
			if err != nil {
				return nil, err
			}
			if ip, err := netip.ParseAddr(host); err == nil {
				if err := checkAddr(ip); err != nil {
					return nil, fmt.Errorf("refusing to connect to %s: %w", host, err)
				}
			}
			return dialer.DialContext(ctx, network, addr)
		},
		TLSHandshakeTimeout:   10 * time.Second,
		ResponseHeaderTimeout: timeout,
	}
	return &http.Client{
		Timeout:   timeout,
		Transport: transport,
		CheckRedirect: func(req *http.Request, via []*http.Request) error {
			if len(via) >= 5 {
				return fmt.Errorf("too many redirects")
			}
			if _, err := CheckURL(req.Context(), req.URL.String()); err != nil {
				return fmt.Errorf("refusing the redirect to %s: %w", req.URL, err)
			}
			return nil
		},
	}
}

// Fetch retrieves one URL and returns it as text.
//
// Everything about it is bounded: the address it may dial, the time it may
// take, the bytes it may read, and the types it will accept. A tool the model
// drives has to be, because the model is not the one who pays for an
// unbounded read of something that turned out to be a disk image.
func Fetch(ctx context.Context, raw string, maxBytes int) (*Page, error) {
	ctx, cancel := context.WithTimeout(ctx, FetchTimeout)
	defer cancel()

	u, err := CheckURL(ctx, raw)
	if err != nil {
		return nil, err
	}
	return fetchWith(ctx, SafeClient(FetchTimeout), u.String(), maxBytes)
}

// fetchWith is Fetch once the URL has been vetted and a client chosen. The
// split exists for the tests: an httptest server lives on loopback, which the
// jail refuses by design, so the reading and the refusing have to be provable
// separately.
func fetchWith(ctx context.Context, client *http.Client, raw string, maxBytes int) (*Page, error) {
	if maxBytes <= 0 {
		maxBytes = DefaultMaxBytes
	}
	if maxBytes > MaxBytes {
		maxBytes = MaxBytes
	}

	req, err := http.NewRequestWithContext(ctx, http.MethodGet, raw, nil)
	if err != nil {
		return nil, err
	}
	// A plausible browser string, because a bare Go user-agent is blocked by
	// enough documentation sites to make the tool look broken, and an honest
	// suffix so an operator reading their logs knows what this is.
	req.Header.Set("User-Agent", "Mozilla/5.0 (compatible; codey/1.0; +https://github.com/)")
	req.Header.Set("Accept", "text/html,application/xhtml+xml,text/plain,application/json;q=0.9,*/*;q=0.1")

	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("fetching %s: %v", raw, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return nil, fmt.Errorf("%s returned %s", raw, resp.Status)
	}

	ctype := resp.Header.Get("Content-Type")
	kind, ok := readableKind(ctype)
	if !ok {
		// Refused before the body is read: there is no point spending a
		// megabyte of someone's connection on a PDF this cannot render.
		return nil, fmt.Errorf("%s is %s, which cannot be read as text", raw, firstField(ctype))
	}

	// One byte past the ceiling, so hitting it is distinguishable from ending
	// exactly on it.
	body, err := io.ReadAll(io.LimitReader(resp.Body, int64(maxBytes)+1))
	if err != nil {
		return nil, fmt.Errorf("reading %s: %v", raw, err)
	}
	truncated := len(body) > maxBytes
	if truncated {
		body = body[:maxBytes]
	}

	text := string(body)
	if kind == kindHTML {
		text, err = HTMLToText(strings.NewReader(text))
		if err != nil {
			return nil, fmt.Errorf("reading %s: %v", raw, err)
		}
	}

	return &Page{
		URL:         resp.Request.URL.String(),
		ContentType: firstField(ctype),
		Text:        strings.TrimSpace(text),
		Bytes:       len(body),
		Truncated:   truncated,
	}, nil
}

type contentKind int

const (
	kindHTML contentKind = iota
	kindPlain
)

// readableKind decides whether a content type is worth reading and whether it
// needs the HTML reducer. An empty type is treated as HTML: servers that omit
// it are usually serving a page, and the reducer degrades to "strip nothing"
// on text that has no tags in it.
func readableKind(ctype string) (contentKind, bool) {
	t := strings.ToLower(firstField(ctype))
	switch {
	case t == "":
		return kindHTML, true
	case strings.Contains(t, "html"), strings.Contains(t, "xml"):
		return kindHTML, true
	case strings.HasPrefix(t, "text/"):
		return kindPlain, true
	case strings.Contains(t, "json"), strings.Contains(t, "javascript"):
		return kindPlain, true
	}
	return kindPlain, false
}

func firstField(ctype string) string {
	if i := strings.IndexByte(ctype, ';'); i >= 0 {
		ctype = ctype[:i]
	}
	return strings.TrimSpace(ctype)
}