package main

import (
	"bytes"
	"flag"
	"fmt"
	"net"
	"net/url"
	"os"
	"strconv"
	"strings"
	"sync/atomic"
	"time"

	"dragontcp/internal/protocol"
)

const maxHeader = 128 * 1024

var requestCounter atomic.Uint32

func readHTTPHeaders(conn net.Conn) ([]byte, []byte, error) {
	buf := make([]byte, 0, 8192)
	tmp := make([]byte, 8192)

	for {
		n, err := conn.Read(tmp)
		if n > 0 {
			buf = append(buf, tmp[:n]...)

			if len(buf) > maxHeader {
				return nil, nil, fmt.Errorf("HTTP headers too large")
			}

			if i := bytes.Index(buf, []byte("\r\n\r\n")); i >= 0 {
				end := i + 4
				return buf[:end], buf[end:], nil
			}
		}

		if err != nil {
			return nil, nil, err
		}
	}
}

func parseHostPort(authority string, defaultPort int) (string, int, error) {
	authority = strings.TrimSpace(authority)

	if host, portText, err := net.SplitHostPort(authority); err == nil {
		port, err := strconv.Atoi(portText)
		return host, port, err
	}

	// Host without port.
	if strings.HasPrefix(authority, "[") && strings.HasSuffix(authority, "]") {
		return strings.Trim(authority, "[]"), defaultPort, nil
	}

	if strings.Count(authority, ":") == 0 {
		return authority, defaultPort, nil
	}

	// Bare IPv6.
	if ip := net.ParseIP(authority); ip != nil {
		return authority, defaultPort, nil
	}

	return "", 0, fmt.Errorf("invalid authority: %s", authority)
}

func rewritePlainHTTPRequest(header []byte) (string, int, []byte, error) {
	text := string(header)
	lines := strings.Split(text, "\r\n")
	if len(lines) == 0 {
		return "", 0, nil, fmt.Errorf("empty request")
	}

	parts := strings.SplitN(lines[0], " ", 3)
	if len(parts) != 3 {
		return "", 0, nil, fmt.Errorf("invalid request line")
	}

	method, target, version := parts[0], parts[1], parts[2]

	var (
		hostHeader string
		headers    []string
	)

	for _, line := range lines[1:] {
		if line == "" {
			continue
		}

		k, v, ok := strings.Cut(line, ":")
		if !ok {
			continue
		}

		lk := strings.ToLower(strings.TrimSpace(k))

		if lk == "host" {
			hostHeader = strings.TrimSpace(v)
		}

		if lk == "connection" ||
			lk == "proxy-connection" ||
			lk == "proxy-authorization" {
			continue
		}

		headers = append(headers, k+": "+strings.TrimSpace(v))
	}

	u, err := url.Parse(target)
	if err != nil {
		return "", 0, nil, err
	}

	var host string
	var port int
	path := target

	if u.Hostname() != "" {
		if strings.ToLower(u.Scheme) != "http" {
			return "", 0, nil, fmt.Errorf("unsupported plain HTTP scheme: %s", u.Scheme)
		}

		host = u.Hostname()
		port = 80

		if u.Port() != "" {
			port, err = strconv.Atoi(u.Port())
			if err != nil {
				return "", 0, nil, err
			}
		}

		path = u.EscapedPath()
		if path == "" {
			path = "/"
		}
		if u.RawQuery != "" {
			path += "?" + u.RawQuery
		}
	} else {
		if hostHeader == "" {
			return "", 0, nil, fmt.Errorf("missing Host header")
		}

		host, port, err = parseHostPort(hostHeader, 80)
		if err != nil {
			return "", 0, nil, err
		}
		if path == "" {
			path = "/"
		}
	}

	var out strings.Builder
	fmt.Fprintf(&out, "%s %s %s\r\n", method, path, version)

	sawHost := false
	for _, h := range headers {
		if strings.HasPrefix(strings.ToLower(h), "host:") {
			sawHost = true
		}
		out.WriteString(h)
		out.WriteString("\r\n")
	}

	if !sawHost {
		if port == 80 {
			fmt.Fprintf(&out, "Host: %s\r\n", host)
		} else {
			fmt.Fprintf(&out, "Host: %s\r\n", net.JoinHostPort(host, strconv.Itoa(port)))
		}
	}

	out.WriteString("Connection: close\r\n\r\n")

	return host, port, []byte(out.String()), nil
}

