package gateway

import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/quic-go/quic-go"

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

const (
	defaultHelloLimit   = 16 * 1024
	defaultControlLimit = 128 * 1024
	applicationError    = quic.ApplicationErrorCode(0x100)
)

var (
	ErrAdmissionRejected = errors.New("gateway admission rejected")
	ErrGatewayDraining   = errors.New("gateway draining")
	ErrGatewayTLS        = errors.New("gateway requires TLS 1.3 client authentication")
	ErrAuthorityExpired  = errors.New("gateway authority expired")
)

type Admission interface {
	Admit(context.Context, protocol.TunnelAdmissionRequest) (protocol.SessionAuthority, error)
	ProviderWork(context.Context, protocol.SessionAuthority) (protocol.ProviderSessionWork, error)
	Release(context.Context, protocol.SessionAuthority) error
}

type AdmissionFunc func(context.Context, protocol.TunnelAdmissionRequest) (protocol.SessionAuthority, error)

func (f AdmissionFunc) Admit(ctx context.Context, request protocol.TunnelAdmissionRequest) (protocol.SessionAuthority, error) {
	return f(ctx, request)
}

func (AdmissionFunc) Release(context.Context, protocol.SessionAuthority) error { return nil }

func (AdmissionFunc) ProviderWork(context.Context, protocol.SessionAuthority) (protocol.ProviderSessionWork, error) {
	return protocol.ProviderSessionWork{}, ErrAdmissionRejected
}

type ProviderStateReporter interface {
	ReportProviderState(context.Context, protocol.ProviderState) error
}

type ServerConfig struct {
	ListenAddress         string
	TLSConfig             *tls.Config
	QUICConfig            *quic.Config
	GatewayID             string
	Capabilities          protocol.CapabilityProfile
	ProviderCapabilities  protocol.CapabilityProfile
	Admission             Admission
	ProviderStateReporter ProviderStateReporter
	Provider              Provider
	ProviderProfile       string
	ProviderIdentity      string
	PacerKbps             int64
}

type Server struct {
	listener  *quic.Listener
	config    ServerConfig
	metrics   *Metrics
	mu        sync.Mutex
	sessions  map[*gatewaySession]struct{}
	draining  atomic.Bool
	closed    atomic.Bool
	closeOnce sync.Once
	workers   sync.WaitGroup
}

func NewServer(config ServerConfig) (*Server, error) {
	if config.ListenAddress == "" {
		config.ListenAddress = "127.0.0.1:0"
	}
	if config.GatewayID == "" || config.Admission == nil || config.Provider == nil {
		return nil, errors.New("gateway id, admission, and provider are required")
	}
	if err := validateServerTLS(config.TLSConfig); err != nil {
		return nil, err
	}
	if config.Capabilities == (protocol.CapabilityProfile{}) {
		config.Capabilities = DefaultCapabilities()
	}
	if config.ProviderCapabilities == (protocol.CapabilityProfile{}) {
		config.ProviderCapabilities = DefaultCapabilities()
	}
	if config.ProviderProfile == "" {
		config.ProviderProfile = ProviderProfileApollo
	}
	if config.PacerKbps < 0 {
		return nil, errors.New("negative pacing limit")
	}
	tlsConfig := config.TLSConfig.Clone()
	if len(tlsConfig.NextProtos) == 0 {
		tlsConfig.NextProtos = []string{"versevdi-gateway-v1"}
	}
	quicConfig := &quic.Config{EnableDatagrams: true, MaxIdleTimeout: 30 * time.Second, MaxIncomingStreams: 2, MaxIncomingUniStreams: 2}
	if config.QUICConfig != nil {
		quicConfig = config.QUICConfig.Clone()
		quicConfig.EnableDatagrams = true
	}
	listener, err := quic.ListenAddr(config.ListenAddress, tlsConfig, quicConfig)
	if err != nil {
		return nil, err
	}
	return &Server{listener: listener, config: config, metrics: &Metrics{}, sessions: make(map[*gatewaySession]struct{})}, nil
}

