package tunnel

import (
	"context"
	"crypto/ed25519"
	"encoding/base64"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"time"

	"github.com/coder/websocket"
)

// Client keeps one tunnel open: a control connection to the gateway, and one
// dialled-back data connection per backend connection the public side needs,
// each bridged to the local instance with two plain io.Copy loops.
type Client struct {
	Gateway   string // gateway origin, e.g. "https://tunnel.11ch.at" (or http:// in tests)
	LocalAddr string // where the local instance listens, e.g. "127.0.0.1:52801"
	Version   string // informative, sent in the hello
	// Relay mode: with Identity set the client says hello with Mode "relay"
	// and answers the gateway's challenge, claiming (or reclaiming) the
	// stable hostname that belongs to this keypair. Nil Identity = a trial.
	Identity ed25519.PrivateKey
}

// Run registers, reports the assigned public host on ready, then serves dials
// until the control connection dies or ctx is cancelled. The error return is
// why the tunnel ended (nil on clean ctx cancellation).
func (cl *Client) Run(ctx context.Context, ready func(host string)) error {
	dialCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
	c, _, err := websocket.Dial(dialCtx, cl.wsURL("/api/tunnel/register"), nil)
	cancel()
	if err != nil {
		return fmt.Errorf("reach the tunnel gateway: %w", err)
	}
	defer c.CloseNow()
	c.SetReadLimit(1 << 16)

	hello := ControlMsg{Type: "hello", Version: cl.Version}
	if cl.Identity != nil {
		hello.Mode = "relay"
	}
	if err := writeJSON(ctx, c, hello); err != nil {
		return err
	}

	// Keepalive: a ping every 30s keeps NATs and the router's idle timers
	// friendly, and notices a dead gateway.
	go func() {
		t := time.NewTicker(30 * time.Second)
		defer t.Stop()
		for {
			select {
			case <-t.C:
				pctx, pcancel := context.WithTimeout(ctx, 10*time.Second)
				_ = c.Ping(pctx)
				pcancel()
			case <-ctx.Done():
				return
			}
		}
	}()

	// One loop for the whole conversation. "dial" can arrive BEFORE "ready" —
	// the gateway probes the backend through the tunnel before it publishes
	// the hostname — so dials are served from the first message on.
	gotReady := false
	for {
		var m ControlMsg
		if err := readJSON(ctx, c, &m); err != nil {
			if ctx.Err() != nil {
				return nil
			}
			return fmt.Errorf("tunnel closed: %w", err)
		}
		switch m.Type {
		case "challenge":
			if cl.Identity == nil {
				return fmt.Errorf("gateway sent a challenge but no relay identity is set")
			}
			pub := base64.RawURLEncoding.EncodeToString(cl.Identity.Public().(ed25519.PublicKey))
			sig := base64.RawURLEncoding.EncodeToString(ed25519.Sign(cl.Identity, []byte(m.Nonce)))
			if err := writeJSON(ctx, c, ControlMsg{Type: "claim", PubKey: pub, Sig: sig}); err != nil {
				return err
			}
		case "dial":
			go cl.serveDial(ctx, m.ID)
		case "ready":
			if !gotReady {
				gotReady = true
				ready(m.Host)
			}
		case "error":
			return fmt.Errorf("%s", m.Error)
		}
	}
}

// serveDial opens the data connection for one id and splices it to the local
// instance.
func (cl *Client) serveDial(ctx context.Context, id string) {
	dialCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
	c, _, err := websocket.Dial(dialCtx, cl.wsURL("/api/tunnel/conn/"+id), nil)
	cancel()
	if err != nil {
		return
	}
	c.SetReadLimit(-1)
	remote := websocket.NetConn(context.Background(), c, websocket.MessageBinary)
	local, err := net.DialTimeout("tcp", cl.LocalAddr, 10*time.Second)
	if err != nil {
		remote.Close()
		return
	}
	go splice(local, remote)
	splice(remote, local)
}

func splice(dst io.WriteCloser, src io.Reader) {
	_, _ = io.Copy(dst, src)
	dst.Close()
}

func (cl *Client) wsURL(path string) string {
	origin := strings.TrimRight(cl.Gateway, "/")
	origin = strings.Replace(origin, "https://", "wss://", 1)
	origin = strings.Replace(origin, "http://", "ws://", 1)
	return origin + path
}

// WaitReachable polls url until it answers or the timeout passes — used by
// `lchat try` so the printed link is live before anyone can click it (the
// router notices new routes within a couple of seconds).
func WaitReachable(url string, timeout time.Duration) error {
	deadline := time.Now().Add(timeout)
	client := &http.Client{Timeout: 5 * time.Second}
	var last error
	for time.Now().Before(deadline) {
		resp, err := client.Get(url)
		if err == nil {
			resp.Body.Close()
			if resp.StatusCode < 500 {
				return nil
			}
			last = fmt.Errorf("HTTP %d", resp.StatusCode)
		} else {
			last = err
		}
		time.Sleep(700 * time.Millisecond)
	}
	return fmt.Errorf("not reachable in %s: %v", timeout, last)
}