package cli

import (
	"context"
	"errors"
	"net"
	"net/http"
	"os/signal"
	"syscall"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/sim"
)

func newSimEngineCmd() *cobra.Command {
	cfg := sim.Default()
	var listen string

	cmd := &cobra.Command{
		Use:   "sim-engine",
		Short: "OpenAI-compatible fake inference engine (for testing)",
		Long: "Simulates the engine behaviours that matter to the router:\n" +
			"continuous batching, KV block accounting, prefix caching, and vLLM-shaped\n" +
			"/metrics. Lets the router be developed and tested without GPUs.",
		RunE: func(cmd *cobra.Command, _ []string) error {
			log := logger()
			eng := sim.New(cfg)
			defer eng.Close()

			srv := &http.Server{
				Addr:              listen,
				Handler:           sim.NewServer(eng, log),
				ReadHeaderTimeout: 10 * time.Second,
			}

			ln, err := net.Listen("tcp", listen)
			if err != nil {
				return err
			}
			log.Info("sim-engine listening",
				"addr", ln.Addr().String(),
				"model", cfg.Model,
				"max_num_seqs", cfg.MaxNumSeqs,
				"num_gpu_blocks", cfg.NumGPUBlocks,
				"prefix_cache_blocks", cfg.PrefixCacheBlocks,
			)

			ctx, stop := signal.NotifyContext(cmd.Context(), syscall.SIGINT, syscall.SIGTERM)
			defer stop()

			errCh := make(chan error, 1)
			go func() { errCh <- srv.Serve(ln) }()

			select {
			case err := <-errCh:
				if errors.Is(err, http.ErrServerClosed) {
					return nil
				}
				return err
			case <-ctx.Done():
				log.Info("shutting down")
				shutCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
				defer cancel()
				return srv.Shutdown(shutCtx)
			}
		},
	}

	f := cmd.Flags()
	f.StringVar(&listen, "listen", ":8000", "address to listen on")
	f.StringVar(&cfg.Model, "model", cfg.Model, "model name to serve")
	f.IntVar(&cfg.MaxNumSeqs, "max-num-seqs", cfg.MaxNumSeqs, "concurrent sequences (continuous batch capacity)")
	f.IntVar(&cfg.BlockSize, "block-size", cfg.BlockSize, "tokens per KV cache block")
	f.IntVar(&cfg.NumGPUBlocks, "num-gpu-blocks", cfg.NumGPUBlocks, "total KV cache blocks")
	f.IntVar(&cfg.PrefixCacheBlocks, "prefix-cache-blocks", cfg.PrefixCacheBlocks, "prefix cache capacity in blocks (0 disables)")
	f.Float64Var(&cfg.PrefillTokensPerSec, "prefill-tps", cfg.PrefillTokensPerSec, "prefill rate in tokens/sec, shared engine-wide")
	f.Float64Var(&cfg.DecodeTokensPerSec, "decode-tps", cfg.DecodeTokensPerSec, "single-sequence decode rate in tokens/sec")
	f.Float64Var(&cfg.BatchContention, "batch-contention", cfg.BatchContention, "per-extra-sequence decode slowdown factor")
	f.IntVar(&cfg.DefaultOutputTokens, "default-output-tokens", cfg.DefaultOutputTokens, "tokens generated when max_tokens is absent")
	f.DurationVar(&cfg.Tick, "tick", cfg.Tick, "simulation step")
	f.Float64Var(&cfg.FailRate, "fail-rate", 0, "fraction of requests failed artificially (chaos testing)")
	f.Int64Var(&cfg.Seed, "seed", 0, "random seed (0 = time-based)")

	return cmd
}