func validateServerTLS(config *tls.Config) error {
	if config == nil || config.MinVersion < tls.VersionTLS13 || config.ClientAuth != tls.RequireAndVerifyClientCert || config.ClientCAs == nil || len(config.Certificates) == 0 {
		return ErrGatewayTLS
	}
	return nil
}

func (s *Server) Addr() net.Addr           { return s.listener.Addr() }
func (s *Server) Metrics() MetricsSnapshot { return s.metrics.Snapshot() }
func (s *Server) Draining() bool           { return s.draining.Load() }

func (s *Server) BeginDrain() {
	s.draining.Store(true)
}

func (s *Server) Serve(ctx context.Context) error {
	if s.closed.Load() {
		return net.ErrClosed
	}
	go func() {
		<-ctx.Done()
		_ = s.Close()
	}()
	for {
		connection, err := s.listener.Accept(ctx)
		if err != nil {
			if s.closed.Load() || errors.Is(err, context.Canceled) || errors.Is(err, net.ErrClosed) {
				return nil
			}
			return err
		}
		s.workers.Add(1)
		go func() {
			defer s.workers.Done()
			s.handleConnection(ctx, connection)
		}()
	}
}

func (s *Server) Close() error {
	var err error
	s.closeOnce.Do(func() {
		s.BeginDrain()
		s.closed.Store(true)
		err = s.listener.Close()
		s.mu.Lock()
		for session := range s.sessions {
			session.cancel()
		}
		s.mu.Unlock()
	})
	s.workers.Wait()
	return err
}

func (s *Server) handleConnection(parent context.Context, connection *quic.Conn) {
	defer connection.CloseWithError(applicationError, "connection closed")
	ctx, cancel := context.WithTimeout(parent, 10*time.Second)
	defer cancel()
	stream, err := connection.AcceptStream(ctx)
	if err != nil {
		return
	}
	requestBytes, err := readWire(stream, defaultHelloLimit)
	if err != nil {
		_ = writeStableError(stream, "invalid_hello", err, false)
		return
	}
	request, err := protocol.DecodeTunnelAdmissionRequest(requestBytes)
	if err != nil {
		_ = writeStableError(stream, "invalid_hello", err, false)
		return
	}
	if s.Draining() {
		_ = writeStableError(stream, "gateway_draining", ErrGatewayDraining, true)
		return
	}
	if request.GatewayID != s.config.GatewayID {
		_ = writeStableError(stream, "wrong_gateway", ErrAdmissionRejected, false)
		return
	}
	authority, err := s.config.Admission.Admit(ctx, request)
	if err != nil {
		s.metrics.AdmissionRejects.Add(1)
		_ = writeStableError(stream, stableAdmissionCode(err), err, errors.Is(err, context.DeadlineExceeded))
		return
	}
	if s.Draining() {
		_ = s.config.Admission.Release(context.Background(), authority)
		_ = writeStableError(stream, "gateway_draining", ErrGatewayDraining, true)
		return
	}
	if err := s.validateAuthority(authority, request); err != nil {
		_ = s.config.Admission.Release(context.Background(), authority)
		_ = writeStableError(stream, "invalid_authority", err, false)
		return
	}
	work, err := s.config.Admission.ProviderWork(ctx, authority)
	if err != nil || s.validateProviderWork(work, authority) != nil {
		_ = s.config.Admission.Release(context.Background(), authority)
		_ = writeStableError(stream, "provider_work_unavailable", ErrAdmissionRejected, err != nil)
		return
	}
	selected, err := IntersectCapabilities(s.config.Capabilities, s.config.ProviderCapabilities, request.Capabilities, authority.Capabilities)
	if err != nil {
		_ = s.config.Admission.Release(context.Background(), authority)
		s.metrics.AdmissionRejects.Add(1)
		_ = writeStableError(stream, "no_capability_overlap", err, false)
		return
	}
	if err := s.reportProviderState(ctx, protocol.ProviderState{Version: "1", SessionID: request.SessionID, State: ProviderStateStarting, CleanupPending: false, Channels: []string{"video", "audio", "input", "feedback"}}); err != nil {
		_ = s.config.Admission.Release(context.Background(), authority)
		_ = writeStableError(stream, "provider_state_unavailable", err, true)
		return
	}
	providerSession, err := s.config.Provider.Start(ctx, LaunchRequest{SessionID: request.SessionID, Capabilities: selected, ProviderProfile: authority.ProviderProfile, ProviderIdentity: work.ProviderIdentity, ProviderWork: work})
	if err != nil {
		s.metrics.ProviderErrors.Add(1)
		_ = s.reportProviderState(context.Background(), protocol.ProviderState{Version: "1", SessionID: request.SessionID, State: ProviderStateFailed, CleanupPending: false, Channels: []string{"video", "audio", "input", "feedback"}})
		_ = s.config.Admission.Release(context.Background(), authority)
		_ = writeStableError(stream, stableProviderCode(err), err, errors.Is(err, context.DeadlineExceeded))
		return
	}
	if err := s.reportProviderState(ctx, providerSession.State()); err != nil {
		_ = providerSession.ReleaseAll(context.Background())
		_ = providerSession.Terminate(context.Background())
		_ = s.config.Admission.Release(context.Background(), authority)
		_ = writeStableError(stream, "provider_state_unavailable", err, true)
		return
	}
	authority.Capabilities = selected
	authorityBytes, err := protocol.EncodeSessionAuthority(authority)
	if err != nil || writeWire(stream, authorityBytes, defaultHelloLimit) != nil {
		_ = providerSession.ReleaseAll(context.Background())
		_ = providerSession.Terminate(context.Background())
		_ = s.config.Admission.Release(context.Background(), authority)
		return
	}
	session := newGatewaySession(s, connection, stream, request, authority, providerSession)
	s.addSession(session)
	s.metrics.ActiveSessions.Add(1)
	defer func() {
		s.removeSession(session)
		s.metrics.ActiveSessions.Add(-1)
	}()
	session.run()
}

