package gateway

import (
	"bytes"
	"context"
	"crypto/rand"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/binary"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

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

// NativeApolloBackend keeps provider sockets inside the gateway process. The
// session-scoped Server work is the sole source of provider endpoint and mTLS
// material; it is never serialized into a client manifest or authority.
type NativeApolloBackend struct {
	Dialer *net.Dialer

	mu      sync.Mutex
	pending map[string]net.Conn
}

func NewNativeApolloBackend() *NativeApolloBackend {
	return &NativeApolloBackend{Dialer: &net.Dialer{Timeout: 5 * time.Second}, pending: make(map[string]net.Conn)}
}

func (b *NativeApolloBackend) Management(ctx context.Context, request LaunchRequest) ([]byte, error) {
	work := request.ProviderWork
	if err := work.Validate(); err != nil || request.SessionID == "" || request.SessionID != work.SessionID || work.ProviderProfile != ProviderProfileApollo {
		return nil, ErrProviderMalformed
	}
	client, err := newPinnedApolloHTTPClient(work)
	if err != nil {
		return nil, err
	}
	return apolloGet(ctx, client, work, "/serverinfo", nil)
}

func newPinnedApolloHTTPClient(work protocol.ProviderSessionWork) (*http.Client, error) {
	tlsConfig, err := pinnedApolloTLSConfig(work)
	if err != nil {
		return nil, err
	}
	return &http.Client{Transport: &http.Transport{TLSClientConfig: tlsConfig}, Timeout: 5 * time.Second}, nil
}

func apolloGet(ctx context.Context, client *http.Client, work protocol.ProviderSessionWork, path string, values url.Values) ([]byte, error) {
	endpoint := url.URL{Scheme: "https", Host: net.JoinHostPort(work.ManagementHost, strconv.FormatInt(work.ManagementPort, 10)), Path: path}
	endpoint.RawQuery = values.Encode()
	request, err := http.NewRequestWithContext(ctx, http.MethodGet, endpoint.String(), nil)
	if err != nil {
		return nil, ErrProviderMalformed
	}
	response, err := client.Do(request)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()
	if response.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("provider management status %d", response.StatusCode)
	}
	return readBounded(response.Body, 64*1024)
}

func pinnedApolloTLSConfig(work protocol.ProviderSessionWork) (*tls.Config, error) {
	identity, ok := providerIdentityFromKey(work.ProviderIdentity)
	if !ok || !strings.HasPrefix(identity.Fingerprint, "sha256:") {
		return nil, ErrProviderIdentity
	}
	pinned, err := hex.DecodeString(strings.TrimPrefix(identity.Fingerprint, "sha256:"))
	if err != nil || len(pinned) != sha256.Size {
		return nil, ErrProviderIdentity
	}
	certificate, err := tls.X509KeyPair([]byte(work.ClientCertificatePem), []byte(work.ClientPrivateKeyPem))
	if err != nil {
		return nil, ErrProviderIdentity
	}
	trust := x509.NewCertPool()
	if !trust.AppendCertsFromPEM([]byte(work.ServerCertificatePem)) {
		return nil, ErrProviderIdentity
	}
	return &tls.Config{
		MinVersion: tls.VersionTLS13, Certificates: []tls.Certificate{certificate}, RootCAs: trust,
		VerifyPeerCertificate: func(rawCertificates [][]byte, _ [][]*x509.Certificate) error {
			if len(rawCertificates) == 0 {
				return ErrProviderIdentity
			}
			digest := sha256.Sum256(rawCertificates[0])
			if !bytes.Equal(digest[:], pinned) {
				return ErrProviderIdentity
			}
			return nil
		},
	}, nil
}

