package main

import (
	"bytes"
	"fmt"
	"go/ast"
	"go/format"
	"go/parser"
	"go/token"
	"go/types"
	"log"
	"os"
	"strings"
)

func main() {
	if len(os.Args) != 3 {
		fmt.Printf("Usage: %s gendb/querier.go gendb/otelwrap.go", os.Args[0])
	}

	input := os.Args[1]
	output := os.Args[2]

	fileSet := token.NewFileSet()
	file, err := parser.ParseFile(fileSet, input, nil, parser.ParseComments)
	if err != nil {
		log.Fatal(err)
	}

	var buf bytes.Buffer

	buf.WriteString(fmt.Sprintf("package %s\n\n", file.Name.Name))
	for _, importSpec := range file.Imports {
		buf.WriteString(fmt.Sprintf("import %s\n", importSpec.Path.Value))
	}
	buf.WriteString("import \"go.opentelemetry.io/otel\"\n")
	buf.WriteString("import \"go.opentelemetry.io/otel/trace\"\n\n")

	for _, decl := range file.Decls {
		iface, ok := decl.(*ast.GenDecl)
		if !ok || iface.Tok != token.TYPE {
			continue
		}

		for _, declSpec := range iface.Specs {
			typeSpec, ok := declSpec.(*ast.TypeSpec)
			if !ok {
				continue
			}

			ifaceType, ok := typeSpec.Type.(*ast.InterfaceType)
			if !ok {
				continue
			}

			typeName := typeSpec.Name.Name
			variable := strings.ToLower(typeName[:1]) + typeName[1:]
			wrapperName := typeName + "OtelWrapper"

			buf.WriteString(fmt.Sprintf("type %s struct {\n", wrapperName))
			buf.WriteString(fmt.Sprintf("\t%s %s\n", variable, typeName))
			buf.WriteString("}\n\n")

			buf.WriteString(fmt.Sprintf("func NewOtelWrapper(%s %s) %s {\n", variable, typeName, typeName))
			buf.WriteString(fmt.Sprintf("\treturn &%s{%s: %s}\n", wrapperName, variable, variable))
			buf.WriteString("}\n\n")

			buf.WriteString(fmt.Sprintf("func (w *%s) GetQuerier() Querier {\n", wrapperName))
			buf.WriteString("\treturn w.querier\n")
			buf.WriteString("}\n\n")

			for _, methodNode := range ifaceType.Methods.List {
				methodName := methodNode.Names[0].Name
				buf.WriteString(fmt.Sprintf("func (w *%s) %s(", wrapperName, methodName))

				method, ok := methodNode.Type.(*ast.FuncType)
				if !ok {
					continue
				}

				for i, param := range method.Params.List {
					if i > 0 {
						buf.WriteString(", ")
					}
					for _, paramName := range param.Names {
						buf.WriteString(fmt.Sprintf("%s %s", paramName.Name, types.ExprString(param.Type)))
					}
				}
				buf.WriteString(")")

				if method.Results != nil {
					buf.WriteString(" (")
					for i, result := range method.Results.List {
						if i > 0 {
							buf.WriteString(", ")
						}
						buf.WriteString(types.ExprString(result.Type))
					}
					buf.WriteString(")")
				}

				// Write method body with OpenTelemetry span
				buf.WriteString(" {\n")
				buf.WriteString("\tctx, span := otel.Tracer(\"db\").Start(ctx, \"" + methodName + "\", trace.WithSpanKind(trace.SpanKindClient))\n")
				buf.WriteString("\tdefer span.End()\n\n")
				buf.WriteString("\treturn w." + variable + "." + methodName + "(")

				// Write method parameters for the wrapped method call
				for i, param := range method.Params.List {
					if i > 0 {
						buf.WriteString(", ")
					}
					for _, paramName := range param.Names {
						buf.WriteString(paramName.Name)
					}
				}
				buf.WriteString(")\n")
				buf.WriteString("}\n\n")
			}
		}
	}

	gofmt, err := format.Source(buf.Bytes())
	if err != nil {
		// Write the unformatted source to file and fail afterwards
		gofmt = buf.Bytes()
		defer log.Fatal(err)
	}

	if err := os.WriteFile(output, gofmt, 0644); err != nil {
		log.Fatal(err)
	}
}