package controlplane

import (
	"context"
	"crypto/sha256"
	"errors"
	"fmt"
	"io"
	"math"
	"time"

	controlplanev1 "proxy-pool/gen/controlplane/v1"
	"proxy-pool/internal/controlplane/snapshotwire"
	proxyDomain "proxy-pool/internal/domain/proxy"
	routingDomain "proxy-pool/internal/domain/routing"
	"proxy-pool/internal/domain/workerruntime"
	"proxy-pool/internal/gateway/snapshot"

	"google.golang.org/grpc"
	"google.golang.org/protobuf/types/known/emptypb"
	"google.golang.org/protobuf/types/known/timestamppb"
)

var (
	ErrInvalidSnapshotWatcher = errors.New("invalid gateway snapshot watcher")
	ErrSnapshotChecksum       = errors.New("worker snapshot checksum mismatch")
	ErrSnapshotDelta          = errors.New("worker snapshot delta requires a full snapshot")
)

type SnapshotStream interface {
	Recv() (*controlplanev1.SnapshotEnvelope, error)
}

type SnapshotRPCClient interface {
	Watch(context.Context, *controlplanev1.WatchSnapshotsRequest) (SnapshotStream, error)
	Acknowledge(context.Context, *controlplanev1.AcknowledgeSnapshotRequest) (*emptypb.Empty, error)
}

type SnapshotWatcherOptions struct {
	ClusterID string
	WorkerID  string
}

type SnapshotWatcher struct {
	client  SnapshotRPCClient
	store   *snapshot.Store
	options SnapshotWatcherOptions
}

func NewSnapshotWatcher(client SnapshotRPCClient, store *snapshot.Store, options SnapshotWatcherOptions) (*SnapshotWatcher, error) {
	if client == nil || store == nil || !workerruntime.ValidIdentifier(options.ClusterID) || !workerruntime.ValidIdentifier(options.WorkerID) {
		return nil, ErrInvalidSnapshotWatcher
	}
	return &SnapshotWatcher{client: client, store: store, options: options}, nil
}

// Watch applies verified full snapshots in stream order. Delta support remains
// intentionally fail-closed until the Gateway can validate a complete routing view.
func (watcher *SnapshotWatcher) Watch(ctx context.Context, sessionID string) error {
	if watcher == nil || ctx == nil || !workerruntime.ValidIdentifier(sessionID) {
		return ErrInvalidSnapshotWatcher
	}
	request := &controlplanev1.WatchSnapshotsRequest{WorkerId: watcher.options.WorkerID, SessionId: sessionID}
	if current := watcher.store.Current(); current != nil {
		request.LastAppliedVersion = current.Version
	}
	stream, err := watcher.client.Watch(ctx, request)
	if err != nil {
		return err
	}
	for {
		envelope, err := stream.Recv()
		if errors.Is(err, io.EOF) {
			return nil
		}
		if err != nil {
			return err
		}
		if err := watcher.applyAndAcknowledge(ctx, sessionID, envelope); err != nil {
			return err
		}
	}
}

func (watcher *SnapshotWatcher) applyAndAcknowledge(ctx context.Context, sessionID string, envelope *controlplanev1.SnapshotEnvelope) error {
	if envelope == nil || envelope.GetFull() == nil {
		return ErrSnapshotDelta
	}
	full := envelope.GetFull()
	applyErr := watcher.applyFull(full)
	if len(full.GetChecksum()) != sha256.Size {
		return errors.Join(applyErr, ErrSnapshotChecksum)
	}
	acknowledgement := &controlplanev1.AcknowledgeSnapshotRequest{
		WorkerId: watcher.options.WorkerID, SessionId: sessionID, Version: full.GetVersion(),
		OwnershipEpoch: full.GetOwnershipEpoch(), Checksum: append([]byte(nil), full.GetChecksum()...),
		Applied: applyErr == nil,
	}
	if applyErr != nil {
		acknowledgement.ErrorCode = "snapshot_apply_failed"
	}
	_, acknowledgeErr := watcher.client.Acknowledge(ctx, acknowledgement)
	if applyErr != nil || acknowledgeErr != nil {
		return errors.Join(applyErr, acknowledgeErr)
	}
	return nil
}

func (watcher *SnapshotWatcher) applyFull(full *controlplanev1.WorkerSnapshot) error {
	if full == nil || full.GetVersion() == 0 || full.GetOwnershipEpoch() == 0 {
		return ErrInvalidSnapshotWatcher
	}
	validUntil, err := requiredFutureTimestamp(full.GetValidUntil(), time.Now().UTC())
	if err != nil {
		return err
	}
	checksum, err := snapshotwire.Checksum(full)
	if err != nil {
		return err
	}
	if len(full.GetChecksum()) != sha256.Size || string(checksum[:]) != string(full.GetChecksum()) {
		return ErrSnapshotChecksum
	}
	proxies, err := wireProxies(full.GetProxies())
	if err != nil {
		return err
	}
	routes, err := wireRouting(full.GetRouting())
	if err != nil {
		return err
	}
	envelope := snapshot.Envelope{
		ClusterID: watcher.options.ClusterID, WorkerID: watcher.options.WorkerID,
		Epoch: full.GetOwnershipEpoch(), Version: full.GetVersion(), Full: true, ValidUntil: validUntil,
		Proxies: proxies, Routing: routes,
	}
	envelope.Checksum = snapshot.ChecksumWithRouting(proxies, routes)
	return watcher.store.Apply(envelope)
}