func (b *NativeApolloBackend) Setup(ctx context.Context, request LaunchRequest) ([]byte, error) {
	work := request.ProviderWork
	if err := work.Validate(); err != nil || request.SessionID == "" || request.SessionID != work.SessionID || work.ProviderProfile != ProviderProfileApollo {
		return nil, ErrProviderMalformed
	}
	client, err := newPinnedApolloHTTPClient(work)
	if err != nil {
		return nil, err
	}
	inventory, err := apolloGet(ctx, client, work, "/applist", url.Values{"uniqueid": []string{work.ClientID}})
	if err != nil || !apolloInventoryContains(inventory, work.ApplicationID) {
		return nil, ErrProviderMalformed
	}
	key := make([]byte, 16)
	if _, err := rand.Read(key); err != nil {
		return nil, err
	}
	var keyID [4]byte
	if _, err := rand.Read(keyID[:]); err != nil {
		return nil, err
	}
	launch, err := apolloGet(ctx, client, work, "/launch", url.Values{
		"uniqueid": {work.ClientID}, "appid": {work.ApplicationID}, "rikey": {hex.EncodeToString(key)},
		"rikeyid": {strconv.FormatUint(uint64(binary.BigEndian.Uint32(keyID[:])), 10)}, "localAudioPlayMode": {"0"},
		"corever": {"1"},
	})
	if err != nil {
		return nil, err
	}
	launchResponse, err := parseApolloLaunchResponse(launch)
	if err != nil {
		return nil, err
	}
	streamURL, err := url.Parse(launchResponse.SessionURL)
	if err != nil || streamURL.Scheme != "rtspenc" || streamURL.Hostname() != work.StreamHost {
		return nil, ErrProviderMalformed
	}
	streamPort, err := strconv.ParseInt(streamURL.Port(), 10, 64)
	if err != nil || streamPort != work.StreamPort {
		return nil, ErrProviderMalformed
	}
	conn, err := b.Dialer.DialContext(ctx, "tcp", net.JoinHostPort(work.StreamHost, strconv.FormatInt(work.StreamPort, 10)))
	if err != nil {
		return nil, err
	}
	if deadline, ok := ctx.Deadline(); ok {
		_ = conn.SetDeadline(deadline)
	}
	codec, err := newEncryptedRTSPCodec(key)
	if err != nil {
		_ = conn.Close()
		return nil, err
	}
	requestText := "SETUP rtsp://" + work.StreamHost + "/streamid=video/0/0 RTSP/1.0\r\nCSeq: 1\r\nTransport: RTP/AVP/TCP;interleaved=0-1\r\n\r\n"
	encoded, err := codec.SealClient([]byte(requestText))
	if err != nil {
		_ = conn.Close()
		return nil, err
	}
	if _, err := conn.Write(encoded); err != nil {
		_ = conn.Close()
		return nil, err
	}
	response, err := readEncryptedRTSPHeaders(conn, codec)
	if err != nil {
		_ = conn.Close()
		return nil, err
	}
	b.mu.Lock()
	b.pending[request.SessionID] = conn
	b.mu.Unlock()
	return response, nil
}

func (b *NativeApolloBackend) Open(_ context.Context, request LaunchRequest, _ RTSPResponse) (ProviderSession, error) {
	b.mu.Lock()
	conn, ok := b.pending[request.SessionID]
	delete(b.pending, request.SessionID)
	b.mu.Unlock()
	if !ok || conn == nil {
		return nil, ErrProviderDisconnected
	}
	session := newNativeApolloSession(conn, request.SessionID)
	go session.readMedia()
	return session, nil
}

func readBounded(reader io.Reader, max int) ([]byte, error) {
	data, err := io.ReadAll(io.LimitReader(reader, int64(max)+1))
	if err != nil {
		return nil, err
	}
	if len(data) > max {
		return nil, ErrProviderMalformed
	}
	return data, nil
}

func readEncryptedRTSPHeaders(conn net.Conn, codec *encryptedRTSPCodec) ([]byte, error) {
	header := make([]byte, encryptedRTSPHeaderSize)
	if _, err := io.ReadFull(conn, header); err != nil {
		return nil, err
	}
	length := binary.BigEndian.Uint32(header[:4]) & 0x7fffffff
	if length == 0 || length > encryptedRTSPMaxPayload {
		return nil, ErrProviderMalformed
	}
	frame := make([]byte, encryptedRTSPHeaderSize+int(length))
	copy(frame, header)
	if _, err := io.ReadFull(conn, frame[encryptedRTSPHeaderSize:]); err != nil {
		return nil, err
	}
	plaintext, err := codec.OpenHost(frame)
	if err != nil || len(plaintext) > 16*1024 || !strings.HasSuffix(string(plaintext), "\r\n\r\n") {
		return nil, ErrProviderMalformed
	}
	return plaintext, nil
}

