package main

import (
	"bytes"
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/sha256"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"math/big"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"
	"neuralhunt/internal/auth"
	"neuralhunt/internal/core"
)

var b64 = base64.RawURLEncoding

type apiClient struct {
	base       string
	hc         *http.Client
	token, cid string
	key        *ecdsa.PrivateKey
}
type taskDTO struct {
	ID         string `json:"id"`
	PublicSeed string `json:"public_seed"`
	RangeBits  int    `json:"range_bits"`
	NextSeq    int64  `json:"next_seq"`
	SubmitSec  int    `json:"client_submit_interval_sec"`
	Paused     bool   `json:"paused"`
}

func pad32(x *big.Int) []byte {
	b := x.Bytes()
	out := make([]byte, 32)
	copy(out[32-len(b):], b)
	return out
}
func signRaw(k *ecdsa.PrivateKey, msg string) (string, error) {
	h := sha256.Sum256([]byte(msg))
	r, s, err := ecdsa.Sign(rand.Reader, k, h[:])
	if err != nil {
		return "", err
	}
	raw := append(pad32(r), pad32(s)...)
	return b64.EncodeToString(raw), nil
}
func (c *apiClient) do(method, path string, body any, out any) error {
	var rd io.Reader
	if body != nil {
		b, _ := json.Marshal(body)
		rd = bytes.NewReader(b)
	}
	req, _ := http.NewRequest(method, c.base+path, rd)
	req.Header.Set("Content-Type", "application/json")
	if c.token != "" {
		req.Header.Set("Authorization", "Bearer "+c.token)
	}
	resp, err := c.hc.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	b, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if resp.StatusCode/100 != 2 {
		return fmt.Errorf("%s: %s", resp.Status, string(b))
	}
	if out != nil {
		return json.Unmarshal(b, out)
	}
	return nil
}

func leadingZeroBitsLoad(b []byte) int {
	n := 0
	for _, x := range b {
		if x == 0 {
			n += 8
			continue
		}
		for m := byte(0x80); m != 0 && x&m == 0; m >>= 1 {
			n++
		}
		break
	}
	return n
}

func solveProofLoad(challenge, cid string, bits int) string {
	if bits <= 0 {
		return ""
	}
	for i := uint64(0); ; i++ {
		counter := fmt.Sprint(i)
		h := sha256.Sum256([]byte("nh-pow-v1|" + challenge + "|" + cid + "|" + counter))
		if leadingZeroBitsLoad(h[:]) >= bits {
			return counter
		}
	}
}
func (c *apiClient) authn() error {
	k, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	if err != nil {
		return err
	}
	c.key = k
	j := auth.PublicJWK{Kty: "EC", Crv: "P-256", X: b64.EncodeToString(pad32(k.X)), Y: b64.EncodeToString(pad32(k.Y)), Ext: true}
	var ch struct {
		ClientID        string `json:"client_id"`
		Challenge       string `json:"challenge"`
		ProofOfWorkBits int    `json:"proof_of_work_bits"`
	}
	if err = c.do("POST", "/api/auth/challenge", map[string]any{"public_jwk": j}, &ch); err != nil {
		return err
	}
	c.cid = ch.ClientID
	sig, _ := signRaw(k, "login|"+ch.Challenge+"|"+c.cid)
	var lg struct {
		Token string `json:"token"`
	}
	pow := solveProofLoad(ch.Challenge, c.cid, ch.ProofOfWorkBits)
	if err = c.do("POST", "/api/auth/login", map[string]any{"public_jwk": j, "challenge": ch.Challenge, "signature": sig, "proof_of_work_counter": pow}, &lg); err != nil {
		return err
	}
	c.token = lg.Token
	return nil
}
func (c *apiClient) current() (taskDTO, error) {
	var t taskDTO
	err := c.do("GET", "/api/tasks/current", nil, &t)
	return t, err
}
func (c *apiClient) ws(ctx context.Context, maxNodes int) (*websocket.Conn, error) {
	u, err := url.Parse(c.base)
	if err != nil {
		return nil, err
	}
	scheme := "ws"
	if u.Scheme == "https" {
		scheme = "wss"
	}
	q := url.Values{}
	q.Set("max_nodes", fmt.Sprint(maxNodes))
	wu := scheme + "://" + u.Host + "/api/ws?" + q.Encode()
	h := http.Header{}
	h.Set("Authorization", "Bearer "+c.token)
	conn, _, err := websocket.DefaultDialer.DialContext(ctx, wu, h)
	return conn, err
}

func main() {
	base := flag.String("url", "http://127.0.0.1:8080", "server URL")
	clients := flag.Int("clients", 1000, "virtual clients")
	ramp := flag.Duration("ramp", 30*time.Second, "connection ramp")
	duration := flag.Duration("duration", 2*time.Minute, "test duration after ramp")
	nodes := flag.Int("max-nodes", 250, "snapshot budget per client")
	flag.Parse()
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	var connected, guesses, errs atomic.Uint64
	var wg sync.WaitGroup
	start := time.Now()
	step := time.Duration(0)
	if *clients > 0 {
		step = *ramp / time.Duration(*clients)
	}
	for i := 0; i < *clients; i++ {
		wg.Add(1)
		go func(i int) {
			defer wg.Done()
			if step > 0 {
				select {
				case <-ctx.Done():
					return
				case <-time.After(step * time.Duration(i)):
				}
			}
			c := &apiClient{base: strings.TrimRight(*base, "/"), hc: &http.Client{Timeout: 10 * time.Second}}
			if err := c.authn(); err != nil {
				errs.Add(1)
				return
			}
			t, err := c.current()
			if err != nil {
				errs.Add(1)
				return
			}
			ws, err := c.ws(ctx, *nodes)
			if err != nil {
				errs.Add(1)
				return
			}
			defer ws.Close()
			connected.Add(1)
			go func() {
				for {
					if _, _, e := ws.ReadMessage(); e != nil {
						return
					}
				}
			}()
			interval := time.Duration(t.SubmitSec) * time.Second
			if interval <= 0 {
				interval = 11 * time.Second
			}
			ticker := time.NewTicker(interval)
			defer ticker.Stop()
			seq := t.NextSeq
			for {
				select {
				case <-ctx.Done():
					return
				case <-ticker.C:
					if t.Paused {
						continue
					}
					guess := core.ExpectedGuess(t.ID, t.PublicSeed, c.cid, seq, t.RangeBits)
					sig, _ := signRaw(c.key, fmt.Sprintf("guess|%s|%d|%s", t.ID, seq, guess))
					var ok bool
					err := c.do("POST", "/api/tasks/"+t.ID+"/guess", map[string]any{"seq": seq, "guess": guess, "signature": sig}, &ok)
					if err != nil {
						errs.Add(1)
						nt, e := c.current()
						if e == nil {
							t = nt
							seq = t.NextSeq
						}
						continue
					}
					guesses.Add(1)
					seq++
					if ok {
						nt, e := c.current()
						if e == nil {
							t = nt
							seq = t.NextSeq
						}
					}
				}
			}
		}(i)
	}
	tick := time.NewTicker(5 * time.Second)
	defer tick.Stop()
	end := time.After(*ramp + *duration)
	for {
		select {
		case <-end:
			cancel()
			wg.Wait()
			fmt.Printf("done connected=%d accepted_guesses=%d errors=%d elapsed=%s\n", connected.Load(), guesses.Load(), errs.Load(), time.Since(start).Round(time.Second))
			return
		case <-tick.C:
			log.Printf("connected=%d guesses=%d errors=%d", connected.Load(), guesses.Load(), errs.Load())
		}
	}
}