// tera-target-session-load exercises the deployed Direct Mode WebSocket path.
// It deliberately uses the same Origin and subprotocol requirements as the web UI.
package main

import (
	"bytes"
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/gorilla/websocket"

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

type options struct {
	baseURL      *url.URL
	origin       string
	username     string
	password     string
	access       string
	attachments  int
	startChannel int
	sendPayload  []byte
	hold         time.Duration
	tlsConfig    *tls.Config
}

type client struct {
	http  *http.Client
	base  *url.URL
	token string
}

type loadConnection struct {
	session domain.DirectSession
	socket  *websocket.Conn
	done    chan error
}

type report struct {
	State         string     `json:"state"`
	Attachments   int        `json:"attachments"`
	StartedAt     time.Time  `json:"started_at"`
	CompletedAt   *time.Time `json:"completed_at,omitempty"`
	HoldSeconds   int        `json:"hold_seconds"`
	SessionDetail string     `json:"session_service_detail,omitempty"`
	WelcomeFrames int        `json:"welcome_frames"`
	SessionLogs   int        `json:"session_logs"`
	TXBytes       int        `json:"tx_bytes,omitempty"`
	TXLogs        int        `json:"tx_logs,omitempty"`
}

type serviceStatus struct {
	Service string `json:"service"`
	Detail  string `json:"detail"`
}

func main() {
	if err := run(); err != nil {
		fmt.Fprintf(os.Stderr, "target Direct Mode load failed: %v\n", err)
		os.Exit(1)
	}
}

func run() error {
	configuration, err := parseOptions()
	if err != nil {
		return err
	}
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	api := &client{
		http: &http.Client{Timeout: 15 * time.Second, Transport: &http.Transport{TLSClientConfig: configuration.tlsConfig}},
		base: configuration.baseURL,
	}
	login := domain.AuthSession{}
	if err := api.request(ctx, http.MethodPost, "/api/v1/auth/login", domain.AuthLoginRequest{
		Username: configuration.username,
		Password: configuration.password,
	}, &login); err != nil {
		return err
	}
	api.token = login.SessionToken
	defer func() {
		logoutCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		_ = api.request(logoutCtx, http.MethodPost, "/api/v1/auth/logout", nil, nil)
	}()

	startedAt := time.Now().UTC()
	connections := make([]loadConnection, 0, configuration.attachments)
	defer func() { cleanup(connections) }()
	for index := 0; index < configuration.attachments; index++ {
		if err := ctx.Err(); err != nil {
			return err
		}
		channel := fmt.Sprintf("CH%02d", (configuration.startChannel-1+index)%48+1)
		session := domain.DirectSession{}
		if err := api.request(ctx, http.MethodPost, "/api/v1/direct/"+channel+"/sessions", domain.DirectSessionRequest{
			Access: configuration.access,
		}, &session); err != nil {
			return fmt.Errorf("create attachment %d/%d: %w", index+1, configuration.attachments, err)
		}
		socket, readDone, err := connect(ctx, configuration, session)
		if err != nil {
			return fmt.Errorf("connect attachment %d/%d (%s): %w", index+1, configuration.attachments, channel, err)
		}
		// The ticket is a one-use secret and is no longer needed after authentication.
		session.StreamTicket = ""
		connections = append(connections, loadConnection{session: session, socket: socket, done: readDone})
	}
	sessionLogs, err := validateSessionLogs(ctx, api, connections)
	if err != nil {
		return err
	}
	txLogs := 0
	if len(configuration.sendPayload) > 0 {
		if err := connections[0].socket.WriteMessage(websocket.BinaryMessage, configuration.sendPayload); err != nil {
			return fmt.Errorf("send Direct Mode payload: %w", err)
		}
		txLogs, err = validateTrafficLog(ctx, api, connections[0].session.Channel, "TX", configuration.sendPayload, startedAt)
		if err != nil {
			return err
		}
	}

	serviceDetail := api.sessionServiceDetail(ctx)
	writeReport(report{
		State: "holding", Attachments: len(connections), StartedAt: startedAt,
		HoldSeconds: int(configuration.hold.Seconds()), SessionDetail: serviceDetail,
		WelcomeFrames: len(connections), SessionLogs: sessionLogs,
		TXBytes: len(configuration.sendPayload), TXLogs: txLogs,
	})

	timer := time.NewTimer(configuration.hold)
	defer timer.Stop()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-timer.C:
			validatedConnections := len(connections)
			cleanup(connections)
			connections = nil
			if err := waitForSessionCleanup(ctx, api); err != nil {
				return err
			}
			completedAt := time.Now().UTC()
			writeReport(report{
				State: "completed", Attachments: validatedConnections, StartedAt: startedAt,
				CompletedAt: &completedAt, HoldSeconds: int(configuration.hold.Seconds()),
				SessionDetail: api.sessionServiceDetail(context.Background()),
				WelcomeFrames: validatedConnections, SessionLogs: sessionLogs,
				TXBytes: len(configuration.sendPayload), TXLogs: txLogs,
			})
			return nil
		default:
		}

		for index := range connections {
			select {
			case readErr := <-connections[index].done:
				if readErr == nil {
					readErr = errors.New("connection closed unexpectedly")
				}
				return fmt.Errorf("attachment %d/%d (%s) ended during hold: %w", index+1, len(connections), connections[index].session.Channel, readErr)
			default:
			}
		}
		time.Sleep(100 * time.Millisecond)
	}
}