type nativeApolloSession struct {
	conn      net.Conn
	sessionID string
	video     chan []byte
	audio     chan []byte
	mu        sync.Mutex
	state     protocol.ProviderState
	closeOnce sync.Once
	done      chan struct{}
	readDone  chan struct{}
}

func newNativeApolloSession(conn net.Conn, sessionID string) *nativeApolloSession {
	return &nativeApolloSession{conn: conn, sessionID: sessionID, video: make(chan []byte, 16), audio: make(chan []byte, 16), state: protocol.ProviderState{Version: "1", SessionID: sessionID, State: ProviderStateStarting, Channels: []string{"video", "audio", "input", "feedback"}}, done: make(chan struct{}), readDone: make(chan struct{})}
}

func (s *nativeApolloSession) Ready(context.Context) error {
	s.mu.Lock()
	s.state.State = ProviderStateReady
	s.mu.Unlock()
	return nil
}

func (s *nativeApolloSession) Video() <-chan []byte { return s.video }
func (s *nativeApolloSession) Audio() <-chan []byte { return s.audio }

func (s *nativeApolloSession) Input(ctx context.Context, event InputEvent) error {
	payload, err := EncodeInputEvent(event)
	if err != nil {
		return err
	}
	return s.writeControl(ctx, ControlPacket{Kind: 1, Sequence: event.Sequence, Payload: payload})
}

func (s *nativeApolloSession) Feedback(ctx context.Context, feedback Feedback) error {
	return s.writeControl(ctx, ControlPacket{Kind: 3, Sequence: feedback.Sequence, Payload: feedback.Payload})
}

func (s *nativeApolloSession) Reconnect(ctx context.Context) error {
	return s.writeControl(ctx, ControlPacket{Kind: 4, Payload: []byte("RECN")})
}

func (s *nativeApolloSession) ReleaseAll(ctx context.Context) error {
	return s.writeControl(ctx, ControlPacket{Kind: 2, Payload: []byte("RELEASE_ALL")})
}

func (s *nativeApolloSession) Terminate(ctx context.Context) error {
	_ = s.writeControl(ctx, ControlPacket{Kind: 5, Payload: []byte("TEAR")})
	var timedOut bool
	s.closeOnce.Do(func() {
		close(s.done)
		_ = s.conn.Close()
		select {
		case <-s.readDone:
		case <-ctx.Done():
			timedOut = true
		}
		if !timedOut {
			close(s.video)
			close(s.audio)
		}
	})
	s.mu.Lock()
	if timedOut {
		s.state.State = ProviderStateCleanup
		s.state.CleanupPending = true
	} else {
		s.state.State = ProviderStateTerminated
	}
	s.mu.Unlock()
	if timedOut {
		return ctx.Err()
	}
	return nil
}

func (s *nativeApolloSession) State() protocol.ProviderState {
	s.mu.Lock()
	defer s.mu.Unlock()
	return s.state
}

func (s *nativeApolloSession) writeControl(ctx context.Context, packet ControlPacket) error {
	encoded, err := EncodeControlPacket(packet)
	if err != nil {
		return err
	}
	if deadline, ok := ctx.Deadline(); ok {
		_ = s.conn.SetWriteDeadline(deadline)
	}
	if _, err := s.conn.Write(encoded); err != nil {
		return err
	}
	return nil
}

func (s *nativeApolloSession) readMedia() {
	defer close(s.readDone)
	header := make([]byte, 4)
	for {
		if _, err := io.ReadFull(s.conn, header); err != nil {
			return
		}
		if header[0] != '$' || (header[1] != 0 && header[1] != 1) {
			return
		}
		length := int(header[2])<<8 | int(header[3])
		if length > 65536 {
			return
		}
		payload := make([]byte, length)
		if _, err := io.ReadFull(s.conn, payload); err != nil {
			return
		}
		if header[1] == 0 {
			pushLatest(s.video, payload)
		} else {
			pushLatest(s.audio, payload)
		}
	}
}

func pushLatest(channel chan []byte, payload []byte) {
	select {
	case channel <- payload:
	default:
		select {
		case <-channel:
		default:
		}
		select {
		case channel <- payload:
		default:
		}
	}
}

var _ ApolloBackend = (*NativeApolloBackend)(nil)
var _ ProviderSession = (*nativeApolloSession)(nil)