package gateway

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/url"
	"strconv"
	"strings"
	"time"

	protocol "git.sechmachine.io.vn/sechmachine/VerseVDI-Protocol/gen/go/protocol"
)

const (
	maxApolloRTSPHeaders = 16 << 10
	maxApolloRTSPBody    = 48 << 10
	maxApolloRTSPPort    = 65535
	apolloEncryptionAll  = 0x07
)

type apolloRTSPMessage struct {
	status  int
	cseq    uint32
	headers map[string]string
	body    []byte
	raw     []byte
}

type apolloRTSPSetup struct {
	sessionID      string
	audioPort      int
	videoPort      int
	controlPort    int
	audioPing      []byte
	videoPing      []byte
	controlConnect uint32
	streamHost     string
	streamPort     int64
	providerWork   protocol.ProviderSessionWork
	streamKey      []byte
	streamKeyID    uint32
}

func (b *NativeApolloBackend) performRTSPHandshake(ctx context.Context, work protocol.ProviderSessionWork, key []byte, keyID uint32, streamURL *url.URL) (*apolloRTSPSetup, []byte, error) {
	if streamURL == nil || streamURL.Scheme != "rtspenc" || streamURL.Hostname() != work.StreamHost || streamURL.Port() != strconv.FormatInt(work.StreamPort, 10) || streamURL.User != nil || streamURL.RawQuery != "" || streamURL.Fragment != "" {
		return nil, nil, ErrProviderMalformed
	}
	codec, err := newEncryptedRTSPCodec(key)
	if err != nil {
		return nil, nil, err
	}
	request := func(method, target, session string, headers []apolloRTSPHeader, body []byte, sequence uint32) (apolloRTSPMessage, error) {
		return b.encryptedRTSPRequest(ctx, work, codec, method, target, session, headers, body, sequence)
	}
	options, err := request("OPTIONS", streamURL.String(), "", nil, nil, 1)
	if err != nil {
		return nil, nil, err
	}
	describe, err := request("DESCRIBE", streamURL.String(), "", []apolloRTSPHeader{{"Accept", "application/sdp"}, {"If-Modified-Since", "Thu, 01 Jan 1970 00:00:00 GMT"}}, nil, 2)
	if err != nil {
		return nil, nil, err
	}
	if err := validateApolloDescribe(describe); err != nil {
		return nil, nil, err
	}
	setupHeaders := []apolloRTSPHeader{{"Transport", "unicast;X-GS-ClientPort=50000-50001"}, {"If-Modified-Since", "Thu, 01 Jan 1970 00:00:00 GMT"}}
	audio, err := request("SETUP", "streamid=audio/0/0", "", setupHeaders, nil, 3)
	if err != nil {
		return nil, nil, err
	}
	sessionID, err := apolloRTSPSession(audio)
	if err != nil {
		return nil, nil, err
	}
	audioPort, err := apolloRTSPServerPort(audio)
	if err != nil {
		return nil, nil, err
	}
	audioPing, err := apolloRTSPPingPayload(audio)
	if err != nil {
		return nil, nil, err
	}
	video, err := request("SETUP", "streamid=video/0/0", sessionID, setupHeaders, nil, 4)
	if err != nil {
		return nil, nil, err
	}
	if err := apolloRTSPMatchSession(video, sessionID); err != nil {
		return nil, nil, err
	}
	videoPort, err := apolloRTSPServerPort(video)
	if err != nil {
		return nil, nil, err
	}
	videoPing, err := apolloRTSPPingPayload(video)
	if err != nil {
		return nil, nil, err
	}
	control, err := request("SETUP", "streamid=control/13/0", sessionID, setupHeaders, nil, 5)
	if err != nil {
		return nil, nil, err
	}
	if err := apolloRTSPMatchSession(control, sessionID); err != nil {
		return nil, nil, err
	}
	controlPort, err := apolloRTSPServerPort(control)
	if err != nil {
		return nil, nil, err
	}
	connectData, err := apolloRTSPConnectData(control)
	if err != nil {
		return nil, nil, err
	}
	announceBody, err := apolloAnnounceProfile(work.StreamPolicy)
	if err != nil {
		return nil, nil, err
	}
	announce, err := request("ANNOUNCE", "streamid=control/13/0", sessionID, []apolloRTSPHeader{{"Content-Type", "application/sdp"}}, announceBody, 6)
	if err != nil {
		return nil, nil, err
	}
	if err := apolloRTSPMatchSession(announce, sessionID); err != nil {
		return nil, nil, err
	}
	play, err := request("PLAY", "/", sessionID, nil, nil, 7)
	if err != nil {
		return nil, nil, err
	}
	if err := apolloRTSPMatchSession(play, sessionID); err != nil {
		return nil, nil, err
	}
	_ = options
	setup := &apolloRTSPSetup{
		sessionID: sessionID, audioPort: audioPort, videoPort: videoPort, controlPort: controlPort,
		audioPing: audioPing, videoPing: videoPing, controlConnect: connectData, streamHost: work.StreamHost,
		streamPort: work.StreamPort, providerWork: work, streamKey: append([]byte(nil), key...), streamKeyID: keyID,
	}
	return setup, append([]byte(nil), control.raw...), nil
}