func parseOptions() (options, error) {
	baseValue := flag.String("base-url", "http://192.168.0.223:8083", "target web base URL")
	originValue := flag.String("origin", "", "WebSocket Origin (defaults to the base URL origin)")
	username := flag.String("username", "terauser", "web username")
	access := flag.String("access", directstream.AccessViewer, "Direct Mode access: viewer or writer")
	attachments := flag.Int("attachments", 192, "number of viewer attachments: 1, 48, 96, or 192")
	startChannel := flag.Int("start-channel", 1, "first Direct Mode channel: 1-48")
	sendText := flag.String("send-text", "", "bounded UTF-8 payload sent by one writer attachment")
	sendCR := flag.Bool("send-cr", false, "append carriage return to send-text")
	holdSeconds := flag.Int("hold-seconds", 60, "hold time in seconds")
	tlsCA := flag.String("tls-ca", "", "PEM CA certificate used to verify an HTTPS/WSS target")
	flag.Parse()

	password := os.Getenv("TERA_TARGET_PASSWORD")
	if password == "" {
		return options{}, errors.New("TERA_TARGET_PASSWORD is required")
	}
	allowedCount := *attachments == 1 || *attachments == 48 || *attachments == 96 || *attachments == 192
	if !allowedCount {
		return options{}, errors.New("attachments must be 1, 48, 96, or 192")
	}
	if *startChannel < 1 || *startChannel > 48 {
		return options{}, errors.New("start-channel must be between 1 and 48")
	}
	if *access != directstream.AccessViewer && *access != directstream.AccessWriter {
		return options{}, errors.New("access must be viewer or writer")
	}
	sendPayload := []byte(*sendText)
	if *sendCR {
		sendPayload = append(sendPayload, '\r')
	}
	if len(sendPayload) > directstream.MaxDataFrameBytes {
		return options{}, fmt.Errorf("send-text payload must not exceed %d bytes", directstream.MaxDataFrameBytes)
	}
	if len(sendPayload) > 0 && (*access != directstream.AccessWriter || *attachments != 1) {
		return options{}, errors.New("send-text requires access=writer and attachments=1")
	}
	if *holdSeconds < 1 || *holdSeconds > 3600 {
		return options{}, errors.New("hold-seconds must be between 1 and 3600")
	}
	base, err := url.Parse(*baseValue)
	if err != nil || (base.Scheme != "http" && base.Scheme != "https") || base.Host == "" {
		return options{}, errors.New("base-url must be an absolute HTTP or HTTPS URL")
	}
	base.Path, base.RawPath, base.RawQuery, base.Fragment = "", "", "", ""
	origin := strings.TrimSpace(*originValue)
	if origin == "" {
		origin = base.Scheme + "://" + base.Host
	}
	tlsConfig, err := trustedTLSConfig(base, strings.TrimSpace(*tlsCA))
	if err != nil {
		return options{}, err
	}
	return options{
		baseURL: base, origin: origin, username: *username, password: password, access: *access,
		attachments: *attachments, startChannel: *startChannel, hold: time.Duration(*holdSeconds) * time.Second,
		sendPayload: sendPayload, tlsConfig: tlsConfig,
	}, nil
}

