package vision

import (
	"context"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/netip"
	"net/url"
)

// 图像 URL 守卫：服务端会代取用户/渠道传入的任意 http(s) URL（F3 之后
// 含 IM 陌生人的来图），因此取回链路必须自己拦下链路本地与云元数据地址，
// 不能指望入口层过滤。这里只收紧「链路本地 + 云元数据」：环回、私网与
// 公网保持原样可达 —— 同源自取需要它们（本服务要能取自己的
// http://<局域网地址>:10942/api/v1/files/…）。
//
// 判定分三层，任一层命中都拒绝：
//  1. 取回前对字面主机校验（guardImageURL）：IP 字面量立即被拒，
//     DNS 名放行到下一层；
//  2. 重定向每一跳重新校验（guardedCheckRedirect）：跨主机跳转不得落到
//     被拒网段；
//  3. 真正建连前校验解析结果（guardedDialContext）：解析到被拒地址的连接
//     根本不会建立，也顺带让「先答公网 IP、再答元数据 IP」的 DNS 重绑失效。
//
// 三层只报「主机 + 网段原因」，绝不把 URL、userinfo 或查询串写进错误。

// maxImageRedirects 是一次取回最多跟随的重定向跳数。与 net/http 默认的 10
// 同形，只是更小：取一张图不需要长跳转链，少跟几跳就少几次对陌生主机的连接。
const maxImageRedirects = 5

// 被拒的云元数据地址以字符串常量给出（netip 的值只能是变量，而本项目约束
// 不允许包级变量，只允许 var ErrXxx），判定时解析一次即可。
const (
	// alibabaIMDSv4 是阿里云 ECS 元数据地址，落在 100.64.0.0/10 共享地址段内。
	alibabaIMDSv4 = "100.100.100.200"
	// awsIMDSv6Prefix 是 AWS 元数据专用 IPv6 前缀，fd00:ec2::254 在其内。
	// IPv6 ULA（fd00::/8）整体仍按私网放行，只拒这一个前缀。
	awsIMDSv6Prefix = "fd00:ec2::/64"
)

// blockedAddressError 是被拒地址的错误：只携带主机与网段原因，
// 绝不携带 URL 或其中的凭据，因此即使被 net/http 的 *url.Error 包裹
// （重定向与拨号失败都走这条路径），也能用 errors.As 取回原因原样呈现。
type blockedAddressError struct {
	host   string
	reason string
}

func (e *blockedAddressError) Error() string {
	return fmt.Sprintf("%s: %s（主机 %s）", ErrBlockedAddress, e.reason, e.host)
}

func (e *blockedAddressError) Unwrap() error { return ErrBlockedAddress }

// blockedImageAddrReason 报告 addr 是否属于被拒的链路本地/云元数据地址，
// 是则给出原因（用于错误文案）。IPv4-mapped 形式
// （::ffff:169.254.169.254）先经 Unmap 归一到 IPv4 再判定，绕不过去；
// IPv6 ULA（fd00::/8）整体仍按私网放行，只拒云元数据专用的 fd00:ec2::/64。
func blockedImageAddrReason(addr netip.Addr) (string, bool) {
	if !addr.IsValid() {
		return "", false
	}
	addr = addr.Unmap()
	// netip 的 IsLinkLocalUnicast 同时覆盖 IPv4 169.254.0.0/16
	// （含云元数据 169.254.169.254）与 IPv6 fe80::/10。
	if addr.IsLinkLocalUnicast() {
		return "链路本地地址 169.254.0.0/16（含云元数据 169.254.169.254）或 fe80::/10", true
	}
	if addr == netip.MustParseAddr(alibabaIMDSv4) {
		return "云元数据地址 " + alibabaIMDSv4 + "（阿里云 ECS）", true
	}
	// AWS IMDSv6：fd00:ec2::254 落在 AWS 元数据专用前缀内。
	if netip.MustParsePrefix(awsIMDSv6Prefix).Contains(addr) {
		return "云元数据地址 " + awsIMDSv6Prefix + "（AWS IMDSv6）", true
	}
	return "", false
}

// guardImageURL 校验一个待取回/待跳转的 URL：协议必须仍是 http/https，
// 且主机不得是被拒地址。主机是 DNS 名时这里放行（字面量判不出解析结果），
// 由 guardedDialContext 在建连前按解析结果继续拦。
func guardImageURL(u *url.URL) error {
	if u == nil {
		return fmt.Errorf("vision: %w: 空的 URL", ErrUnsupportedSource)
	}
	if u.Scheme != "http" && u.Scheme != "https" {
		return fmt.Errorf("vision: %w: 仅支持 http/https", ErrUnsupportedSource)
	}
	host := u.Hostname()
	if host == "" {
		return fmt.Errorf("vision: %w: URL 缺少主机", ErrFetch)
	}
	addr, err := netip.ParseAddr(host)
	if err != nil {
		return nil
	}
	reason, blocked := blockedImageAddrReason(addr)
	if !blocked {
		return nil
	}
	return &blockedAddressError{host: host, reason: reason}
}