type apolloRTSPHeader struct{ key, value string }

func (b *NativeApolloBackend) encryptedRTSPRequest(ctx context.Context, work protocol.ProviderSessionWork, codec *encryptedRTSPCodec, method, target, session string, headers []apolloRTSPHeader, body []byte, sequence uint32) (apolloRTSPMessage, error) {
	if b == nil || b.Dialer == nil || codec == nil || sequence == 0 || len(body) > maxApolloRTSPBody || method == "" || target == "" {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	conn, err := b.Dialer.DialContext(ctx, "tcp", net.JoinHostPort(work.StreamHost, strconv.FormatInt(work.StreamPort, 10)))
	if err != nil {
		return apolloRTSPMessage{}, err
	}
	defer conn.Close()
	deadline := time.Now().Add(5 * time.Second)
	if contextDeadline, ok := ctx.Deadline(); ok && contextDeadline.Before(deadline) {
		deadline = contextDeadline
	}
	if err := conn.SetDeadline(deadline); err != nil {
		return apolloRTSPMessage{}, err
	}
	plaintext, err := buildApolloRTSPRequest(method, target, session, headers, body, sequence)
	if err != nil {
		return apolloRTSPMessage{}, err
	}
	frame, err := codec.SealClient(plaintext)
	if err != nil {
		return apolloRTSPMessage{}, err
	}
	if _, err := conn.Write(frame); err != nil {
		return apolloRTSPMessage{}, err
	}
	response, err := readEncryptedRTSPMessage(conn, codec)
	if err != nil {
		return apolloRTSPMessage{}, err
	}
	if response.status != 200 || response.cseq != sequence {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	return response, nil
}

func buildApolloRTSPRequest(method, target, session string, headers []apolloRTSPHeader, body []byte, sequence uint32) ([]byte, error) {
	if method == "" || target == "" || strings.ContainsAny(method, "\r\n ") || strings.ContainsAny(target, "\r\n") || sequence == 0 {
		return nil, ErrProviderMalformed
	}
	var builder strings.Builder
	builder.Grow(256 + len(body))
	fmt.Fprintf(&builder, "%s %s RTSP/1.0\r\nCSeq: %d\r\n", method, target, sequence)
	if session != "" {
		if !validApolloRTSPToken(session) {
			return nil, ErrProviderMalformed
		}
		fmt.Fprintf(&builder, "Session: %s\r\n", session)
	}
	seen := map[string]struct{}{"cseq": {}, "session": {}}
	for _, header := range headers {
		key := strings.ToLower(header.key)
		if !validApolloRTSPToken(header.key) || header.value == "" || len(header.value) > 1024 || strings.ContainsAny(header.value, "\r\n") {
			return nil, ErrProviderMalformed
		}
		if _, ok := seen[key]; ok {
			return nil, ErrProviderMalformed
		}
		seen[key] = struct{}{}
		fmt.Fprintf(&builder, "%s: %s\r\n", header.key, header.value)
	}
	if len(body) != 0 {
		if _, ok := seen["content-length"]; ok {
			return nil, ErrProviderMalformed
		}
		fmt.Fprintf(&builder, "Content-Length: %d\r\n", len(body))
	}
	builder.WriteString("\r\n")
	builder.Write(body)
	return []byte(builder.String()), nil
}

func readEncryptedRTSPMessage(conn net.Conn, codec *encryptedRTSPCodec) (apolloRTSPMessage, error) {
	header := make([]byte, encryptedRTSPHeaderSize)
	if _, err := io.ReadFull(conn, header); err != nil {
		return apolloRTSPMessage{}, err
	}
	length := uint64(header[0]&0x7f)<<24 | uint64(header[1])<<16 | uint64(header[2])<<8 | uint64(header[3])
	if header[0]&0x80 == 0 || length == 0 || length > encryptedRTSPMaxPayload {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	frame := make([]byte, encryptedRTSPHeaderSize+int(length))
	copy(frame, header)
	if _, err := io.ReadFull(conn, frame[encryptedRTSPHeaderSize:]); err != nil {
		return apolloRTSPMessage{}, err
	}
	plaintext, err := codec.OpenHost(frame)
	if err != nil {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	return parseApolloRTSPMessage(plaintext)
}

func parseApolloRTSPMessage(data []byte) (apolloRTSPMessage, error) {
	if len(data) == 0 || len(data) > encryptedRTSPMaxPayload {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	headerEnd := strings.Index(string(data), "\r\n\r\n")
	if headerEnd < 0 || headerEnd+4 > maxApolloRTSPHeaders {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	lines := strings.Split(string(data[:headerEnd]), "\r\n")
	if len(lines) < 1 {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	parts := strings.SplitN(lines[0], " ", 3)
	if len(parts) != 3 || parts[0] != "RTSP/1.0" || len(parts[2]) == 0 || len(parts[2]) > 128 {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	status, err := strconv.Atoi(parts[1])
	if err != nil || status < 100 || status > 599 {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	message := apolloRTSPMessage{status: status, headers: make(map[string]string), raw: append([]byte(nil), data...)}
	for _, line := range lines[1:] {
		key, value, ok := strings.Cut(line, ":")
		key = strings.ToLower(strings.TrimSpace(key))
		value = strings.TrimSpace(value)
		if !ok || !validApolloRTSPToken(key) || value == "" || len(value) > 1024 {
			return apolloRTSPMessage{}, ErrProviderMalformed
		}
		if _, duplicate := message.headers[key]; duplicate {
			return apolloRTSPMessage{}, ErrProviderMalformed
		}
		message.headers[key] = value
	}
	cseq, ok := message.headers["cseq"]
	if !ok {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	parsedCSeq, err := strconv.ParseUint(cseq, 10, 32)
	if err != nil || parsedCSeq == 0 {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	message.cseq = uint32(parsedCSeq)
	message.body = append([]byte(nil), data[headerEnd+4:]...)
	if len(message.body) > maxApolloRTSPBody {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	if length, hasLength := message.headers["content-length"]; hasLength {
		declared, err := strconv.ParseUint(length, 10, 16)
		if err != nil || int(declared) != len(message.body) {
			return apolloRTSPMessage{}, ErrProviderMalformed
		}
	} else if len(message.body) != 0 {
		return apolloRTSPMessage{}, ErrProviderMalformed
	}
	return message, nil
}

func validateApolloDescribe(message apolloRTSPMessage) error {
	if message.headers["content-type"] != "application/sdp" || len(message.body) == 0 {
		return ErrProviderMalformed
	}
	body := string(message.body)
	lineEnding := "\n"
	if strings.Contains(body, "\r\n") {
		if strings.Contains(strings.ReplaceAll(body, "\r\n", ""), "\n") {
			return ErrProviderMalformed
		}
		lineEnding = "\r\n"
	}
	if !strings.HasSuffix(body, lineEnding) {
		return ErrProviderMalformed
	}
	attributes := map[string]string{}
	seen := map[string]struct{}{}
	stage := 0
	stereo := false
	for _, line := range strings.Split(strings.TrimSuffix(body, lineEnding), lineEnding) {
		if line == "" {
			return ErrProviderMalformed
		}
		if !strings.HasPrefix(line, "a=") {
			if line == "sprop-parameter-sets=AAAAAU" && stage >= 3 && stage <= 4 {
				stage = 5
				continue
			}
			return ErrProviderMalformed
		}
		key, value, ok := strings.Cut(strings.TrimPrefix(line, "a="), ":")
		if !ok || key == "" || value == "" || len(key) > 128 || len(value) > 256 {
			return ErrProviderMalformed
		}
		if key == "fmtp" {
			if stage < 3 || !validApolloSurroundParameters(value) {
				return ErrProviderMalformed
			}
			if _, duplicate := seen["fmtp:"+value]; duplicate {
				return ErrProviderMalformed
			}
			seen["fmtp:"+value] = struct{}{}
			if value == "97 surround-params=21101" {
				stereo = true
			}
			stage = 6
			continue
		}
		if key == "rtpmap" && value == "98 AV1/90000" && stage >= 3 && stage <= 5 {
			if _, duplicate := seen[key]; duplicate {
				return ErrProviderMalformed
			}
			seen[key] = struct{}{}
			stage = 6
			continue
		}
		if key == "x-nv-video[0].refPicInvalidation" && value == "1" && stage == 3 {
			if _, duplicate := seen[key]; duplicate {
				return ErrProviderMalformed
			}
			seen[key] = struct{}{}
			stage = 4
			continue
		}
		if key != "x-ss-general.featureFlags" && key != "x-ss-general.encryptionSupported" && key != "x-ss-general.encryptionRequested" {
			return ErrProviderMalformed
		}
		if _, duplicate := attributes[key]; duplicate || (key == "x-ss-general.featureFlags" && stage != 0) || (key == "x-ss-general.encryptionSupported" && stage != 1) || (key == "x-ss-general.encryptionRequested" && stage != 2) {
			return ErrProviderMalformed
		}
		attributes[key] = value
		stage++
	}
	featureFlags, featureFlagsOK := attributes["x-ss-general.featureFlags"]
	supported, supportedOK := attributes["x-ss-general.encryptionSupported"]
	requested, requestedOK := attributes["x-ss-general.encryptionRequested"]
	if !featureFlagsOK || !supportedOK || !requestedOK || !stereo {
		return ErrProviderMalformed
	}
	if _, err := strconv.ParseUint(featureFlags, 10, 32); err != nil {
		return ErrProviderMalformed
	}
	supportedFlags, err := strconv.ParseUint(supported, 10, 32)
	if err != nil || supportedFlags&apolloEncryptionAll != apolloEncryptionAll {
		return ErrProviderMalformed
	}
	requestedFlags, err := strconv.ParseUint(requested, 10, 32)
	if err != nil || requestedFlags&^supportedFlags != 0 || requestedFlags&1 == 0 {
		return ErrProviderMalformed
	}
	return nil
}

func validApolloSurroundParameters(value string) bool {
	const prefix = "97 surround-params="
	if !strings.HasPrefix(value, prefix) {
		return false
	}
	parameters := strings.TrimPrefix(value, prefix)
	if len(parameters) < 5 || len(parameters) > 11 {
		return false
	}
	channels := int(parameters[0] - '0')
	streams := int(parameters[1] - '0')
	coupled := int(parameters[2] - '0')
	if (channels != 2 && channels != 6 && channels != 8) || len(parameters) != channels+3 || streams+coupled != channels || streams == 0 {
		return false
	}
	used := [8]bool{}
	for _, character := range parameters[3:] {
		if character < '0' || int(character-'0') >= channels || used[character-'0'] {
			return false
		}
		used[character-'0'] = true
	}
	return true
}

func apolloRTSPSession(message apolloRTSPMessage) (string, error) {
	value, ok := message.headers["session"]
	if !ok {
		return "", ErrProviderMalformed
	}
	token, _, _ := strings.Cut(value, ";")
	token = strings.TrimSpace(token)
	if !validApolloRTSPToken(token) || len(token) > 256 {
		return "", ErrProviderMalformed
	}
	return token, nil
}

func apolloRTSPMatchSession(message apolloRTSPMessage, expected string) error {
	actual, err := apolloRTSPSession(message)
	if err != nil || actual != expected {
		return ErrProviderMalformed
	}
	return nil
}

func apolloRTSPServerPort(message apolloRTSPMessage) (int, error) {
	transport, ok := message.headers["transport"]
	if !ok {
		return 0, ErrProviderMalformed
	}
	parts := strings.Split(transport, ";")
	if len(parts) != 2 || parts[0] != "unicast" {
		return 0, ErrProviderMalformed
	}
	key, value, ok := strings.Cut(parts[1], "=")
	port, err := strconv.Atoi(value)
	if !ok || key != "server_port" || err != nil || port < 1 || port > maxApolloRTSPPort {
		return 0, ErrProviderMalformed
	}
	return port, nil
}

func apolloRTSPPingPayload(message apolloRTSPMessage) ([]byte, error) {
	payload, ok := message.headers["x-ss-ping-payload"]
	if !ok || len(payload) != 16 || !validApolloRTSPToken(payload) {
		return nil, ErrProviderMalformed
	}
	return []byte(payload), nil
}

func apolloRTSPConnectData(message apolloRTSPMessage) (uint32, error) {
	value, ok := message.headers["x-ss-connect-data"]
	if !ok {
		return 0, ErrProviderMalformed
	}
	if value == "" || strings.Trim(value, "0123456789") != "" {
		return 0, ErrProviderMalformed
	}
	parsed, err := strconv.ParseUint(value, 10, 32)
	if err != nil || parsed == 0 {
		return 0, ErrProviderMalformed
	}
	return uint32(parsed), nil
}

func apolloAnnounceProfile(policy protocol.ProviderStreamPolicy) ([]byte, error) {
	if err := validateApolloStreamPolicy(policy); err != nil {
		return nil, err
	}
	format, supportsHEVC := int64(0), int64(0)
	if policy.Codec == "HEVC" {
		format, supportsHEVC = 1, 1
	}
	maximumBitrate := policy.BitrateKbps * 80 / 100
	if maximumBitrate > 100000 {
		maximumBitrate = 100000
	}
	return []byte("v=0\r\n" +
		"o=android 0 0 IN IP4 0.0.0.0\r\n" +
		"s=NVIDIA Streaming Client\r\n" +
		fmt.Sprintf("a=x-nv-video[0].clientViewportWd:%d\r\n", policy.ResolutionWidth) +
		fmt.Sprintf("a=x-nv-video[0].clientViewportHt:%d\r\n", policy.ResolutionHeight) +
		fmt.Sprintf("a=x-nv-video[0].maxFPS:%d\r\n", policy.Fps) +
		"a=x-nv-video[0].packetSize:1024\r\n" +
		"a=x-nv-video[0].videoEncoderSlicesPerFrame:1\r\n" +
		"a=x-nv-video[0].maxNumReferenceFrames:0\r\n" +
		fmt.Sprintf("a=x-nv-clientSupportHevc:%d\r\n", supportsHEVC) +
		fmt.Sprintf("a=x-nv-vqos[0].bitStreamFormat:%d\r\n", format) +
		fmt.Sprintf("a=x-nv-video[0].initialBitrateKbps:%d\r\n", maximumBitrate) +
		fmt.Sprintf("a=x-nv-video[0].initialPeakBitrateKbps:%d\r\n", maximumBitrate) +
		fmt.Sprintf("a=x-nv-vqos[0].bw.minimumBitrateKbps:%d\r\n", maximumBitrate) +
		fmt.Sprintf("a=x-nv-vqos[0].bw.maximumBitrateKbps:%d\r\n", maximumBitrate) +
		"a=x-nv-vqos[0].fec.minRequiredFecPackets:2\r\n" +
		"a=x-nv-vqos[0].qosTrafficType:5\r\n" +
		"a=x-nv-audio.surround.numChannels:2\r\n" +
		"a=x-nv-audio.surround.channelMask:3\r\n" +
		"a=x-nv-audio.surround.AudioQuality:0\r\n" +
		"a=x-nv-aqos.packetDuration:5\r\n" +
		"a=x-nv-aqos.qosTrafficType:4\r\n" +
		"a=x-nv-general.useReliableUdp:13\r\n" +
		"a=x-nv-general.featureFlags:167\r\n" +
		"a=x-ml-general.featureFlags:0\r\n" +
		fmt.Sprintf("a=x-ml-video.configuredBitrateKbps:%d\r\n", policy.BitrateKbps) +
		"a=x-ss-general.encryptionEnabled:7\r\n" +
		"a=x-ss-video[0].chromaSamplingType:0\r\n" +
		"a=x-ss-video[0].intraRefresh:0\r\n"), nil
}

func validateApolloStreamPolicy(policy protocol.ProviderStreamPolicy) error {
	if err := policy.Validate(); err != nil || !policy.AudioEnabled || policy.BitrateKbps > 125000 ||
		(policy.Codec != "H264" && policy.Codec != "HEVC") {
		return ErrProviderMalformed
	}
	if (policy.Codec == "H264" && (policy.ResolutionWidth > 4096 || policy.ResolutionHeight > 4096)) ||
		(policy.Codec == "HEVC" && (policy.ResolutionWidth > 8192 || policy.ResolutionHeight > 8192)) {
		return ErrProviderMalformed
	}
	return nil
}

func validateApolloProviderStreamPolicy(info ManagementInfo, policy protocol.ProviderStreamPolicy) error {
	if validateApolloStreamPolicy(policy) != nil || !info.HasServerCodecModeSupport || !info.HasMaxLumaPixelsHEVC {
		return ErrProviderMalformed
	}
	switch policy.Codec {
	case "H264":
		if info.ServerCodecModeSupport&0x1 == 0 {
			return ErrProviderMalformed
		}
	case "HEVC":
		luma := uint64(policy.ResolutionWidth) * uint64(policy.ResolutionHeight)
		if info.ServerCodecModeSupport&0x100 == 0 || info.MaxLumaPixelsHEVC == 0 || luma > info.MaxLumaPixelsHEVC {
			return ErrProviderMalformed
		}
	default:
		return ErrProviderMalformed
	}
	return nil
}

func validApolloRTSPToken(value string) bool {
	if value == "" || len(value) > 128 {
		return false
	}
	for _, character := range value {
		if character <= 0x20 || character >= 0x7f || strings.ContainsRune("()<>@,;:\\\"/[]?={}", character) {
			return false
		}
	}
	return true
}