package hook

import (
	"bufio"
	"bytes"
	"encoding/base64"
	"errors"
	"fmt"
	"io"
	"regexp"
	"strconv"
	"time"
	"tunnel/pkg/server/env"
	"tunnel/pkg/server/opts"
	"tunnel/pkg/server/queue"
)

var addrRe = regexp.MustCompile("^[0-9a-zA-Z-.]+:[0-9]+$")
var respRe = regexp.MustCompile("^([^ ]+) +([0-9]+) +(.*)$")

var errBadHttpResponse = errors.New("bad HTTP response")

type proxyHook struct {
	addr string
	auth string
}

type proxy struct {
	addr string
	auth string
	ok   chan struct{}
	fail chan struct{}
}

func (p *proxy) Send(rq, wq queue.Q) error {
	var out bytes.Buffer

	fmt.Fprintf(&out, "CONNECT %s HTTP/1.0\r\n", p.addr)
	if p.auth != "" {
		encoded := base64.StdEncoding.EncodeToString([]byte(p.auth))
		fmt.Fprintf(&out, "Proxy-Authorization: Basic %s\r\n", encoded)
	}
	fmt.Fprintf(&out, "\r\n")

	wq <- out.Bytes()

	select {
	case <-p.fail:
		return nil
	case <-p.ok:
	}

	return queue.Copy(rq, wq)
}

func parseProxyResponse(s string) error {
	var version string
	var code int
	var desc string

	if m := respRe.FindStringSubmatch(s); m == nil {
		return errBadHttpResponse
	} else {
		version = m[1]
		if c, err := strconv.Atoi(m[2]); err != nil {
			return errBadHttpResponse
		} else {
			code = c
		}
		desc = m[3]
	}

	if version != "HTTP/1.0" && version != "HTTP/1.1" {
		return errBadHttpResponse
	}

	if code != 200 {
		return fmt.Errorf("connect failed: %d %s", code, desc)
	}

	return nil
}

func (p *proxy) Recv(rq, wq queue.Q) (err error) {
	defer func() {
		if err != nil {
			close(p.fail)
		}
	}()

	s := bufio.NewScanner(rq.TimeoutReader(5 * time.Second))

	var resp bool

	for s.Scan() {
		line := s.Text()

		if !resp {
			if err := parseProxyResponse(line); err != nil {
				return err
			}
			resp = true
			continue
		}

		if line == "" {
			break
		}
	}

	if err := s.Err(); err != nil {
		return err
	} else if !resp {
		return io.ErrUnexpectedEOF
	}

	close(p.ok)

	return queue.Copy(rq, wq)
}

func (h *proxyHook) Open(env env.Env) (interface{}, error) {
	p := &proxy{
		addr: h.addr,
		auth: h.auth,
		ok:   make(chan struct{}),
		fail: make(chan struct{}),
	}

	if p.auth == "" {
		p.auth = getHookVar(env, "proxy.auth")
	}

	return p, nil
}

func newProxyHook(opts opts.Opts, env env.Env) (hook, error) {
	h := &proxyHook{}

	if addr, ok := opts["addr"]; !ok {
		return nil, errors.New("proxy: missing addr")
	} else if !addrRe.MatchString(addr) {
		return nil, errors.New("proxy: invalid addr")
	} else {
		h.addr = addr
	}

	h.auth = opts["auth"]

	return h, nil
}

func init() {
	register("proxy", newProxyHook)
}