package main

import (
	"bytes"
	"fmt"
	"go/format"
	"os"
	"strings"
)

type Build struct {
	Suffix     string
	Tags       string
	Interfaces Interfaces
}

func (b *Build) MustBuild() {
	prefix := "wrap_generated"
	if b.Suffix != "" {
		prefix += "_" + b.Suffix
	}

	b.Implementation().MustWriteFile(prefix + ".go")
	b.Tests().MustWriteFile(prefix + "_test.go")
}

func (b *Build) writeHeader(g *Generator) {
	if b.Tags != "" {
		g.Printf(`// +build %s
`, b.Tags)
	}
	g.buf.WriteString(`// Code generated by "httpsnoop/codegen"; DO NOT EDIT.

package httpsnoop

`)
}

func (b *Build) Implementation() *Generator {
	ifaces := b.Interfaces
	// subIfaces has all interfaces except http.ResponseWriter
	subIfaces := ifaces[1:]

	var g Generator
	// Package header
	b.writeHeader(&g)
	g.Printf("import (\n")
	g.Printf(`"net/http"` + "\n")
	g.Printf(`"io"` + "\n")
	g.Printf(`"net"` + "\n")
	g.Printf(`"bufio"` + "\n")
	g.Printf(`"time"` + "\n")
	g.Printf(")\n")
	g.Printf("\n")

	// Hook funcs
	for _, iface := range ifaces {
		for _, fn := range iface.Funcs {
			g.Printf("// %s is part of the %s interface.\n", fn.Type(), iface.Name)
			g.Printf("type %s func(%s) (%s)\n", fn.Type(), fn.Args, fn.Returns)
			g.Printf("\n")
		}
	}

	// Hooks struct
	g.Printf(`
// Hooks defines a set of method interceptors for methods included in
// http.ResponseWriter as well as some others. You can think of them as
// middleware for the function calls they target. See Wrap for more details.
//
// For each method, the exact matching hook takes precedence. For example,
// WriteString calls the WriteString hook when it is configured, even if a
// Write hook is also configured. If the exact hook is not configured, most
// methods call through to the underlying ResponseWriter directly.
//
// Two compatibility fallbacks preserve the behavior users had before Wrap
// learned about newer optional interfaces:
//   - If the underlying ResponseWriter implements io.StringWriter and
//     WriteString is called, but only the Write hook is configured, WriteString
//     is routed through the Write hook with []byte(s). If neither hook is
//     configured, WriteString calls the underlying WriteString method directly.
//   - If the underlying ResponseWriter implements both http.Flusher and
//     FlushError, and FlushError is called, but only the Flush hook is
//     configured, FlushError is routed through the Flush hook while preserving
//     the error returned by the underlying FlushError method. If neither hook is
//     configured, FlushError calls the underlying FlushError method directly.
type Hooks struct {
`)
	for _, iface := range ifaces {
		for _, fn := range iface.Funcs {
			g.Printf("%s func(%s) %s\n", fn.Name, fn.Type(), fn.Type())
		}
	}
	g.Printf("}\n")

	// Wrap func
	docList := make([]string, len(subIfaces))
	for i, iface := range subIfaces {
		docList[i] = "// - " + iface.Name
	}
	g.Printf(`
// Wrap returns a wrapped version of w that provides the exact same interface
// as w. Specifically if w implements any combination of:
//
%s
//
// The wrapped version will implement the exact same combination. If no hooks
// are set, the wrapped version also behaves exactly as w. Hooks targeting
// methods not supported by w are ignored. Any other hooks will intercept the
// method they target and may modify the call's arguments and/or return values.
// The CaptureMetrics implementation serves as a working example for how the
// hooks can be used.
`, strings.Join(docList, "\n"))
	g.Printf("func Wrap(w http.ResponseWriter, hooks Hooks) http.ResponseWriter {\n")
	g.Printf("state := &rwState{w: w}\n")

	// Precompute hook chains once per Wrap call and
	// build a uint8 combo index so the switch compiles to a jump table.
	g.Printf("var combo uint16\n")
	for _, fn := range ifaces[0].Funcs {
		g.Printf("if hooks.%s != nil {\n", fn.Name)
		g.Printf("state.%s = hooks.%s(w.%s)\n", fieldName(fn.Name), fn.Name, fn.Name)
		g.Printf("}\n")
	}

	for i, iface := range subIfaces {
		g.Printf("if t%[1]d, i%[1]d := w.(%s); i%[1]d {\n", i, iface.Name)
		bit := len(subIfaces) - i - 1
		g.Printf("combo |= 1<<%d\n", bit)
		for _, fn := range iface.Funcs {
			g.Printf("if hooks.%s != nil {\n", fn.Name)
			g.Printf("state.%s = hooks.%s(t%d.%s)\n", fieldName(fn.Name), fn.Name, i, fn.Name)
			if fn.Name == "FlushError" {
				// http.ResponseController.Flush prefers FlushError over Flush.
				// Preserve existing Flush hooks when wrapping writers that expose both.
				g.Printf("} else if state.flush != nil {\n")
				g.Printf("state.flushError = func() (err error) { hooks.Flush(func() { err = t%d.FlushError() })(); return err }\n", i)
			} else if fn.Name == "WriteString" {
				// io.WriteString prefers WriteString over Write. Preserve existing Write
				// hooks when wrapping writers that expose both methods.
				g.Printf("} else if state.write != nil {\n")
				g.Printf("state.writeString = func(s string) (int, error) { return state.write([]byte(s)) }\n")
			}
			g.Printf("}\n")
		}
		g.Printf("}\n")
	}

	g.Printf("switch combo {\n")
	combinations := 1 << uint(len(subIfaces))
	for c := 0; c < combinations; c++ {
		g.Printf("case %d: return (*rw%d)(state)\n", c, c)
	}
	g.Printf("}\n")
	g.Printf("panic(\"unreachable\")")
	g.Printf("}\n\n")

	// rwState holds the underlying writer plus the precomputed hooks.
	// All variant types are type-definitions over rwState, so a single *rwState
	// allocation can be reinterpreted as any variant via pointer conversion.
	g.Printf("type rwState struct {\n")
	g.Printf("w http.ResponseWriter\n")
	for _, iface := range ifaces {
		for _, fn := range iface.Funcs {
			g.Printf("%s %s\n", fieldName(fn.Name), fn.Type())
		}
	}
	g.Printf("}\n\n")

	// do<Name> helpers on *rwState
	// These actual dispatch logic, defined once and called by the variant types.
	for _, iface := range ifaces {
		for _, fn := range iface.Funcs {
			g.Printf("func (r *rwState) do%s(%s) (%s) {\n", fn.Name, fn.Args, fn.Returns)
			g.Printf("if r.%s != nil {\n", fieldName(fn.Name))
			if fn.Returns != "" {
				g.Printf("return r.%s(%s)\n", fieldName(fn.Name), fn.Args.Names())
			} else {
				g.Printf("r.%s(%s)\n", fieldName(fn.Name), fn.Args.Names())
				g.Printf("return\n")
			}
			g.Printf("}\n")

			receiver := "r.w"
			if iface.Name != "http.ResponseWriter" {
				receiver = fmt.Sprintf("r.w.(%s)", iface.Name)
			}
			if fn.Returns != "" {
				g.Printf("return %s.%s(%s)\n", receiver, fn.Name, fn.Args.Names())
			} else {
				g.Printf("%s.%s(%s)\n", receiver, fn.Name, fn.Args.Names())
			}
			g.Printf("}\n\n")
		}
	}

	// Variant types, each is a type with the same memory layout as rwState,
	// but exposing exactly the method set required by its combination of interfaces.
	// This allows (*rwN)(state) to be a zero-cost pointer conversion.
	emitVariantMethod := func(c int, fn *InterfaceFunc) {
		g.Printf("func (w *rw%d) %s(%s) (%s) {\n", c, fn.Name, fn.Args, fn.Returns)
		if fn.Returns != "" {
			g.Printf("return (*rwState)(w).do%s(%s)\n", fn.Name, fn.Args.Names())
		} else {
			g.Printf("(*rwState)(w).do%s(%s)\n", fn.Name, fn.Args.Names())
		}
		g.Printf("}\n")
	}
	for c := 0; c < combinations; c++ {
		supported := []string{"http.ResponseWriter"}
		for j, iface := range subIfaces {
			if c&(1<<uint(len(subIfaces)-j-1)) > 0 {
				supported = append(supported, iface.Name)
			}
		}
		g.Printf("// combination %d/%d: %s\n", c+1, combinations, strings.Join(supported, ", "))
		g.Printf("type rw%d rwState\n", c)
		g.Printf("func (w *rw%d) Unwrap() http.ResponseWriter { return w.w }\n", c)
		for _, fn := range ifaces[0].Funcs {
			emitVariantMethod(c, fn)
		}
		for j, iface := range subIfaces {
			if c&(1<<uint(len(subIfaces)-j-1)) > 0 {
				for _, fn := range iface.Funcs {
					emitVariantMethod(c, fn)
				}
			}
		}
		g.Printf("\n")
	}
	g.Printf(`
type Unwrapper interface {
	Unwrap() http.ResponseWriter
}

// Unwrap returns the underlying http.ResponseWriter from within zero or more
// layers of httpsnoop wrappers.
func Unwrap(w http.ResponseWriter) http.ResponseWriter {
	if rw, ok := w.(Unwrapper); ok {
		// recurse until rw.Unwrap() returns a non-Unwrapper
		return Unwrap(rw.Unwrap())
	}
	return w
}
`)
	return &g
}

