// Package agent implements the client side of the RVBox Agent protocol.
package agent

import (
	"context"
	"errors"
	"fmt"

	rvboxv1 "github.com/rvbox/rvbox/gen/go/rvbox/v1"
	"github.com/rvbox/rvbox/internal/agentproto"
	"google.golang.org/protobuf/proto"
)

var (
	ErrProtocolHandshake = errors.New("invalid agent session handshake")
	ErrSessionClosed     = errors.New("agent session transport closed")
)

// Transport has exactly the binary-message operations the protocol needs. A
// WebSocket implementation owns its deadline/ping details outside this layer.
type Transport interface {
	Write(context.Context, []byte) error
	Read(context.Context) ([]byte, error)
	Close() error
}

type Session struct {
	ID         string
	Generation uint64
	Protocol   *rvboxv1.ProtocolVersion
}

func Handshake(ctx context.Context, transport Transport, hello *rvboxv1.ClientHello, limits agentproto.Limits) (Session, error) {
	if transport == nil || hello == nil {
		return Session{}, ErrProtocolHandshake
	}
	request := &rvboxv1.AgentEnvelope{Payload: &rvboxv1.AgentEnvelope_ClientHello{ClientHello: hello}}
	if err := agentproto.ValidateEnvelope(request, limits, rvboxv1.Platform_PLATFORM_WINDOWS); err != nil {
		return Session{}, fmt.Errorf("validate ClientHello: %w", err)
	}
	encoded, err := proto.Marshal(request)
	if err != nil {
		return Session{}, err
	}
	if err := transport.Write(ctx, encoded); err != nil {
		return Session{}, err
	}
	received, err := transport.Read(ctx)
	if err != nil {
		return Session{}, err
	}
	welcome, err := agentproto.DecodeEnvelope(received, limits, rvboxv1.Platform_PLATFORM_WINDOWS)
	if err != nil {
		return Session{}, fmt.Errorf("decode ServerWelcome: %w", err)
	}
	payload := welcome.GetServerWelcome()
	if payload == nil || welcome.GetSessionId() == "" || welcome.GetSessionGeneration() == 0 || payload.GetSelectedProtocol() == nil {
		return Session{}, ErrProtocolHandshake
	}
	return Session{ID: welcome.GetSessionId(), Generation: welcome.GetSessionGeneration(), Protocol: proto.Clone(payload.GetSelectedProtocol()).(*rvboxv1.ProtocolVersion)}, nil
}

// Reconcile sends the complete retained client state after welcome. The caller
// must apply the returned result durably before accepting fresh dispatch.
func Reconcile(ctx context.Context, transport Transport, session Session, snapshot *rvboxv1.ReconcileSnapshot, limits agentproto.Limits) (*rvboxv1.ReconcileResult, error) {
	if transport == nil || session.ID == "" || session.Generation == 0 || snapshot == nil {
		return nil, ErrProtocolHandshake
	}
	request := &rvboxv1.AgentEnvelope{SessionId: session.ID, SessionGeneration: session.Generation, Payload: &rvboxv1.AgentEnvelope_ReconcileSnapshot{ReconcileSnapshot: snapshot}}
	if err := agentproto.ValidateEnvelope(request, limits, rvboxv1.Platform_PLATFORM_WINDOWS); err != nil {
		return nil, fmt.Errorf("validate reconciliation snapshot: %w", err)
	}
	encoded, err := proto.Marshal(request)
	if err != nil {
		return nil, err
	}
	if err := transport.Write(ctx, encoded); err != nil {
		return nil, err
	}
	received, err := transport.Read(ctx)
	if err != nil {
		return nil, err
	}
	result, err := agentproto.DecodeEnvelope(received, limits, rvboxv1.Platform_PLATFORM_WINDOWS)
	if err != nil {
		return nil, fmt.Errorf("decode ReconcileResult: %w", err)
	}
	if result.GetSessionId() != session.ID || result.GetSessionGeneration() != session.Generation || result.GetReconcileResult() == nil {
		return nil, ErrProtocolHandshake
	}
	return proto.Clone(result.GetReconcileResult()).(*rvboxv1.ReconcileResult), nil
}