package main

import (
	"bufio"
	"context"
	"flag"
	"fmt"
	"log"
	"net"
	"net/http"
	"net/netip"
	"strings"
	"sync"
	"time"

	"github.com/go-gost/plugin/bypass/proto"
	"google.golang.org/grpc"
)

type server struct {
	proto.UnimplementedBypassServer
	mu     sync.RWMutex
	list   []netip.Prefix
	lookup func(context.Context, string, string) ([]netip.Addr, error)
}

func (s *server) Bypass(ctx context.Context, in *proto.BypassRequest) (*proto.BypassReply, error) {
	// GOST 服务级插件返回 Ok=true 表示拒绝，外层 whitelist 参数不生效。
	reply := &proto.BypassReply{Ok: true}
	host := in.GetAddr()
	if h, _, err := net.SplitHostPort(host); err == nil {
		host = h
	}
	var ips []netip.Addr
	if ip, err := netip.ParseAddr(host); err == nil {
		ips = []netip.Addr{ip}
	} else {
		ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
		defer cancel()
		var err error
		ips, err = s.lookup(ctx, "ip", host)
		if err != nil || len(ips) == 0 {
			return reply, nil
		}
	}
	s.mu.RLock()
	defer s.mu.RUnlock()
	for _, ip := range ips {
		allowed := false
		for _, prefix := range s.list {
			if prefix.Contains(ip.Unmap()) {
				allowed = true
				break
			}
		}
		// 任一解析结果不在白名单内就拒绝，避免混合地址被放行。
		if !allowed {
			return reply, nil
		}
	}
	reply.Ok = false
	return reply, nil
}

func (s *server) refresh(url string) error {
	client := http.Client{Timeout: 30 * time.Second}
	resp, err := client.Get(url)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("list returned HTTP %d", resp.StatusCode)
	}
	var list []netip.Prefix
	scanner := bufio.NewScanner(resp.Body)
	for scanner.Scan() {
		line, _, _ := strings.Cut(scanner.Text(), "#")
		line = strings.TrimSpace(line)
		if line == "" {
			continue
		}
		prefix, err := netip.ParsePrefix(line)
		if err != nil {
			return fmt.Errorf("invalid CIDR %q: %w", line, err)
		}
		list = append(list, prefix.Masked())
	}
	if err := scanner.Err(); err != nil {
		return err
	}
	if len(list) == 0 {
		return fmt.Errorf("empty CIDR list")
	}
	// 完整加载成功才替换，下载或解析失败继续使用旧列表。
	s.mu.Lock()
	s.list = list
	s.mu.Unlock()
	log.Printf("loaded %d CIDR rules", len(list))
	return nil
}

func main() {
	addr := flag.String("listen-addr", "127.0.0.1:18080", "gRPC listen address")
	url := flag.String("list-url", "https://china-operator-ip.yfgao.com/china.txt", "CIDR list URL")
	period := flag.Duration("refresh", 24*time.Hour, "CIDR refresh interval")
	flag.Parse()
	if *period <= 0 {
		log.Fatal("refresh must be positive")
	}
	s := &server{lookup: net.DefaultResolver.LookupNetIP}
	if err := s.refresh(*url); err != nil {
		log.Fatal(err)
	}
	go func() {
		for range time.Tick(*period) {
			if err := s.refresh(*url); err != nil {
				log.Printf("refresh failed, keeping previous list: %v", err)
			}
		}
	}()
	listener, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatal(err)
	}
	rpc := grpc.NewServer()
	proto.RegisterBypassServer(rpc, s)
	log.Printf("listening on %s", *addr)
	log.Fatal(rpc.Serve(listener))
}