func (s *Server) reportProviderState(ctx context.Context, state protocol.ProviderState) error {
	if s.config.ProviderStateReporter == nil {
		return nil
	}
	if err := state.Validate(); err != nil {
		return err
	}
	return s.config.ProviderStateReporter.ReportProviderState(ctx, state)
}

func (s *Server) validateAuthority(authority protocol.SessionAuthority, request protocol.TunnelAdmissionRequest) error {
	if err := authority.Validate(); err != nil {
		return err
	}
	if authority.SessionID != request.SessionID || authority.GatewayID != request.GatewayID || authority.Audience != request.Audience || authority.ProviderProfile != s.config.ProviderProfile {
		return ErrAdmissionRejected
	}
	expires, err := time.Parse(time.RFC3339Nano, authority.ExpiresAt)
	if err != nil || !time.Now().Before(expires) {
		return ErrAuthorityExpired
	}
	return nil
}

func (s *Server) validateProviderWork(work protocol.ProviderSessionWork, authority protocol.SessionAuthority) error {
	if err := work.Validate(); err != nil {
		return err
	}
	if work.SessionID != authority.SessionID || work.GatewayID != authority.GatewayID ||
		work.ReconnectSequence != authority.ReconnectSequence || work.ExpiresAt != authority.ExpiresAt ||
		work.ProviderProfile != authority.ProviderProfile {
		return ErrAdmissionRejected
	}
	return nil
}

func (s *Server) addSession(session *gatewaySession) {
	s.mu.Lock()
	s.sessions[session] = struct{}{}
	s.mu.Unlock()
}

func (s *Server) removeSession(session *gatewaySession) {
	s.mu.Lock()
	delete(s.sessions, session)
	s.mu.Unlock()
}

type gatewaySession struct {
	server      *Server
	connection  *quic.Conn
	control     *quic.Stream
	request     protocol.TunnelAdmissionRequest
	authority   protocol.SessionAuthority
	provider    ProviderSession
	pacer       *Pacer
	ctx         context.Context
	cancel      context.CancelFunc
	cleanupOnce sync.Once
	inputMu     sync.Mutex
	pressed     map[string]struct{}
	sequence    atomic.Uint32
	result      chan error
}