func trustedTLSConfig(base *url.URL, caPath string) (*tls.Config, error) {
	if base.Scheme != "https" {
		if caPath != "" {
			return nil, errors.New("tls-ca is only valid with an HTTPS base URL")
		}
		return nil, nil
	}
	if caPath == "" {
		return nil, nil
	}
	payload, err := os.ReadFile(caPath)
	if err != nil {
		return nil, fmt.Errorf("read tls-ca: %w", err)
	}
	roots, err := x509.SystemCertPool()
	if err != nil || roots == nil {
		roots = x509.NewCertPool()
	}
	if !roots.AppendCertsFromPEM(payload) {
		return nil, errors.New("tls-ca does not contain a PEM certificate")
	}
	return &tls.Config{MinVersion: tls.VersionTLS12, RootCAs: roots}, nil
}

func (c *client) request(ctx context.Context, method, path string, requestBody, responseBody any) error {
	var body io.Reader
	if requestBody != nil {
		encoded, err := json.Marshal(requestBody)
		if err != nil {
			return err
		}
		body = bytes.NewReader(encoded)
	}
	reference, err := url.Parse(path)
	if err != nil {
		return err
	}
	endpoint := c.base.ResolveReference(reference)
	request, err := http.NewRequestWithContext(ctx, method, endpoint.String(), body)
	if err != nil {
		return err
	}
	if requestBody != nil {
		request.Header.Set("Content-Type", "application/json")
	}
	if c.token != "" {
		request.Header.Set("Authorization", "Bearer "+c.token)
	}
	response, err := c.http.Do(request)
	if err != nil {
		return err
	}
	defer response.Body.Close()
	payload, err := io.ReadAll(io.LimitReader(response.Body, 1<<20))
	if err != nil {
		return err
	}
	if response.StatusCode < http.StatusOK || response.StatusCode >= http.StatusMultipleChoices {
		return fmt.Errorf("%s %s returned HTTP %d: %s", method, path, response.StatusCode, strings.TrimSpace(string(payload)))
	}
	if responseBody == nil || len(payload) == 0 {
		return nil
	}
	if err := json.Unmarshal(payload, responseBody); err != nil {
		return fmt.Errorf("decode %s %s response: %w", method, path, err)
	}
	return nil
}

func (c *client) sessionServiceDetail(ctx context.Context) string {
	status := []serviceStatus{}
	if err := c.request(ctx, http.MethodGet, "/api/v1/system/services?limit=100", nil, &status); err != nil {
		return "unavailable: " + err.Error()
	}
	for _, service := range status {
		if service.Service == "tera-sessiond" {
			return service.Detail
		}
	}
	return "tera-sessiond missing"
}

func waitForSessionCleanup(ctx context.Context, api *client) error {
	deadline := time.Now().Add(10 * time.Second)
	for {
		detail := api.sessionServiceDetail(ctx)
		if strings.Contains(detail, "attachments=0") && strings.Contains(detail, "active_channels=0") && strings.Contains(detail, "serial_leases=0") {
			return nil
		}
		if time.Now().After(deadline) {
			return fmt.Errorf("Direct Mode resources did not return to zero: %s", detail)
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(100 * time.Millisecond):
		}
	}
}

func connect(ctx context.Context, configuration options, session domain.DirectSession) (*websocket.Conn, chan error, error) {
	streamURL, err := configuration.baseURL.Parse(session.StreamURL)
	if err != nil {
		return nil, nil, err
	}
	if streamURL.Scheme == "https" {
		streamURL.Scheme = "wss"
	} else {
		streamURL.Scheme = "ws"
	}
	dialer := websocket.Dialer{
		HandshakeTimeout: 10 * time.Second,
		Subprotocols:     []string{session.StreamProtocol},
		TLSClientConfig:  configuration.tlsConfig,
	}
	header := http.Header{"Origin": []string{configuration.origin}}
	socket, response, err := dialer.DialContext(ctx, streamURL.String(), header)
	if err != nil {
		status := 0
		if response != nil {
			status = response.StatusCode
		}
		return nil, nil, fmt.Errorf("WebSocket upgrade returned HTTP %d: %w", status, err)
	}
	if socket.Subprotocol() != session.StreamProtocol {
		socket.Close()
		return nil, nil, fmt.Errorf("negotiated subprotocol %q, want %q", socket.Subprotocol(), session.StreamProtocol)
	}
	if err := socket.WriteJSON(directstream.ControlFrame{Type: directstream.ControlAuthenticate, Ticket: session.StreamTicket}); err != nil {
		socket.Close()
		return nil, nil, err
	}
	_ = socket.SetReadDeadline(time.Now().Add(10 * time.Second))
	var status directstream.WebSocketStatusFrame
	if err := socket.ReadJSON(&status); err != nil {
		socket.Close()
		return nil, nil, err
	}
	if status.Type != "session.status" || status.Channel != session.Channel || status.SessionID != session.SessionID {
		socket.Close()
		return nil, nil, fmt.Errorf("unexpected authentication status type=%q channel=%q", status.Type, status.Channel)
	}
	messageType, welcome, err := socket.ReadMessage()
	if err != nil {
		socket.Close()
		return nil, nil, err
	}
	if messageType != websocket.BinaryMessage ||
		!strings.Contains(string(welcome), "Welcome to TERA CORA-X48 PRO K") ||
		!strings.Contains(string(welcome), "Serial channel: "+session.Channel) ||
		!strings.Contains(string(welcome), "Session: "+session.SessionID) ||
		!messageContainsSourceIP(string(welcome)) {
		socket.Close()
		return nil, nil, fmt.Errorf("invalid Direct Mode welcome frame for %s", session.Channel)
	}
	_ = socket.SetReadDeadline(time.Time{})
	done := make(chan error, 1)
	go func() {
		for {
			if _, _, err := socket.ReadMessage(); err != nil {
				done <- err
				return
			}
		}
	}()
	return socket, done, nil
}