func (b *Build) Tests() *Generator {
	ifaces := b.Interfaces
	// @TODO dedupe
	// subIfaces has all interfaces except http.ResponseWriter
	subIfaces := ifaces[1:]

	var g Generator
	// Package header
	b.writeHeader(&g)
	g.Printf("import (\n")
	g.Printf(`"net/http"` + "\n")
	g.Printf(`"io"` + "\n")
	g.Printf(`"testing"` + "\n")
	g.Printf(")\n")
	g.Printf("\n")

	// TestWrap func
	g.Printf("func TestWrap(t *testing.T) {\n")
	combinations := 1 << uint(len(subIfaces))
	for i := 0; i < combinations; i++ {
		fields := make([]string, 0, len(subIfaces))
		fields = append(fields, "http.ResponseWriter")
		expected := make([]bool, len(ifaces))
		expected[0] = true
		for j, iface := range subIfaces {
			ok := i&(1<<uint(len(subIfaces)-j-1)) > 0
			expected[j+1] = ok
			if ok {
				fields = append(fields, iface.Name)
			}
		}
		g.Printf("// combination %d/%d\n", i+1, combinations)
		g.Printf("{\n")
		g.Printf(`t.Log("%s")`+"\n", strings.Join(fields, ", "))
		g.Printf("inner := struct{\n%s\n}{}\n", strings.Join(fields, "\n"))
		g.Printf("w := Wrap(inner, Hooks{})\n")
		for i, iface := range ifaces {
			g.Printf("if _, ok := w.(%s); ok != %t {\n", iface.Name, expected[i])
			g.Printf("t.Error(\"unexpected interface\");\n")
			g.Printf("}\n")
		}
		g.Printf(`
if w, ok := w.(Unwrapper); ok {
  if w.Unwrap() != inner {
    t.Error("w.Unwrap() failed")
  }
} else {
	t.Error("Unwrapper interface not implemented")
}`)
		g.Printf("}\n")
		g.Printf("\n")
	}
	g.Printf("}\n")
	return &g
}