func newGatewaySession(server *Server, connection *quic.Conn, control *quic.Stream, request protocol.TunnelAdmissionRequest, authority protocol.SessionAuthority, provider ProviderSession) *gatewaySession {
	ctx, cancel := context.WithCancel(context.Background())
	return &gatewaySession{server: server, connection: connection, control: control, request: request, authority: authority, provider: provider, pacer: NewPacer(server.config.PacerKbps), ctx: ctx, cancel: cancel, pressed: make(map[string]struct{}), result: make(chan error, 3)}
}

func (s *gatewaySession) run() {
	defer s.cleanup()
	deadline, err := time.Parse(time.RFC3339Nano, s.authority.ExpiresAt)
	if err != nil {
		return
	}
	if deadline.Before(time.Now()) {
		return
	}
	timer := time.NewTimer(time.Until(deadline))
	defer timer.Stop()
	go s.controlLoop()
	go s.datagramLoop()
	go s.mediaLoop()
	select {
	case <-timer.C:
		s.server.metrics.InputRejected.Add(1)
	case <-s.ctx.Done():
	case <-s.result:
	}
	s.cancel()
}

func (s *gatewaySession) controlLoop() {
	for {
		data, err := readWire(s.control, defaultControlLimit)
		if err != nil {
			s.result <- err
			return
		}
		frame, err := protocol.DecodeChannelFrame(data)
		if err != nil {
			s.result <- err
			return
		}
		payload, err := base64.StdEncoding.DecodeString(frame.Payload)
		if err != nil || len(payload) > maxFrameSize {
			s.result <- ErrFramePayloadLimit
			return
		}
		switch frame.FlowID {
		case "control":
			if err := s.handleControl(payload); err != nil {
				s.result <- err
				return
			}
		case "input":
			if err := s.handleInput(payload); err != nil {
				s.result <- err
				return
			}
		default:
			s.result <- ErrFrameChannel
			return
		}
	}
}

func (s *gatewaySession) datagramLoop() {
	for {
		data, err := s.connection.ReceiveDatagram(s.ctx)
		if err != nil {
			s.result <- err
			return
		}
		frame, err := DecodeFrame(data)
		if err != nil {
			s.result <- err
			return
		}
		switch frame.Channel {
		case ChannelInput:
			if err := s.handleInput(frame.Payload); err != nil {
				s.result <- err
				return
			}
		case ChannelText:
			if len(frame.Payload) > 4096 {
				s.result <- ErrFramePayloadLimit
				return
			}
			if err := s.provider.Feedback(s.ctx, Feedback{Sequence: frame.Sequence, Payload: append([]byte(nil), frame.Payload...)}); err != nil {
				s.result <- err
				return
			}
		default:
			s.result <- ErrFrameChannel
			return
		}
	}
}

func (s *gatewaySession) mediaLoop() {
	video, audio := s.provider.Video(), s.provider.Audio()
	for video != nil || audio != nil {
		select {
		case <-s.ctx.Done():
			return
		case payload, ok := <-video:
			if !ok {
				video = nil
				continue
			}
			if err := s.sendMedia(ChannelVideo, payload); err != nil {
				s.result <- err
				return
			}
		case payload, ok := <-audio:
			if !ok {
				audio = nil
				continue
			}
			if err := s.sendMedia(ChannelAudio, payload); err != nil {
				s.result <- err
				return
			}
		}
	}
	s.result <- ErrProviderDisconnected
}

func (s *gatewaySession) sendMedia(channel byte, payload []byte) error {
	frames, err := FragmentPayload(channel, s.sequence.Add(1), uint64(time.Now().UnixMilli()), payload)
	if err != nil {
		return err
	}
	for _, frame := range frames {
		encoded, err := EncodeFrame(frame)
		if err != nil {
			return err
		}
		if err := s.pacer.Wait(s.ctx, len(encoded)); err != nil {
			return err
		}
		if err := s.connection.SendDatagram(encoded); err != nil {
			return err
		}
	}
	return nil
}

