internal/agent/server.go

// Package agent provides gRPC server implementation for the Agent service
package agent

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"sync"
	"time"

	ragv1 "github.com/AmaniQuery/amaniquery/pkg/proto/gen/ragv1"

	"go.uber.org/zap"
	"google.golang.org/grpc"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
)

// Server implements the AgentService gRPC server
type Server struct {
	ragv1.UnimplementedAgentServiceServer
	orchestrator *Orchestrator
	logger       *zap.Logger
	mu           sync.RWMutex
}

// NewServer creates a new Agent gRPC server
func NewServer(deps *Dependencies, logger *zap.Logger) *Server {
	return &Server{
		orchestrator: NewOrchestrator(deps, logger),
		logger:       logger,
	}
}

// ProcessQuery handles a single query via gRPC (unary RPC)
func (s *Server) ProcessQuery(ctx context.Context, req *ragv1.QueryRequest) (*ragv1.QueryResponse, error) {
	if req == nil || req.Query == "" {
		return nil, status.Error(codes.InvalidArgument, "query is required")
	}

	s.logger.Info("processing query",
		zap.String("query", req.Query),
		zap.String("session_id", req.SessionId),
		zap.String("user_id", req.UserId),
	)

	// Convert proto request to internal type
	internalReq := s.protoToQueryRequest(req)

	// Process query through orchestrator
	response, err := s.orchestrator.ProcessQuery(ctx, internalReq)
	if err != nil {
		s.logger.Error("query processing failed",
			zap.Error(err),
			zap.String("query", req.Query),
		)
		return nil, status.Error(codes.Internal, "query processing failed: "+err.Error())
	}

	// Convert internal response to proto
	return s.queryResponseToProto(response), nil
}

// ProcessQueryStream handles a query with streaming response via gRPC (server streaming RPC)
func (s *Server) ProcessQueryStream(req *ragv1.QueryRequest, stream ragv1.AgentService_ProcessQueryStreamServer) error {
	if req == nil || req.Query == "" {
		return status.Error(codes.InvalidArgument, "query is required")
	}

	s.logger.Info("processing streaming query",
		zap.String("query", req.Query),
		zap.String("session_id", req.SessionId),
	)

	// Convert proto request to internal type
	internalReq := s.protoToQueryRequest(req)

	// Create a wrapper that implements ResponseStream
	streamWrapper := &grpcStreamWrapper{
		stream: stream,
		logger: s.logger,
	}

	return s.orchestrator.ProcessQueryStream(stream.Context(), internalReq, streamWrapper)
}

// CreateAgent creates a new agent with specific configuration
func (s *Server) CreateAgent(ctx context.Context, req *ragv1.CreateAgentRequest) (*ragv1.Agent, error) {
	if req == nil || req.Name == "" {
		return nil, status.Error(codes.InvalidArgument, "agent name is required")
	}

	s.logger.Info("creating agent", zap.String("name", req.Name))

	// TODO: Implement agent creation logic with persistence
	// For now, return a mock agent
	return &ragv1.Agent{
		Id:          generateAgentID(),
		Name:        req.Name,
		Description: req.Description,
		CreatedAt:   currentTimestamp(),
		Config:      req.Config,
	}, nil
}

// ExecutePlan executes a pre-defined execution plan
func (s *Server) ExecutePlan(ctx context.Context, req *ragv1.ExecutionPlan) (*ragv1.PlanResult, error) {
	if req == nil || req.Id == "" {
		return nil, status.Error(codes.InvalidArgument, "plan ID is required")
	}

	s.logger.Info("executing plan",
		zap.String("plan_id", req.Id),
		zap.Int("steps", len(req.Steps)),
	)

	// TODO: Implement multi-step plan execution
	return &ragv1.PlanResult{
		PlanId:          req.Id,
		Status:          ragv1.ExecutionStatus_EXECUTION_STATUS_COMPLETED,
		StepResults:     make([]*ragv1.StepResult, 0),
		FinalAnswer:     "Plan execution not yet implemented",
		ExecutionTimeMs: 0,
	}, nil
}

// GetQueryStatus returns the status of an ongoing query
func (s *Server) GetQueryStatus(ctx context.Context, req *ragv1.QueryStatusRequest) (*ragv1.QueryStatus, error) {
	if req == nil || req.QueryId == "" {
		return nil, status.Error(codes.InvalidArgument, "query_id is required")
	}

	// TODO: Implement query status tracking
	return &ragv1.QueryStatus{
		QueryId:     req.QueryId,
		Status:      ragv1.ExecutionStatus_EXECUTION_STATUS_COMPLETED,
		Progress:    100,
		CurrentStep: "completed",
	}, nil
}

// RegisterTool registers a tool with the orchestrator
func (s *Server) RegisterTool(name string, tool Tool) error {
	return s.orchestrator.RegisterTool(name, tool)
}

// grpcStreamWrapper wraps gRPC stream to implement ResponseStream interface
type grpcStreamWrapper struct {
	stream ragv1.AgentService_ProcessQueryStreamServer
	logger *zap.Logger
}