type Interfaces []*Interface

type Interface struct {
	Name  string
	Funcs []*InterfaceFunc
}

type InterfaceFunc struct {
	Name    string
	Args    FuncArgs
	Returns string
}

type FuncArgs []*FuncArg

func (fa FuncArgs) String() string {
	args := make([]string, len(fa))
	for i, a := range fa {
		args[i] = a.Name + " " + a.Type
	}
	return strings.Join(args, ", ")
}

func (fa FuncArgs) Names() string {
	args := make([]string, len(fa))
	for i, a := range fa {
		args[i] = a.Name
	}
	return strings.Join(args, ", ")
}

type FuncArg struct {
	Name string
	Type string
}

func (fn *InterfaceFunc) Type() string {
	return fn.Name + "Func"
}

func fieldName(s string) string {
	if s == "" {
		return s
	}
	return strings.ToLower(s[:1]) + s[1:]
}

type Generator struct {
	buf bytes.Buffer
}

func (g *Generator) Printf(s string, args ...interface{}) {
	_, _ = fmt.Fprintf(&g.buf, s, args...)
}

func (g *Generator) WriteFile(name string) error {
	src, err := g.Format()
	if err != nil {
		return fmt.Errorf("format: %s: %s:\n\n%s\n", name, err, g.Bytes())
	} else if err := os.WriteFile(name, src, 0o644); err != nil {
		return err
	}
	return nil
}

