package com.smtaiserver.smtaiserver.javaai.llm.qwen;

import org.apache.commons.pool2.BasePooledObjectFactory;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.embeddings.TextEmbedding;
import com.alibaba.dashscope.utils.Constants;
import com.smtaiserver.smtaiserver.database.SMTDatabase.DBRecord;
import com.smtaiserver.smtaiserver.javaai.llm.core.SMTLLMConnect;
import com.smtaiserver.smtaiserver.javaai.llm.core.SMTLLMFactory;
import com.smtservlet.util.Json;

public class SMTLLMFactoryQwen extends SMTLLMFactory 
{
	public class SMTQwenLLMPoolFactory extends BasePooledObjectFactory<Generation> 
	{
	    @Override
	    public Generation create() throws Exception {
	        return new Generation();
	    }

	    @Override
	    public PooledObject<Generation> wrap(Generation obj) {
	        return new DefaultPooledObject<>(obj);
	    } 
	}
	
	public class SMTQwenEmbedPoolFactory extends BasePooledObjectFactory<TextEmbedding> 
	{
	    @Override
	    public TextEmbedding create() throws Exception {
	        return new TextEmbedding();
	    }

	    @Override
	    public PooledObject<TextEmbedding> wrap(TextEmbedding obj) {
	        return new DefaultPooledObject<>(obj);
	    } 
	}
	
	////////////////////////////////////////////////////////////////////////////////////////
	private GenericObjectPool<TextEmbedding> 	_qwenEmbedPool;
	private GenericObjectPool<Generation> 		_qwennLLMPool;
	private Json								_jsonFactory;
	private Json								_jsonConnect;
	
	@Override
	public void initInstance(DBRecord rec) throws Exception
	{
		super.initInstance(rec);
		
		_jsonFactory = Json.read(rec.getString("factory_args"));
		_jsonConnect = Json.read(rec.getString("connect_args"));
		
		Constants.apiKey= _jsonFactory.getJson("key").asString(); //"sk-57653d19816c44ed81c731ebadeaa3d9";
		{
	        GenericObjectPoolConfig<Generation> config = new GenericObjectPoolConfig<>();
	        config.setMaxTotal(32);
	        config.setMinIdle(32);
	        _qwennLLMPool = new GenericObjectPool<>(new SMTQwenLLMPoolFactory(), config);
		}
		{
	        GenericObjectPoolConfig<TextEmbedding> config = new GenericObjectPoolConfig<>();
	        config.setMaxTotal(32);
	        config.setMinIdle(32);
	        _qwenEmbedPool = new GenericObjectPool<>(new SMTQwenEmbedPoolFactory(), config);
	   }
	}
	
	public Json getConnectJson()
	{
		return _jsonConnect;
	}
	
	public SMTLLMConnect allocLLM() throws Exception
	{
		Generation gen = _qwennLLMPool.borrowObject();
		return new SMTLLMConnectQwen(gen, this);
	}
	
	public TextEmbedding allocLLMEmbedding() throws Exception
	{
		return _qwenEmbedPool.borrowObject();
	}
	
	public void freeLLMEmbedding(TextEmbedding embed)
	{
		_qwenEmbedPool.returnObject(embed);
	}
	
	public void returnGenToPool(Generation gen)
	{
		_qwennLLMPool.returnObject(gen);
	}
}