// tera-direct-load records repeatable in-process Direct Mode broker resource evidence.
package main

import (
	"context"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"runtime/metrics"
	"strconv"
	"strings"
	"sync"
	"time"

	"tera-clone/services/tera-api/internal/directstream"
)

type loadEvidence struct {
	AttachmentsRequested int       `json:"attachments_requested"`
	ActiveAttachments    int       `json:"active_attachments"`
	ActiveChannels       int       `json:"active_channels"`
	QueuedBytes          uint64    `json:"queued_bytes"`
	QueueCapacityBytes   uint64    `json:"queue_capacity_bytes"`
	ResidentBytes        uint64    `json:"resident_bytes"`
	MemorySource         string    `json:"memory_source"`
	HeapAllocBytes       uint64    `json:"heap_alloc_bytes"`
	GoSystemBytes        uint64    `json:"go_system_bytes"`
	CPUSeconds           float64   `json:"cpu_seconds"`
	WallMilliseconds     float64   `json:"wall_milliseconds"`
	RXBytes              uint64    `json:"rx_bytes"`
	FanoutBytes          uint64    `json:"fanout_bytes"`
	SlowConsumerDrops    uint64    `json:"slow_consumer_drops"`
	RejectedAttachments  uint64    `json:"rejected_attachments"`
	TransportReconnects  uint64    `json:"transport_reconnects"`
	GeneratedAt          time.Time `json:"generated_at"`
}

type loadReport struct {
	SchemaVersion int            `json:"schema_version"`
	Runtime       string         `json:"runtime"`
	Profiles      []loadEvidence `json:"profiles"`
}

func main() {
	output := flag.String("output", "", "optional JSON evidence output path")
	payloadBytes := flag.Int("payload-bytes", 4096, "queued RX bytes per active channel")
	flag.Parse()
	if *payloadBytes < 1 || *payloadBytes > 64*1024 {
		fmt.Fprintln(os.Stderr, "payload-bytes must be between 1 and 65536")
		os.Exit(2)
	}
	report := loadReport{SchemaVersion: 1, Runtime: runtime.GOOS + "/" + runtime.GOARCH}
	for _, attachments := range []int{1, 48, 96, 192} {
		evidence, err := runLoadProfile(attachments, *payloadBytes)
		if err != nil {
			fmt.Fprintf(os.Stderr, "load profile %d failed: %v\n", attachments, err)
			os.Exit(1)
		}
		report.Profiles = append(report.Profiles, evidence)
	}
	payload, err := json.MarshalIndent(report, "", "  ")
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	if *output != "" {
		if err := os.MkdirAll(filepath.Dir(*output), 0o750); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
		if err := os.WriteFile(*output, append(payload, '\n'), 0o640); err != nil {
			fmt.Fprintln(os.Stderr, err)
			os.Exit(1)
		}
	}
	fmt.Println(string(payload))
}

func runLoadProfile(attachments, payloadBytes int) (loadEvidence, error) {
	if attachments != 1 && attachments != 48 && attachments != 96 && attachments != 192 {
		return loadEvidence{}, errors.New("attachment profile must be 1, 48, 96, or 192")
	}
	if payloadBytes < 1 || payloadBytes > 64*1024 {
		return loadEvidence{}, errors.New("payload size is outside the broker queue budget")
	}
	runtime.GC()
	cpuBefore := processCPUSeconds()
	started := time.Now()

	factory := newLoadTransportFactory()
	broker, err := directstream.NewBroker(factory, directstream.BrokerOptions{})
	if err != nil {
		return loadEvidence{}, err
	}
	defer broker.Close()
	warmup, err := broker.Attach(context.Background(), directstream.AttachRequest{
		SessionID: "load-warmup", Channel: "CH01", Actor: "load-test", Access: directstream.AccessViewer,
	})
	if err != nil {
		return loadEvidence{}, err
	}
	if err := warmup.Close(); err != nil {
		return loadEvidence{}, err
	}

	activeChannels := 48
	if attachments == 1 {
		activeChannels = 1
	}
	attached := make([]*directstream.Attachment, 0, attachments)
	defer func() {
		for _, attachment := range attached {
			_ = attachment.Close()
		}
	}()
	for index := 0; index < attachments; index++ {
		channel := fmt.Sprintf("CH%02d", index%activeChannels+1)
		attachment, attachErr := broker.Attach(context.Background(), directstream.AttachRequest{
			SessionID: fmt.Sprintf("load-%03d-%03d", attachments, index),
			Channel:   channel, Actor: "load-test", Access: directstream.AccessViewer,
		})
		if attachErr != nil {
			return loadEvidence{}, attachErr
		}
		attached = append(attached, attachment)
	}
	payload := make([]byte, payloadBytes)
	for index := range payload {
		payload[index] = byte(index)
	}
	for channel := 1; channel <= activeChannels; channel++ {
		if err := factory.emit(fmt.Sprintf("CH%02d", channel), payload); err != nil {
			return loadEvidence{}, err
		}
	}
	expectedRX := uint64(activeChannels * payloadBytes)
	expectedQueued := uint64(attachments * payloadBytes)
	deadline := time.Now().Add(3 * time.Second)
	var snapshot directstream.BrokerSnapshot
	for {
		snapshot = broker.Snapshot()
		queued := queuedBytes(snapshot)
		if snapshot.RXBytes == expectedRX && queued == expectedQueued {
			break
		}
		if time.Now().After(deadline) {
			return loadEvidence{}, fmt.Errorf("load did not settle: rx=%d/%d queued=%d/%d", snapshot.RXBytes, expectedRX, queued, expectedQueued)
		}
		time.Sleep(time.Millisecond)
	}
	// Exercise the same bounded snapshot path used by health polling so the CPU
	// measurement has useful resolution even for the one-attachment profile.
	observationSamples := 2000 * 48 / activeChannels
	for sample := 0; sample < observationSamples; sample++ {
		snapshot = broker.Snapshot()
	}
	var memory runtime.MemStats
	runtime.ReadMemStats(&memory)
	resident, source := residentBytes(memory.Sys)
	return loadEvidence{
		AttachmentsRequested: attachments,
		ActiveAttachments:    snapshot.TotalAttachments, ActiveChannels: snapshot.ActiveChannels,
		QueuedBytes: queuedBytes(snapshot), QueueCapacityBytes: uint64(attachments * 64 * 1024),
		ResidentBytes: resident, MemorySource: source,
		HeapAllocBytes: memory.HeapAlloc, GoSystemBytes: memory.Sys,
		CPUSeconds:       processCPUSeconds() - cpuBefore,
		WallMilliseconds: float64(time.Since(started).Microseconds()) / 1000,
		RXBytes:          snapshot.RXBytes, FanoutBytes: snapshot.FanoutBytes,
		SlowConsumerDrops:   snapshot.SlowConsumerDisconnects,
		RejectedAttachments: snapshot.RejectedAttachments,
		TransportReconnects: snapshot.TransportReconnects,
		GeneratedAt:         time.Now().UTC(),
	}, nil
}

