package models

import (
	"context"
	"log"
	"strconv"

	"github.com/milvus-io/milvus-sdk-go/v2/client"
)

type MilvusConnection struct {
	Client client.Client
}

func (m *MilvusConnection) GetClient() client.Client {
	return m.Client
}

func (m *MilvusConnection) Init() (err error) {
	log.Println("Milvus connection init")
	host := Viper.Get("milvus.host").(string)
	port := Viper.Get("milvus.port").(int)
	address_of_milvus := host + ":" + strconv.Itoa(port)
	log.Println("Milvus address:", address_of_milvus)
	m.Client, err = client.NewGrpcClient(context.Background(), address_of_milvus)
	if err != nil {
		log.Println("Milvus connection failed:", err)
		return
	}
	log.Println("Milvus connection success")
	return
}

func (m *MilvusConnection) GetCollection(collection_name string) (collection *client.Client, err error) {
	if m.Client == nil {
		m.Init()
	}
	err = m.Client.LoadCollection(context.Background(), collection_name, false)
	if err != nil {
		log.Println("Milvus load collection failed:", err)
		return
	}
	collection = &m.Client
	return
}