// Send implements ResponseStream.Send by converting internal chunk to protobuf
func (w *grpcStreamWrapper) Send(chunk *ResponseChunk) error {
	w.logger.Debug("sending stream chunk",
		zap.String("type", chunk.Type),
		zap.Bool("is_final", chunk.IsFinal),
		zap.Int("content_length", len(chunk.Content)),
	)

	// Convert chunk type string to proto enum
	chunkType := w.stringToChunkType(chunk.Type)

	// Convert sources to proto format
	protoSources := make([]*ragv1.Source, len(chunk.Sources))
	for i, src := range chunk.Sources {
		protoSources[i] = &ragv1.Source{
			Id:       src.ID,
			Title:    src.Title,
			Content:  src.Content,
			Url:      src.URL,
			Score:    src.Score,
			Metadata: src.Metadata,
			Location: src.Location,
		}
	}

	// Create proto chunk and send
	protoChunk := &ragv1.QueryResponseChunk{
		Type:    chunkType,
		Content: chunk.Content,
		Sources: protoSources,
		IsFinal: chunk.IsFinal,
	}

	return w.stream.Send(protoChunk)
}

// stringToChunkType converts string chunk type to proto enum
func (w *grpcStreamWrapper) stringToChunkType(t string) ragv1.ChunkType {
	switch t {
	case "thinking":
		return ragv1.ChunkType_CHUNK_TYPE_THINKING
	case "retrieval":
		return ragv1.ChunkType_CHUNK_TYPE_RETRIEVAL
	case "generation":
		return ragv1.ChunkType_CHUNK_TYPE_GENERATION
	case "complete":
		return ragv1.ChunkType_CHUNK_TYPE_COMPLETE
	case "error":
		return ragv1.ChunkType_CHUNK_TYPE_ERROR
	default:
		return ragv1.ChunkType_CHUNK_TYPE_UNSPECIFIED
	}
}

// protoToQueryRequest converts proto QueryRequest to internal type
func (s *Server) protoToQueryRequest(req *ragv1.QueryRequest) *QueryRequest {
	// Convert conversation history
	history := make([]Message, len(req.ConversationHistory))
	for i, msg := range req.ConversationHistory {
		history[i] = Message{
			Role:    s.protoRoleToString(msg.Role),
			Content: msg.Content,
		}
	}

	// Set default options if not provided
	maxSources := int(req.Options.GetMaxSources())
	if maxSources == 0 {
		maxSources = 10
	}
	maxTokens := int(req.Options.GetMaxTokens())
	if maxTokens == 0 {
		maxTokens = 4096
	}
	temperature := req.Options.GetTemperature()
	if temperature == 0 {
		temperature = 0.7
	}

	return &QueryRequest{
		Query:               req.Query,
		SessionID:           req.SessionId,
		UserID:              req.UserId,
		ConversationHistory: history,
		Metadata:            req.Metadata,
		Options: QueryOptions{
			MaxSources:     maxSources,
			UseCache:       req.Options.GetUseCache(),
			EnableAgentic:  req.Options.GetEnableAgentic(),
			KnowledgeBases: req.Options.GetKnowledgeBases(),
			Temperature:    temperature,
			MaxTokens:      maxTokens,
		},
	}
}

// queryResponseToProto converts internal QueryResponse to proto type
func (s *Server) queryResponseToProto(resp *QueryResponse) *ragv1.QueryResponse {
	// Convert sources
	protoSources := make([]*ragv1.Source, len(resp.Sources))
	for i, src := range resp.Sources {
		protoSources[i] = &ragv1.Source{
			Id:       src.ID,
			Title:    src.Title,
			Content:  src.Content,
			Url:      src.URL,
			Score:    src.Score,
			Metadata: src.Metadata,
			Location: src.Location,
		}
	}

	return &ragv1.QueryResponse{
		Answer:  resp.Answer,
		Sources: protoSources,
		Confidence: resp.Confidence,
		Metadata: &ragv1.QueryMetadata{
			ProcessingTimeMs: resp.Metadata.ProcessingTimeMs,
			ChunksRetrieved:  int32(resp.Metadata.ChunksRetrieved),
			TokensUsed:       int32(resp.Metadata.TokensUsed),
			CacheHit:         resp.Metadata.CacheHit,
			RoutingStrategy:  resp.Metadata.RoutingStrategy,
			TraceId:          resp.Metadata.TraceID,
		},
		FollowUpQuestions: resp.FollowUpQuestions,
	}
}

// protoRoleToString converts proto MessageRole to string
func (s *Server) protoRoleToString(role ragv1.MessageRole) string {
	switch role {
	case ragv1.MessageRole_MESSAGE_ROLE_USER:
		return "user"
	case ragv1.MessageRole_MESSAGE_ROLE_ASSISTANT:
		return "assistant"
	case ragv1.MessageRole_MESSAGE_ROLE_SYSTEM:
		return "system"
	default:
		return "user"
	}
}

// RegisterAgentServiceServer registers the server with gRPC using generated code
func RegisterAgentServiceServer(s *grpc.Server, srv *Server) {
	ragv1.RegisterAgentServiceServer(s, srv)
}

// Helper functions

func generateAgentID() string {
	b := make([]byte, 16)
	rand.Read(b)
	return "agent_" + hex.EncodeToString(b)
}

func currentTimestamp() int64 {
	return time.Now().Unix()
}