func queuedBytes(snapshot directstream.BrokerSnapshot) uint64 {
	var total uint64
	for _, channel := range snapshot.Channels {
		total += uint64(channel.QueuedBytes)
	}
	return total
}

func processCPUSeconds() float64 {
	samples := []metrics.Sample{{Name: "/cpu/classes/total:cpu-seconds"}}
	metrics.Read(samples)
	if samples[0].Value.Kind() != metrics.KindFloat64 {
		return 0
	}
	return samples[0].Value.Float64()
}

func residentBytes(fallback uint64) (uint64, string) {
	if runtime.GOOS != "linux" {
		return fallback, "go_runtime_sys"
	}
	payload, err := os.ReadFile("/proc/self/statm")
	if err != nil {
		return fallback, "go_runtime_sys"
	}
	fields := strings.Fields(string(payload))
	if len(fields) < 2 {
		return fallback, "go_runtime_sys"
	}
	pages, err := strconv.ParseUint(fields[1], 10, 64)
	if err != nil {
		return fallback, "go_runtime_sys"
	}
	return pages * uint64(os.Getpagesize()), "proc_statm_rss"
}

type loadTransportFactory struct {
	mu         sync.Mutex
	transports map[string]*loadTransport
}

func newLoadTransportFactory() *loadTransportFactory {
	return &loadTransportFactory{transports: make(map[string]*loadTransport)}
}

func (f *loadTransportFactory) Open(_ context.Context, channel string) (io.ReadWriteCloser, error) {
	transport := newLoadTransport()
	f.mu.Lock()
	f.transports[channel] = transport
	f.mu.Unlock()
	return transport, nil
}

func (f *loadTransportFactory) emit(channel string, payload []byte) error {
	f.mu.Lock()
	transport := f.transports[channel]
	f.mu.Unlock()
	if transport == nil {
		return errors.New("load transport is missing for " + channel)
	}
	return transport.emit(payload)
}

type loadTransport struct {
	rx        chan []byte
	closed    chan struct{}
	closeOnce sync.Once
	pending   []byte
}

func newLoadTransport() *loadTransport {
	return &loadTransport{rx: make(chan []byte, 1), closed: make(chan struct{})}
}

func (t *loadTransport) Read(buffer []byte) (int, error) {
	for len(t.pending) == 0 {
		select {
		case payload := <-t.rx:
			t.pending = payload
		case <-t.closed:
			return 0, io.EOF
		}
	}
	count := copy(buffer, t.pending)
	t.pending = t.pending[count:]
	return count, nil
}

func (t *loadTransport) Write(payload []byte) (int, error) {
	select {
	case <-t.closed:
		return 0, io.ErrClosedPipe
	default:
		return len(payload), nil
	}
}

func (t *loadTransport) Close() error {
	t.closeOnce.Do(func() { close(t.closed) })
	return nil
}

func (t *loadTransport) emit(payload []byte) error {
	copyOfPayload := append([]byte(nil), payload...)
	select {
	case t.rx <- copyOfPayload:
		return nil
	case <-t.closed:
		return io.ErrClosedPipe
	}
}