// guardedImageClient 返回带守卫的 http.Client **副本**，绝不改动传入的
// client（调用方可能把它共享给上游 chat 请求等其他用途）：重定向逐跳校验
// 并限跳，未设超时时取下限超时（复用 imageFetchTimeout），传输层换成
// 「建连前按解析结果校验」的副本。
func guardedImageClient(base *http.Client) *http.Client {
	c := &http.Client{}
	if base != nil {
		*c = *base
	}
	if c.Timeout <= 0 {
		c.Timeout = imageFetchTimeout
	}
	c.Transport = guardedImageTransport(base)
	c.CheckRedirect = guardedCheckRedirect
	return c
}

// guardedImageTransport 返回带拨号守卫的传输层。用 *http.Transport 的副本
// 是为了把校验下沉到真正建连的一刻（连接池是逐 Transport 的，因此副本也
// 意味着每次取图新建连接；换取的是一张图都取不到被拒网段）。调用方给了
// 非 *http.Transport 的自定义 RoundTripper 时保留它 —— 此时字面主机与
// 重定向两层守卫仍然生效。走环境代理时拨号目标是代理本身（不是图片主机），
// 那种形态下拦下被拒网段的是前两层：字面主机与重定向逐跳校验。
func guardedImageTransport(base *http.Client) http.RoundTripper {
	if base != nil && base.Transport != nil {
		bt, ok := base.Transport.(*http.Transport)
		if !ok {
			return base.Transport
		}
		return guardedTransport(bt)
	}
	dt, ok := http.DefaultTransport.(*http.Transport)
	if !ok {
		return http.DefaultTransport
	}
	return guardedTransport(dt)
}

// guardedTransport 复制 t 并把拨号改为守卫版本。
func guardedTransport(t *http.Transport) *http.Transport {
	clone := t.Clone()
	clone.DialContext = guardedDialContext(clone.DialContext)
	return clone
}

// guardedCheckRedirect 校验重定向的每一跳：跳数有上限，跳转目标必须仍是
// http/https 且不是被拒地址。跨主机跳转因此不能把取回引到 169.254.169.254
// 或 fd00:ec2::254 上；即便这一层被绕过，第三层的拨号守卫也不会真正建连。
func guardedCheckRedirect(req *http.Request, via []*http.Request) error {
	if len(via) >= maxImageRedirects {
		// 这里不包 ErrFetch：可由 fetchURL 统一包一层，
		// 免得同一句话被套上两层前缀。
		return fmt.Errorf("重定向超过 %d 跳", maxImageRedirects)
	}
	return guardImageURL(req.URL)
}

// dialFunc 与 net.Dialer.DialContext 同形，便于测试注入假拨号器。
type dialFunc func(ctx context.Context, network, addr string) (net.Conn, error)

// guardedDialContext 在真正建连**之前**校验目标地址的解析结果：只要解析
// 结果里有被拒地址，整次连接就被拒（宁可不取这张图，也不让进程去连元数据
// 服务）；随后用校验过的 IP 直接连接，拨号器不会再解析一次，「先答公网 IP、
// 再答元数据 IP」的 DNS 重绑因此无从生效。
func guardedDialContext(base dialFunc) dialFunc {
	if base == nil {
		dialer := &net.Dialer{Timeout: imageFetchTimeout}
		base = dialer.DialContext
	}
	return func(ctx context.Context, network, addr string) (net.Conn, error) {
		host, port, err := net.SplitHostPort(addr)
		if err != nil {
			return nil, fmt.Errorf("vision: %w: 目标地址无法解析", ErrFetch)
		}
		addrs, err := resolveImageHost(ctx, host)
		if err != nil {
			return nil, err
		}
		for _, a := range addrs {
			if reason, blocked := blockedImageAddrReason(a); blocked {
				return nil, &blockedAddressError{host: host, reason: reason}
			}
		}
		var lastErr error
		for _, a := range addrs {
			conn, err := base(ctx, network, net.JoinHostPort(a.Unmap().String(), port))
			if err == nil {
				return conn, nil
			}
			lastErr = err
			if ctx.Err() != nil {
				break
			}
		}
		if lastErr == nil {
			return nil, fmt.Errorf("vision: %w: 主机 %s 没有可用地址", ErrFetch, host)
		}
		return nil, lastErr
	}
}

// resolveImageHost 把主机名解析为 IP 列表；主机已是 IP 字面量时直接返回它，
// 不再查一次 DNS。
func resolveImageHost(ctx context.Context, host string) ([]netip.Addr, error) {
	if addr, err := netip.ParseAddr(host); err == nil {
		return []netip.Addr{addr}, nil
	}
	addrs, err := net.DefaultResolver.LookupNetIP(ctx, "ip", host)
	if err != nil {
		return nil, fmt.Errorf("vision: %w: 解析主机 %s 失败: %v", ErrFetch, host, err)
	}
	if len(addrs) == 0 {
		return nil, fmt.Errorf("vision: %w: 主机 %s 没有解析结果", ErrFetch, host)
	}
	return addrs, nil
}

// blockedAddressOf 从（可能被 *url.Error 包裹的）错误里取回被拒地址错误，
// 使调用方能直接呈现「主机 + 网段原因」而不是带回整条 URL 的错误串。
func blockedAddressOf(err error) (*blockedAddressError, bool) {
	var blocked *blockedAddressError
	if errors.As(err, &blocked) {
		return blocked, true
	}
	return nil, false
}