package api

import (
	"context"

	pb "gitlab.ost.ch/marketminds/marketminds/pkg/protobuf"
	"google.golang.org/grpc"
)

type NEntity struct {
	Name     string
	Type     string
	Saliency float32
}

type NERResponse struct {
	Entities []*NEntity
}

func NewNERResponse(entities []*pb.Entity) *NERResponse {
	return &NERResponse{
		Entities: func() []*NEntity {
			var nEntities []*NEntity
			for _, entity := range entities {
				nEntities = append(nEntities, &NEntity{
					Name:     entity.Name,
					Type:     entity.Type,
					Saliency: entity.Salience,
				})
			}
			return nEntities
		}(),
	}
}

type AnalyzeNER struct {
	Client pb.NamedEntityRecognizerClient
}

func NewAnalyzeNER(conn *grpc.ClientConn) *AnalyzeNER {
	client := pb.NewNamedEntityRecognizerClient(conn)
	return &AnalyzeNER{Client: client}
}

func (an *AnalyzeNER) Analyze(text string) (*NERResponse, error) {
	request := &pb.RecognizeEntitiesRequest{
		Text: text,
	}

	response, err := an.Client.RecognizeEntities(context.Background(), request)
	if err != nil {
		return nil, err
	}

	return NewNERResponse(response.Entities), nil
}
