package main

import (
	"context"
	"database/sql"
	"database/sql/driver"
	"errors"
	"fmt"
	"os"
	"reflect"
	"runtime"
	"runtime/debug"
	"runtime/pprof"
	"strings"
	"sync"
	"syscall"
	"time"
)

type fakeDriver struct{}
type fakeConn struct{}

func (fakeDriver) Open(string) (driver.Conn, error)  { return fakeConn{}, nil }
func (fakeConn) Prepare(string) (driver.Stmt, error) { return nil, errors.New("not supported") }
func (fakeConn) Close() error                        { return nil }
func (fakeConn) Begin() (driver.Tx, error)           { return nil, errors.New("not supported") }
func (fakeConn) ExecContext(ctx context.Context, _ string, _ []driver.NamedValue) (driver.Result, error) {
	timer := time.NewTimer(30 * time.Millisecond)
	defer timer.Stop()
	select {
	case <-ctx.Done():
		return nil, ctx.Err()
	case <-timer.C:
		return driver.RowsAffected(1), nil
	}
}

var registerDriver sync.Once

func pool() error {
	registerDriver.Do(func() { sql.Register("fault-demo", fakeDriver{}) })
	db, err := sql.Open("fault-demo", "")
	if err != nil {
		return err
	}
	defer db.Close()
	db.SetMaxOpenConns(1)
	c, err := db.Conn(context.Background())
	if err != nil {
		return err
	}
	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
	_, waitErr := db.Conn(ctx)
	cancel()
	s := db.Stats()
	fmt.Printf("pool: deadline=%v inuse=%d waitcount=%d waited=%v\n", errors.Is(waitErr, context.DeadlineExceeded), s.InUse, s.WaitCount, s.WaitDuration > 0)
	if !errors.Is(waitErr, context.DeadlineExceeded) || s.WaitCount != 1 {
		c.Close()
		return errors.New("pool assertion")
	}
	if err := c.Close(); err != nil {
		return err
	}
	c, err = db.Conn(context.Background())
	if err != nil {
		return err
	}
	fmt.Println("pool: acquired after close=true")
	if err := c.Close(); err != nil {
		return err
	}
	before := db.Stats().WaitCount
	ctx, cancel = context.WithTimeout(context.Background(), 10*time.Millisecond)
	_, queryErr := db.ExecContext(ctx, "simulated slow operation")
	cancel()
	fmt.Printf("query: deadline=%v pool_wait_delta=%d\n", errors.Is(queryErr, context.DeadlineExceeded), db.Stats().WaitCount-before)
	if !errors.Is(queryErr, context.DeadlineExceeded) || db.Stats().WaitCount != before {
		return errors.New("query assertion")
	}
	return nil
}

var held [][]byte
var trash []byte

func memory() error {
	runtime.GC()
	var start, live, end runtime.MemStats
	runtime.ReadMemStats(&start)
	for range 64 {
		b := make([]byte, 64<<10)
		for i := range b {
			b[i] = 1
		}
		held = append(held, b)
	}
	runtime.GC()
	runtime.ReadMemStats(&live)
	held = nil
	runtime.GC()
	runtime.GC()
	runtime.ReadMemStats(&end)
	fmt.Printf("heap: retained_growth_over_3MiB=%v released_over_3MiB=%v\n", live.HeapAlloc > start.HeapAlloc+3<<20, live.HeapAlloc > end.HeapAlloc+3<<20)
	if live.HeapAlloc <= start.HeapAlloc+3<<20 || live.HeapAlloc <= end.HeapAlloc+3<<20 {
		return errors.New("heap assertion")
	}
	old := debug.SetGCPercent(20)
	defer debug.SetGCPercent(old)
	runtime.ReadMemStats(&start)
	for range 256 {
		trash = make([]byte, 128<<10)
	}
	trash = nil
	runtime.ReadMemStats(&end)
	fmt.Printf("gc: allocated_over_32MiB=%v automatic_collections=%v\n", end.TotalAlloc-start.TotalAlloc >= 32<<20, end.NumGC > start.NumGC)
	if end.TotalAlloc-start.TotalAlloc < 32<<20 || end.NumGC <= start.NumGC {
		return errors.New("gc assertion")
	}
	return nil
}

//go:noinline
func stuck(ctx context.Context, ready chan<- struct{}, wg *sync.WaitGroup) {
	defer wg.Done()
	ready <- struct{}{}
	<-ctx.Done()
}
func goroutines() error {
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	ready := make(chan struct{}, 3)
	var wg sync.WaitGroup
	for range 3 {
		wg.Add(1)
		go stuck(ctx, ready, &wg)
	}
	for range 3 {
		<-ready
	}
	var b strings.Builder
	if err := pprof.Lookup("goroutine").WriteTo(&b, 2); err != nil {
		return err
	}
	name := runtime.FuncForPC(reflect.ValueOf(stuck).Pointer()).Name()
	found := strings.Contains(b.String(), name+"(")
	fmt.Printf("goroutine: stuck_stack_found=%v\n", found)
	cancel()
	wg.Wait()
	fmt.Println("goroutine: joined=3")
	if !found {
		return errors.New("missing stuck stack")
	}
	return nil
}

var total uint64

//go:noinline
func busy() {
	until := time.Now().Add(200 * time.Millisecond)
	x := uint64(1)
	for time.Now().Before(until) {
		for range 10000 {
			x = x*1664525 + 1013904223
		}
	}
	total = x
}
func cpu(path string) error {
	f, err := os.Create(path)
	if err != nil {
		return err
	}
	if err := pprof.StartCPUProfile(f); err != nil {
		f.Close()
		return err
	}
	busy()
	pprof.StopCPUProfile()
	return f.Close()
}
func fd() error {
	var lim syscall.Rlimit
	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &lim); err != nil {
		return err
	}
	if lim.Cur > 64 {
		lim.Cur = 64
	}
	if err := syscall.Setrlimit(syscall.RLIMIT_NOFILE, &lim); err != nil {
		return err
	}
	files := []*os.File{}
	defer func() {
		for _, f := range files {
			f.Close()
		}
	}()
	for range 128 {
		f, err := os.Open("/dev/null")
		if err != nil {
			fmt.Printf("fd: emfile=%v opened=%d\n", errors.Is(err, syscall.EMFILE), len(files))
			if errors.Is(err, syscall.EMFILE) {
				return nil
			}
			return err
		}
		files = append(files, f)
	}
	return errors.New("limit not reached within bound")
}
func run() error {
	if len(os.Args) < 2 {
		return errors.New("mode required")
	}
	switch os.Args[1] {
	case "pool":
		return pool()
	case "memory":
		return memory()
	case "goroutine":
		return goroutines()
	case "cpu":
		if len(os.Args) != 3 {
			return errors.New("profile path required")
		}
		return cpu(os.Args[2])
	case "fd":
		return fd()
	case "deadlock":
		var ch chan int
		<-ch
	default:
		return errors.New("unknown mode")
	}
	return nil
}
func main() {
	if err := run(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}