func validateSessionLogs(ctx context.Context, api *client, connections []loadConnection) (int, error) {
	expected := make(map[string]map[string]bool)
	for _, connection := range connections {
		if expected[connection.session.Channel] == nil {
			expected[connection.session.Channel] = make(map[string]bool)
		}
		expected[connection.session.Channel][connection.session.SessionID] = false
	}
	matched := 0
	for channel, sessions := range expected {
		events := []domain.PortLogEvent{}
		path := "/api/v1/logs/port?channel=" + url.QueryEscape(channel) + "&event_type=session&limit=100"
		if err := api.request(ctx, http.MethodGet, path, nil, &events); err != nil {
			return matched, fmt.Errorf("read %s session logs: %w", channel, err)
		}
		for _, event := range events {
			if event.EventType != "session" || !strings.Contains(event.Message, "Welcome to TERA CORA-X48 PRO K") ||
				!strings.Contains(event.Message, "Serial channel: "+channel) || !messageContainsSourceIP(event.Message) {
				continue
			}
			for sessionID, found := range sessions {
				if !found && strings.Contains(event.Message, "Session: "+sessionID) {
					sessions[sessionID] = true
					matched++
				}
			}
		}
		for sessionID, found := range sessions {
			if !found {
				return matched, fmt.Errorf("session log missing for %s %s", channel, sessionID)
			}
		}
	}
	return matched, nil
}

func validateTrafficLog(ctx context.Context, api *client, channel, direction string, payload []byte, since time.Time) (int, error) {
	deadline := time.Now().Add(5 * time.Second)
	for {
		events := []domain.PortLogEvent{}
		path := "/api/v1/logs/port?channel=" + url.QueryEscape(channel) + "&event_type=traffic&direction=" + url.QueryEscape(direction) + "&limit=100"
		if err := api.request(ctx, http.MethodGet, path, nil, &events); err != nil {
			return 0, fmt.Errorf("read %s traffic logs: %w", channel, err)
		}
		matched := 0
		for _, event := range events {
			if event.Channel == channel && event.EventType == "traffic" && event.Direction == direction &&
				event.PayloadBytes == len(payload) && !event.Timestamp.Before(since.Add(-5*time.Second)) &&
				(len(event.Payload) == 0 || bytes.Equal(event.Payload, payload)) {
				matched++
			}
		}
		if matched > 0 {
			return matched, nil
		}
		if time.Now().After(deadline) {
			return 0, fmt.Errorf("%s %s traffic log for %d bytes was not observed", channel, direction, len(payload))
		}
		select {
		case <-ctx.Done():
			return 0, ctx.Err()
		case <-time.After(100 * time.Millisecond):
		}
	}
}

func messageContainsSourceIP(message string) bool {
	const marker = "Connected from: "
	index := strings.Index(message, marker)
	if index < 0 {
		return false
	}
	fields := strings.Fields(message[index+len(marker):])
	return len(fields) > 0 && net.ParseIP(fields[0]) != nil
}

func cleanup(connections []loadConnection) {
	for index := range connections {
		connection := &connections[index]
		_ = connection.socket.WriteControl(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, "load test complete"), time.Now().Add(time.Second))
		_ = connection.socket.Close()
	}
}

func writeReport(value report) {
	encoder := json.NewEncoder(os.Stdout)
	encoder.SetEscapeHTML(false)
	_ = encoder.Encode(value)
}