func wireRouting(source []*controlplanev1.RoutingRule) ([]routingDomain.Rule, error) {
	result := make([]routingDomain.Rule, 0, len(source))
	names := make(map[string]struct{}, len(source))
	for _, rule := range source {
		if rule == nil || !workerruntime.ValidIdentifier(rule.GetName()) {
			return nil, ErrInvalidSnapshotWatcher
		}
		if _, duplicate := names[rule.GetName()]; duplicate {
			return nil, ErrInvalidSnapshotWatcher
		}
		names[rule.GetName()] = struct{}{}
		if !rule.GetEnabled() {
			continue
		}
		upstreams, err := wireRoutingUpstreams(rule.GetUpstreams())
		if err != nil {
			return nil, err
		}
		strategy, err := wireRoutingStrategy(rule.GetStrategy(), upstreams)
		if err != nil {
			return nil, err
		}
		action, err := wireUnavailableAction(rule.GetOnUnavailable())
		if err != nil {
			return nil, err
		}
		result = append(result, routingDomain.Rule{
			Name: rule.GetName(),
			Match: routingDomain.Match{
				HostRegex: rule.GetHostRegex(), Methods: append([]string(nil), rule.GetMethods()...),
				PathRegex: rule.GetPathRegex(), Headers: cloneRoutingHeaders(rule.GetHeaders()),
			},
			Upstreams: upstreams, Action: routingDomain.ActionProxy, Strategy: strategy, OnUnavailable: action,
		})
	}
	return result, nil
}

func wireRoutingUpstreams(source []string) ([]string, error) {
	if len(source) == 0 {
		return nil, ErrInvalidSnapshotWatcher
	}
	result := make([]string, len(source))
	seen := make(map[string]struct{}, len(source))
	for index, upstream := range source {
		if !workerruntime.ValidIdentifier(upstream) {
			return nil, ErrInvalidSnapshotWatcher
		}
		if _, duplicate := seen[upstream]; duplicate {
			return nil, ErrInvalidSnapshotWatcher
		}
		seen[upstream] = struct{}{}
		result[index] = upstream
	}
	return result, nil
}

func wireRoutingStrategy(source *controlplanev1.RoutingStrategy, upstreams []string) (routingDomain.Strategy, error) {
	if source == nil {
		return routingDomain.Strategy{}, ErrInvalidSnapshotWatcher
	}
	strategy := routingDomain.Strategy{CurrentUpstream: source.GetCurrentUpstream()}
	switch source.GetType() {
	case controlplanev1.StrategyType_STRATEGY_TYPE_SEQUENTIAL:
		strategy.Type = routingDomain.StrategySequential
		if !containsUpstream(upstreams, strategy.CurrentUpstream) {
			return routingDomain.Strategy{}, ErrInvalidSnapshotWatcher
		}
	case controlplanev1.StrategyType_STRATEGY_TYPE_RANDOM:
		strategy.Type = routingDomain.StrategyRandom
	case controlplanev1.StrategyType_STRATEGY_TYPE_ROUND_ROBIN:
		strategy.Type = routingDomain.StrategyRoundRobin
	case controlplanev1.StrategyType_STRATEGY_TYPE_WEIGHTED:
		strategy.Type = routingDomain.StrategyWeighted
		strategy.Weights = make(map[string]uint32, len(upstreams))
		for _, upstream := range upstreams {
			weight, exists := source.GetWeights()[upstream]
			if !exists || weight == 0 {
				return routingDomain.Strategy{}, ErrInvalidSnapshotWatcher
			}
			strategy.Weights[upstream] = weight
		}
		for upstream := range source.GetWeights() {
			if !containsUpstream(upstreams, upstream) {
				return routingDomain.Strategy{}, ErrInvalidSnapshotWatcher
			}
		}
	case controlplanev1.StrategyType_STRATEGY_TYPE_LEAST_CONNECTIONS:
		strategy.Type = routingDomain.StrategyLeastConnections
	default:
		return routingDomain.Strategy{}, ErrInvalidSnapshotWatcher
	}
	return strategy, nil
}

