package replica

import (
	"context"
	"fmt"
	"net"
	"net/url"
	"time"

	"github.com/go-git/go-git/v5/plumbing/transport"
	"golang.org/x/net/proxy"
)

const deadlineScheme = "tug-replica-deadline"

func init() {
	proxy.RegisterDialerType(deadlineScheme, newDeadlineDialer)
}

type deadlineDialer struct {
	total time.Duration
	fetch *fetchLimit
}

func newDeadlineDialer(u *url.URL, _ proxy.Dialer) (proxy.Dialer, error) {
	if id := u.Query().Get("fetch"); id != "" {
		limit, ok := activeFetches.Load(id)
		if !ok {
			return nil, fmt.Errorf("replica: fetch has ended")
		}
		return deadlineDialer{total: limit.(*fetchLimit).total, fetch: limit.(*fetchLimit)}, nil
	}
	d, err := time.ParseDuration(u.Query().Get("d"))
	if err != nil {
		return nil, err
	}
	return deadlineDialer{total: d}, nil
}

func (d deadlineDialer) Dial(network, addr string) (net.Conn, error) {
	return d.DialContext(context.Background(), network, addr)
}

func (d deadlineDialer) DialContext(ctx context.Context, network, addr string) (net.Conn, error) {
	if d.fetch != nil {
		ctx = d.fetch.ctx
	}
	var nd net.Dialer
	c, err := nd.DialContext(ctx, network, addr)
	if err != nil {
		return nil, err
	}
	if d.total > 0 {
		if err := c.SetDeadline(time.Now().Add(d.total)); err != nil {
			c.Close()
			return nil, err
		}
	}
	if d.fetch != nil {
		return &fetchConn{Conn: c, limit: d.fetch, stop: context.AfterFunc(ctx, func() { c.Close() })}, nil
	}
	return c, nil
}

func proxyFor(auth transport.AuthMethod) transport.ProxyOptions {
	a, ok := auth.(*SSH)
	if ok && a.fetch != "" {
		return transport.ProxyOptions{URL: deadlineScheme + "://deadline/?fetch=" + a.fetch}
	}
	if !ok || a.Timeout <= 0 {
		return transport.ProxyOptions{}
	}
	return transport.ProxyOptions{URL: deadlineScheme + "://deadline/?d=" + a.Timeout.String()}
}