package http3

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"fmt"
	"io"
	"math/big"
	"net"
	"time"

	"grpccanary/lib/grpc/http3/protoapi"

	"github.com/quic-go/quic-go"
	"google.golang.org/grpc"
	"google.golang.org/grpc/reflection"
)

// 1. self-signed certificate generation
func generateTLSConfig() (*tls.Config, error) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, err
	}

	template := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			Organization: []string{"gRPC HTTP3 Canary"},
		},
		NotBefore: time.Now(),
		NotAfter:  time.Now().Add(365 * 24 * time.Hour),
		KeyUsage:  x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		DNSNames:  []string{"localhost"},
		IPAddresses: []net.IP{net.ParseIP("127.0.0.1"), net.ParseIP("::1")},
	}

	certDER, err := x509.CreateCertificate(rand.Reader, &template, &template, &key.PublicKey, key)
	if err != nil {
		return nil, err
	}

	certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})

	privBytes, err := x509.MarshalPKCS8PrivateKey(key)
	if err != nil {
		return nil, err
	}
	keyPEM := pem.EncodeToMemory(&pem.Block{Type: "PRIVATE KEY", Bytes: privBytes})

	tlsCert, err := tls.X509KeyPair(certPEM, keyPEM)
	if err != nil {
		return nil, err
	}

	return &tls.Config{
		Certificates: []tls.Certificate{tlsCert},
		NextProtos:   []string{"grpc-http3-canary"},
	}, nil
}

// 2. net.Conn wrapper for quic.Stream
type quicNetConn struct {
	*quic.Stream
	conn *quic.Conn
}

func (c *quicNetConn) LocalAddr() net.Addr {
	return c.conn.LocalAddr()
}

func (c *quicNetConn) RemoteAddr() net.Addr {
	return c.conn.RemoteAddr()
}

// 3. net.Listener wrapper for quic.Listener using non-blocking channels
type quicListener struct {
	lis      *quic.Listener
	connChan chan net.Conn
	errChan  chan error
	ctx      context.Context
	cancel   context.CancelFunc
}

func NewQuicListener(lis *quic.Listener) *quicListener {
	ctx, cancel := context.WithCancel(context.Background())
	ql := &quicListener{
		lis:      lis,
		connChan: make(chan net.Conn, 100),
		errChan:  make(chan error, 10),
		ctx:      ctx,
		cancel:   cancel,
	}
	go ql.listenLoop()
	return ql
}

func (ql *quicListener) listenLoop() {
	for {
		qconn, err := ql.lis.Accept(ql.ctx)
		if err != nil {
			select {
			case ql.errChan <- err:
			default:
			}
			return
		}
		go ql.acceptStreams(qconn)
	}
}

func (ql *quicListener) acceptStreams(qconn *quic.Conn) {
	for {
		stream, err := qconn.AcceptStream(ql.ctx)
		if err != nil {
			// Stop checking this connection when it closes
			return
		}
		ql.connChan <- &quicNetConn{Stream: stream, conn: qconn}
	}
}

func (ql *quicListener) Accept() (net.Conn, error) {
	select {
	case conn := <-ql.connChan:
		return conn, nil
	case err := <-ql.errChan:
		return nil, err
	case <-ql.ctx.Done():
		return nil, io.EOF
	}
}

func (ql *quicListener) Close() error {
	ql.cancel()
	return ql.lis.Close()
}

func (ql *quicListener) Addr() net.Addr {
	return ql.lis.Addr()
}

// 4. Http3Service Server Implementation
type Http3Server struct {
	protoapi.UnimplementedHttp3ServiceServer
}

func (Http3Server) Ping(ctx context.Context, r *protoapi.PingRequest) (*protoapi.PingResponse, error) {
	return &protoapi.PingResponse{
		Message:   "Pong: " + r.GetMessage(),
		Transport: "quic",
	}, nil
}

func (Http3Server) StreamPing(stream protoapi.Http3Service_StreamPingServer) error {
	for {
		req, err := stream.Recv()
		if err == io.EOF {
			return nil
		}
		if err != nil {
			return err
		}
		err = stream.Send(&protoapi.PingResponse{
			Message:   "Pong Stream: " + req.GetMessage(),
			Transport: "quic",
		})
		if err != nil {
			return err
		}
	}
}

func ServerRun(addr string) (*quic.Listener, func(), error) {
	tlsConf, err := generateTLSConfig()
	if err != nil {
		return nil, nil, err
	}

	lis, err := quic.ListenAddr(addr, tlsConf, &quic.Config{
		KeepAlivePeriod: 10 * time.Second,
	})
	if err != nil {
		return nil, nil, err
	}

	qlis := NewQuicListener(lis)
	server := grpc.NewServer()
	protoapi.RegisterHttp3ServiceServer(server, Http3Server{})

	reflection.Register(server)

	go func() {
		fmt.Printf("[HTTP3 Server] Serving gRPC on UDP/QUIC %s...\n", addr)
		if err := server.Serve(qlis); err != nil {
			fmt.Println("[HTTP3 Server] Server closed:", err)
		}
	}()

	cleanup := func() {
		server.GracefulStop()
		qlis.Close()
	}

	return lis, cleanup, nil
}