package gateway

import (
	"bufio"
	"bytes"
	"context"
	"crypto/sha256"
	"crypto/tls"
	"crypto/x509"
	"encoding/hex"
	"fmt"
	"io"
	"net"
	"net/http"
	"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 || work.ProviderProfile != ProviderProfileApollo {
		return nil, ErrProviderMalformed
	}
	tlsConfig, err := pinnedApolloTLSConfig(work)
	if err != nil {
		return nil, err
	}
	client := &http.Client{Transport: &http.Transport{TLSClientConfig: tlsConfig}, Timeout: 5 * time.Second}
	managementURL := "https://" + net.JoinHostPort(work.ManagementHost, strconv.FormatInt(work.ManagementPort, 10)) + "/serverinfo"
	httpRequest, err := http.NewRequestWithContext(ctx, http.MethodGet, managementURL, nil)
	if err != nil {
		return nil, ErrProviderMalformed
	}
	response, err := client.Do(httpRequest)
	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 == "" {
		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)
	}
	requestText := "SETUP rtsp://" + work.StreamHost + "/streamid=video/0/0 RTSP/1.0\r\nCSeq: 1\r\nTransport: RTP/AVP/TCP;interleaved=0-1\r\nSession: " + request.SessionID + "\r\n\r\n"
	if _, err := io.WriteString(conn, requestText); err != nil {
		_ = conn.Close()
		return nil, err
	}
	response, err := readRTSPHeaders(conn, 16*1024)
	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 := b.pending[request.SessionID]
	delete(b.pending, request.SessionID)
	b.mu.Unlock()
	if 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 readRTSPHeaders(conn net.Conn, max int) ([]byte, error) {
	reader := bufio.NewReaderSize(conn, 4096)
	var response []byte
	for len(response) < max {
		line, err := reader.ReadBytes('\n')
		if err != nil {
			return nil, err
		}
		response = append(response, line...)
		if strings.HasSuffix(string(response), "\r\n\r\n") {
			return response, nil
		}
	}
	return nil, ErrProviderMalformed
}

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)