func (s *gatewaySession) handleControl(payload []byte) error {
	if len(payload) < 4 {
		return ErrProviderMalformed
	}
	switch string(payload[:4]) {
	case "TERM":
		return errors.New("client requested termination")
	case "RECN":
		return s.provider.Reconnect(s.ctx)
	case "FBRK":
		return s.provider.Feedback(s.ctx, Feedback{Payload: append([]byte(nil), payload[4:]...)})
	default:
		return ErrProviderMalformed
	}
}

func (s *gatewaySession) handleInput(payload []byte) error {
	event, err := DecodeInputEvent(payload)
	if err != nil {
		return err
	}
	if err := s.provider.Input(s.ctx, event); err != nil {
		s.server.metrics.InputRejected.Add(1)
		return err
	}
	key := fmt.Sprintf("%s:%d", event.Device, event.Code)
	s.inputMu.Lock()
	if event.Pressed {
		s.pressed[key] = struct{}{}
	} else {
		delete(s.pressed, key)
	}
	s.inputMu.Unlock()
	return nil
}

func (s *gatewaySession) cleanup() {
	s.cleanupOnce.Do(func() {
		s.cancel()
		cleanupCtx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		releaseInputsErr := s.provider.ReleaseAll(cleanupCtx)
		if releaseInputsErr != nil {
			s.server.metrics.ProviderErrors.Add(1)
		}
		terminateErr := s.provider.Terminate(cleanupCtx)
		if terminateErr != nil {
			s.server.metrics.ProviderErrors.Add(1)
		}
		state := s.provider.State()
		if releaseInputsErr != nil || terminateErr != nil {
			state.State = ProviderStateCleanup
			state.CleanupPending = true
			if err := s.server.reportProviderState(cleanupCtx, state); err != nil {
				s.server.metrics.ProviderErrors.Add(1)
			} else if err := s.server.config.Admission.Release(cleanupCtx, s.authority); err != nil {
				s.server.metrics.ProviderErrors.Add(1)
			}
			_ = s.connection.CloseWithError(applicationError, "session closed")
			return
		}
		if err := s.server.config.Admission.Release(cleanupCtx, s.authority); err != nil {
			s.server.metrics.ProviderErrors.Add(1)
		}
		if err := s.server.reportProviderState(cleanupCtx, state); err != nil {
			s.server.metrics.ProviderErrors.Add(1)
		}
		_ = s.connection.CloseWithError(applicationError, "session closed")
	})
}

func writeStableError(writer io.Writer, code string, err error, retryable bool) error {
	message := err.Error()
	if len(message) > 256 {
		message = message[:256]
	}
	payload, encodeErr := protocol.EncodeStableError(protocol.StableError{Version: "1", Code: code, Message: message, Retryable: retryable})
	if encodeErr != nil {
		return encodeErr
	}
	return writeWire(writer, payload, defaultHelloLimit)
}

func stableAdmissionCode(err error) string {
	if errors.Is(err, ErrGatewayDraining) {
		return "gateway_draining"
	}
	if errors.Is(err, ErrAuthorityExpired) {
		return "expired_grant"
	}
	return "admission_rejected"
}

func stableProviderCode(err error) string {
	switch {
	case errors.Is(err, ErrProviderIdentity):
		return "provider_identity_rejected"
	case errors.Is(err, ErrProviderMalformed):
		return "provider_malformed"
	case errors.Is(err, ErrProviderTimeout), errors.Is(err, context.DeadlineExceeded):
		return "provider_timeout"
	default:
		return "provider_unavailable"
	}
}

func writeWire(writer io.Writer, payload []byte, max int) error {
	if len(payload) > max {
		return ErrFrameSize
	}
	header := [4]byte{}
	binary.BigEndian.PutUint32(header[:], uint32(len(payload)))
	if _, err := writer.Write(header[:]); err != nil {
		return err
	}
	_, err := writer.Write(payload)
	return err
}