func openDragonTCPTunnel(serverAddr, token, targetHost string, targetPort int, transport string, tcpBuffer int) (net.Conn, error) {
	d := net.Dialer{
		Timeout:   10 * time.Second,
		KeepAlive: 30 * time.Second,
	}

	conn, err := d.Dial("tcp", serverAddr)
	if err != nil {
		return nil, err
	}

	protocol.TuneTCP(conn)
	protocol.TuneTCPBuffer(conn, tcpBuffer)
	_ = conn.SetDeadline(time.Now().Add(15 * time.Second))

	// Correlation only; cryptographic randomness is unnecessary here.
	requestID := requestCounter.Add(1)

	var command []byte
	if transport == "raw" {
		command = []byte(fmt.Sprintf("TUNNEL2 %s %s %d RAW", token, targetHost, targetPort))
	} else {
		// Legacy XOR command remains compatible with the older server.
		command = []byte(fmt.Sprintf("TUNNEL %s %s %d", token, targetHost, targetPort))
	}

	if err := protocol.WriteRequestFrame(conn, requestID, command); err != nil {
		conn.Close()
		return nil, err
	}

	responseID, response, err := protocol.ReadResponseFrame(conn)
	if err != nil {
		conn.Close()
		return nil, err
	}

	if responseID != requestID {
		conn.Close()
		return nil, fmt.Errorf("request ID mismatch")
	}

	if string(response) != "CONNECTED" {
		conn.Close()
		return nil, fmt.Errorf("%s", response)
	}

	_ = conn.SetDeadline(time.Time{})
	return conn, nil
}

func writeHTTPError(conn net.Conn, code int, reason, detail string) {
	if detail == "" {
		detail = reason
	}

	body := []byte(detail)

	fmt.Fprintf(
		conn,
		"HTTP/1.1 %d %s\r\nContent-Type: text/plain; charset=utf-8\r\nContent-Length: %d\r\nConnection: close\r\n\r\n",
		code,
		reason,
		len(body),
	)
	_, _ = conn.Write(body)
}

func handleLocal(conn net.Conn, serverAddr, token, transport string, tcpBuffer int, chunkOpts chunkClientOptions, slots chan struct{}) {
	defer func() {
		<-slots
		_ = conn.Close()
	}()

	protocol.TuneTCP(conn)
	protocol.TuneTCPBuffer(conn, tcpBuffer)
	_ = conn.SetDeadline(time.Now().Add(15 * time.Second))

	header, extra, err := readHTTPHeaders(conn)
	if err != nil {
		return
	}

	firstLine := strings.SplitN(string(header), "\r\n", 2)[0]
	parts := strings.SplitN(firstLine, " ", 3)

	if len(parts) != 3 {
		writeHTTPError(conn, 400, "Bad Request", "invalid HTTP request line")
		return
	}

	method, target := parts[0], parts[1]

	if strings.EqualFold(method, "CONNECT") {
		host, port, err := parseHostPort(target, 443)
		if err != nil {
			writeHTTPError(conn, 400, "Bad Request", err.Error())
			return
		}

		var remote net.Conn
		if transport == "chunk" {
			remote, err = openChunkTunnel(serverAddr, token, host, port, chunkOpts)
		} else {
			remote, err = openDragonTCPTunnel(serverAddr, token, host, port, transport, tcpBuffer)
		}
		if err != nil {
			writeHTTPError(conn, 502, "Bad Gateway", err.Error())
			return
		}
		defer remote.Close()

		_, _ = conn.Write([]byte(
			"HTTP/1.1 200 Connection Established\r\n" +
				"Proxy-Agent: dragontcp-proxy/2.0\r\n\r\n",
		))

		if len(extra) > 0 {
			if transport == "xor" {
				protocol.XorInPlace(extra)
			}
			if _, err := remote.Write(extra); err != nil {
				return
			}
		}

		_ = conn.SetDeadline(time.Time{})
		if transport == "xor" {
			protocol.RelayXOR(conn, remote)
		} else {
			// raw and chunk connections expose a normal plaintext net.Conn.
			protocol.RelayRaw(conn, remote)
		}
		return
	}

	host, port, rewritten, err := rewritePlainHTTPRequest(header)
	if err != nil {
		writeHTTPError(conn, 400, "Bad Request", err.Error())
		return
	}

	var remote net.Conn
	if transport == "chunk" {
		remote, err = openChunkTunnel(serverAddr, token, host, port, chunkOpts)
	} else {
		remote, err = openDragonTCPTunnel(serverAddr, token, host, port, transport, tcpBuffer)
	}
	if err != nil {
		writeHTTPError(conn, 502, "Bad Gateway", err.Error())
		return
	}
	defer remote.Close()

	initial := make([]byte, 0, len(rewritten)+len(extra))
	initial = append(initial, rewritten...)
	initial = append(initial, extra...)
	if transport == "xor" {
		protocol.XorInPlace(initial)
	}

	if _, err := remote.Write(initial); err != nil {
		return
	}

	_ = conn.SetDeadline(time.Time{})
	if transport == "xor" {
		protocol.RelayXOR(conn, remote)
	} else {
		protocol.RelayRaw(conn, remote)
	}
}