func wireUnavailableAction(action controlplanev1.UnavailableAction) (routingDomain.OnUnavailableAction, error) {
	switch action {
	case controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_REJECT:
		return routingDomain.OnUnavailableReject, nil
	case controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_WAIT:
		return routingDomain.OnUnavailableWait, nil
	case controlplanev1.UnavailableAction_UNAVAILABLE_ACTION_DIRECT:
		return routingDomain.OnUnavailableDirect, nil
	default:
		return "", fmt.Errorf("%w: unsupported unavailable action", ErrInvalidSnapshotWatcher)
	}
}

func cloneRoutingHeaders(source map[string]string) map[string]string {
	if source == nil {
		return nil
	}
	result := make(map[string]string, len(source))
	for name, value := range source {
		result[name] = value
	}
	return result
}

func containsUpstream(upstreams []string, target string) bool {
	for _, upstream := range upstreams {
		if upstream == target {
			return true
		}
	}
	return false
}

func wireProxies(source []*controlplanev1.OwnedProxy) ([]proxyDomain.Proxy, error) {
	proxies := make([]proxyDomain.Proxy, len(source))
	for index, item := range source {
		proxy, err := wireProxy(item)
		if err != nil {
			return nil, err
		}
		proxies[index] = proxy
	}
	return proxies, nil
}

func wireProxy(source *controlplanev1.OwnedProxy) (proxyDomain.Proxy, error) {
	if source == nil || !workerruntime.ValidIdentifier(source.GetId()) || !workerruntime.ValidIdentifier(source.GetUpstream()) ||
		source.GetHost() == "" || source.GetPort() == 0 || source.GetPort() > math.MaxUint16 || source.GetMaxConcurrency() == 0 {
		return proxyDomain.Proxy{}, ErrInvalidSnapshotWatcher
	}
	scheme, ok := wireScheme(source.GetProtocol())
	if !ok {
		return proxyDomain.Proxy{}, ErrInvalidSnapshotWatcher
	}
	expiresAt, err := wireTimestamp(source.GetExpiresAt())
	if err != nil {
		return proxyDomain.Proxy{}, err
	}
	usableUntil, err := wireTimestamp(source.GetUsableUntil())
	if err != nil {
		return proxyDomain.Proxy{}, err
	}
	tags := make(map[string]string, len(source.GetTags()))
	for key, value := range source.GetTags() {
		tags[key] = value
	}
	return proxyDomain.Proxy{
		ID: source.GetId(), Scheme: scheme, Host: source.GetHost(), Port: uint16(source.GetPort()),
		Username: source.GetUsername(), CredentialVersion: source.GetCredentialVersion(), SecretRef: source.GetSecretRef(),
		SourceUpstream: source.GetUpstream(), ExpiresAt: expiresAt, UsableUntil: usableUntil,
		MaxConcurrency: int64(source.GetMaxConcurrency()), State: proxyDomain.StateAvailable, Tags: tags,
	}, nil
}

func wireScheme(protocol controlplanev1.ProxyProtocol) (proxyDomain.Scheme, bool) {
	switch protocol {
	case controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTP:
		return proxyDomain.SchemeHTTP, true
	case controlplanev1.ProxyProtocol_PROXY_PROTOCOL_HTTPS:
		return proxyDomain.SchemeHTTPS, true
	case controlplanev1.ProxyProtocol_PROXY_PROTOCOL_SOCKS5:
		return proxyDomain.SchemeSOCKS5, true
	default:
		return "", false
	}
}

func wireTimestamp(value *timestamppb.Timestamp) (*time.Time, error) {
	if value == nil {
		return nil, nil
	}
	if err := value.CheckValid(); err != nil {
		return nil, ErrInvalidSnapshotWatcher
	}
	converted := value.AsTime().UTC()
	return &converted, nil
}

func requiredFutureTimestamp(value *timestamppb.Timestamp, now time.Time) (time.Time, error) {
	if value == nil || value.CheckValid() != nil || now.IsZero() {
		return time.Time{}, ErrInvalidSnapshotWatcher
	}
	converted := value.AsTime().UTC()
	if !converted.After(now) {
		return time.Time{}, ErrInvalidSnapshotWatcher
	}
	return converted, nil
}

type generatedSnapshotRPCClient struct {
	client controlplanev1.WorkerControlPlaneClient
}

func NewGeneratedSnapshotRPCClient(client controlplanev1.WorkerControlPlaneClient) SnapshotRPCClient {
	return generatedSnapshotRPCClient{client: client}
}

func (client generatedSnapshotRPCClient) Watch(ctx context.Context, request *controlplanev1.WatchSnapshotsRequest) (SnapshotStream, error) {
	return client.client.WatchSnapshots(ctx, request, grpc.WaitForReady(true))
}

func (client generatedSnapshotRPCClient) Acknowledge(ctx context.Context, request *controlplanev1.AcknowledgeSnapshotRequest) (*emptypb.Empty, error) {
	return client.client.AcknowledgeSnapshot(ctx, request, grpc.WaitForReady(true))
}