func (g *Generator) MustWriteFile(name string) {
	if err := g.WriteFile(name); err != nil {
		fatalf("%s", err)
	}
}

func (g *Generator) Bytes() []byte {
	return g.buf.Bytes()
}

func (g *Generator) Format() ([]byte, error) {
	return format.Source(g.Bytes())
}

func main() {
	ifaces := Interfaces{
		{
			Name: "http.ResponseWriter",
			Funcs: []*InterfaceFunc{
				{"Header", nil, "http.Header"},
				{"WriteHeader", FuncArgs{{"code", "int"}}, ""},
				{"Write", FuncArgs{{"b", "[]byte"}}, "int, error"},
			},
		},
		{
			Name: "http.Flusher",
			Funcs: []*InterfaceFunc{
				{"Flush", nil, ""},
			},
		},
		{
			Name: "httpFlushError", // Introduced in Go 1.20.
			Funcs: []*InterfaceFunc{
				{"FlushError", nil, "error"},
			},
		},
		{
			Name: "http.CloseNotifier",
			Funcs: []*InterfaceFunc{
				{"CloseNotify", nil, "<-chan bool"},
			},
		},
		{
			Name: "http.Hijacker",
			Funcs: []*InterfaceFunc{
				{"Hijack", nil, "net.Conn, *bufio.ReadWriter, error"},
			},
		},
		{
			Name: "io.ReaderFrom",
			Funcs: []*InterfaceFunc{
				{"ReadFrom", FuncArgs{{"src", "io.Reader"}}, "int64, error"},
			},
		},
		{
			Name: "deadliner", // Introduced in Go 1.20.
			Funcs: []*InterfaceFunc{
				{"SetReadDeadline", FuncArgs{{"deadline", "time.Time"}}, "error"},
				{"SetWriteDeadline", FuncArgs{{"deadline", "time.Time"}}, "error"},
			},
		},
		{
			Name: "fullDuplexEnabler", // Introduced in Go 1.21.
			Funcs: []*InterfaceFunc{
				{"EnableFullDuplex", nil, "error"},
			},
		},
		{
			Name: "http.Pusher",
			Funcs: []*InterfaceFunc{
				{"Push", FuncArgs{
					{"target", "string"},
					{"opts", "*http.PushOptions"},
				}, "error"},
			},
		}, {
			Name: "io.StringWriter",
			Funcs: []*InterfaceFunc{
				{"WriteString", FuncArgs{{"s", "string"}}, "int, error"},
			},
		},
	}
	builds := []Build{
		{
			Interfaces: ifaces,
		},
	}
	for _, build := range builds {
		build.MustBuild()
	}
}

func fatalf(s string, args ...interface{}) {
	_, _ = fmt.Fprintf(os.Stderr, s+"\n", args...)
	os.Exit(1)
}