func readWire(reader io.Reader, max int) ([]byte, error) {
	var header [4]byte
	if _, err := io.ReadFull(reader, header[:]); err != nil {
		return nil, err
	}
	length := binary.BigEndian.Uint32(header[:])
	if length > uint32(max) {
		return nil, ErrFrameSize
	}
	payload := make([]byte, int(length))
	if _, err := io.ReadFull(reader, payload); err != nil {
		return nil, err
	}
	return payload, nil
}

type Client struct {
	connection *quic.Conn
	control    *quic.Stream
	Authority  protocol.SessionAuthority
}

func Dial(ctx context.Context, address string, tlsConfig *tls.Config, request protocol.TunnelAdmissionRequest) (*Client, error) {
	if tlsConfig == nil || tlsConfig.MinVersion < tls.VersionTLS13 || tlsConfig.RootCAs == nil || len(tlsConfig.RootCAs.Subjects()) == 0 || len(tlsConfig.Certificates) == 0 {
		return nil, ErrGatewayTLS
	}
	tlsConfig = tlsConfig.Clone()
	if len(tlsConfig.NextProtos) == 0 {
		tlsConfig.NextProtos = []string{"versevdi-gateway-v1"}
	}
	quicConfig := &quic.Config{EnableDatagrams: true, MaxIdleTimeout: 30 * time.Second}
	connection, err := quic.DialAddr(ctx, address, tlsConfig, quicConfig)
	if err != nil {
		return nil, err
	}
	stream, err := connection.OpenStreamSync(ctx)
	if err != nil {
		_ = connection.CloseWithError(applicationError, "stream unavailable")
		return nil, err
	}
	payload, err := protocol.EncodeTunnelAdmissionRequest(request)
	if err != nil {
		_ = connection.CloseWithError(applicationError, "invalid hello")
		return nil, err
	}
	if err := writeWire(stream, payload, defaultHelloLimit); err != nil {
		_ = connection.CloseWithError(applicationError, "invalid hello")
		return nil, err
	}
	response, err := readWire(stream, defaultHelloLimit)
	if err != nil {
		_ = connection.CloseWithError(applicationError, "no authority")
		return nil, err
	}
	authority, authorityErr := protocol.DecodeSessionAuthority(response)
	if authorityErr != nil {
		stable, stableErr := protocol.DecodeStableError(response)
		if stableErr == nil {
			return nil, fmt.Errorf("%s: %s", stable.Code, stable.Message)
		}
		return nil, authorityErr
	}
	return &Client{connection: connection, control: stream, Authority: authority}, nil
}

func (c *Client) SendInput(event InputEvent) error {
	payload, err := EncodeInputEvent(event)
	if err != nil {
		return err
	}
	frame := protocol.ChannelFrame{Version: "1", FlowID: "input", Sequence: int64(event.Sequence), Flags: 0, FragmentIndex: 0, FragmentCount: 1, TimestampMs: time.Now().UnixMilli(), Payload: base64.StdEncoding.EncodeToString(payload)}
	encoded, err := protocol.EncodeChannelFrame(frame)
	if err != nil {
		return err
	}
	return writeWire(c.control, encoded, defaultControlLimit)
}

func (c *Client) SendControl(payload []byte) error {
	frame := protocol.ChannelFrame{Version: "1", FlowID: "control", Sequence: 0, Flags: 0, FragmentIndex: 0, FragmentCount: 1, TimestampMs: time.Now().UnixMilli(), Payload: base64.StdEncoding.EncodeToString(payload)}
	encoded, err := protocol.EncodeChannelFrame(frame)
	if err != nil {
		return err
	}
	return writeWire(c.control, encoded, defaultControlLimit)
}

func (c *Client) ReceiveFrame(ctx context.Context) (Frame, error) {
	data, err := c.connection.ReceiveDatagram(ctx)
	if err != nil {
		return Frame{}, err
	}
	return DecodeFrame(data)
}

func (c *Client) Close() error {
	return c.connection.CloseWithError(applicationError, "client closed")
}