func main() {
	var (
		listenHost       = flag.String("listen-host", "127.0.0.1", "local proxy listen host")
		listenPort       = flag.Int("listen-port", 8080, "local proxy listen port")
		serverHost       = flag.String("server-host", "", "remote DragonTCP server host")
		serverPort       = flag.Int("server-port", 53, "remote DragonTCP server port")
		token            = flag.String("token", "", "optional shared token")
		maxConnections   = flag.Int("max-connections", 20000, "max simultaneous proxy connections")
		transport        = flag.String("transport", "chunk", "transport: chunk (DragonTCP binary adaptive transport)")
		tcpBuffer        = flag.Int("tcp-buffer", 0, "optional TCP read/write buffer bytes; 0 keeps OS autotuning")
		chunkStart       = flag.Int("chunk-start", 1048576, "initial adaptive chunk payload bytes")
		chunkMin         = flag.Int("chunk-min", 32, "minimum adaptive chunk payload bytes")
		chunkMax         = flag.Int("chunk-max", 1048576, "maximum adaptive chunk payload bytes (up to 1 MiB)")
		chunkAdaptive    = flag.Bool("chunk-adaptive", true, "automatically shrink on failures and grow after stable success")
		chunkSuccesses   = flag.Int("chunk-grow-after", 16, "successful data records required before increasing chunk size")
		chunkAdaptLog    = flag.Bool("chunk-adapt-log", true, "print adaptive chunk size changes")
		chunkSizeLegacy  = flag.Int("chunk-size", 0, "legacy fixed chunk size; nonzero disables adaptation")
		chunkPollers     = flag.Int("chunk-pollers", 1, "reserved compatibility setting; binary transport uses one download worker")
		chunkConcurrency = flag.Int("chunk-concurrency", 1, "maximum adaptive download pipeline depth (1-256); 1 keeps concurrency fixed at one")
		chunkReconnect   = flag.Int("chunk-reconnect-every", 0, "force reconnect after N logical requests; 0 = persistent/automatic")
		chunkPollDelay   = flag.Duration("chunk-poll-delay", 2*time.Millisecond, "delay after an empty chunk poll")
		chunkTimeout     = flag.Duration("chunk-timeout", 2*time.Second, "per-record transaction timeout before adaptive shrink")
	)
	flag.Parse()

	if *serverHost == "" {
		fmt.Fprintln(os.Stderr, "--server-host is required")
		os.Exit(2)
	}

	*transport = strings.ToLower(*transport)
	if *transport != "chunk" {
		fmt.Fprintln(os.Stderr, "DragonTCP requires --transport chunk (binary adaptive TCP/53 transport)")
		os.Exit(2)
	}
	if *chunkSizeLegacy != 0 {
		if *chunkSizeLegacy < 32 || *chunkSizeLegacy > protocol.MaxChunkPayload {
			fmt.Fprintf(os.Stderr, "--chunk-size must be between 32 and %d\n", protocol.MaxChunkPayload)
			os.Exit(2)
		}
		*chunkStart = *chunkSizeLegacy
		*chunkMin = *chunkSizeLegacy
		*chunkMax = *chunkSizeLegacy
		*chunkAdaptive = false
	}
	if *chunkMin < 32 || *chunkMax > protocol.MaxChunkPayload || *chunkMin > *chunkStart || *chunkStart > *chunkMax {
		fmt.Fprintf(os.Stderr, "require 32 <= --chunk-min <= --chunk-start <= --chunk-max <= %d\n", protocol.MaxChunkPayload)
		os.Exit(2)
	}
	if *chunkSuccesses < 1 {
		fmt.Fprintln(os.Stderr, "--chunk-grow-after must be at least 1")
		os.Exit(2)
	}
	if *chunkPollers < 1 || *chunkPollers > 128 {
		fmt.Fprintln(os.Stderr, "--chunk-pollers must be between 1 and 128")
		os.Exit(2)
	}
	if *chunkConcurrency < 1 || *chunkConcurrency > 256 {
		fmt.Fprintln(os.Stderr, "--chunk-concurrency must be between 1 and 256")
		os.Exit(2)
	}
	if *chunkReconnect < 0 {
		fmt.Fprintln(os.Stderr, "--chunk-reconnect-every must be 0 or greater")
		os.Exit(2)
	}
	chunkOpts := chunkClientOptions{
		startSize:      *chunkStart,
		minSize:        *chunkMin,
		maxSize:        *chunkMax,
		adaptive:       *chunkAdaptive,
		adaptSuccesses: *chunkSuccesses,
		adaptLog:       *chunkAdaptLog,
		pollers:        *chunkPollers,
		reconnectEvery: *chunkReconnect,
		pollDelay:      *chunkPollDelay,
		txnTimeout:     *chunkTimeout,
		tcpBuffer:      *tcpBuffer,
		maxPipeline:    *chunkConcurrency,
	}

	listenAddr := net.JoinHostPort(*listenHost, strconv.Itoa(*listenPort))
	serverAddr := net.JoinHostPort(*serverHost, strconv.Itoa(*serverPort))

	ln, err := net.Listen("tcp", listenAddr)
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	defer ln.Close()

	fmt.Printf("local Go HTTP proxy listening on %s\n", listenAddr)
	fmt.Printf("remote DragonTCP endpoint=%s\n", serverAddr)
	fmt.Printf("max_connections=%d transport=%s tcp_buffer=%d\n", *maxConnections, *transport, *tcpBuffer)
	if *transport == "chunk" {
		fmt.Printf(
			"adaptive_chunk=%v start=%d min=%d max=%d grow_after=%d pollers=%d concurrency=%d reconnect_every=%d timeout=%s\n",
			*chunkAdaptive,
			*chunkStart,
			*chunkMin,
			*chunkMax,
			*chunkSuccesses,
			*chunkPollers,
			*chunkConcurrency,
			*chunkReconnect,
			chunkTimeout.String(),
		)
	}

	slots := make(chan struct{}, *maxConnections)

	for {
		conn, err := ln.Accept()
		if err != nil {
			fmt.Fprintln(os.Stderr, "accept:", err)
			continue
		}

		select {
		case slots <- struct{}{}:
			go handleLocal(conn, serverAddr, *token, *transport, *tcpBuffer, chunkOpts, slots)
		default:
			writeHTTPError(
				conn,
				503,
				"Service Unavailable",
				"proxy connection limit reached",
			)
			_ = conn.Close()
		}
	}
}