package com.smtaiserver.smtaiserver.control;

import com.smtaiserver.smtaiserver.core.SMTAIServerApp;
import com.smtaiserver.smtaiserver.core.SMTAIServerRequest;
import com.smtaiserver.smtaiserver.core.SMTCheckChatStreamView;
import com.smtaiserver.smtaiserver.database.SMTDatabase;
import com.smtaiserver.smtaiserver.database.SMTDatabase.DBRecord;
import com.smtaiserver.smtaiserver.database.SMTDatabase.DBRecords;
import com.smtaiserver.smtaiserver.javaai.llm.core.SMTLLMConnect;
import com.smtaiserver.smtaiserver.javaai.qwen.agent.SMTQwenAgent;
import com.smtservlet.util.SMTJsonWriter;
import com.smtservlet.util.SMTStatic;
import java.io.BufferedReader;
import java.io.StringReader;
import java.util.*;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.multipart.MultipartFile;
import org.springframework.web.servlet.ModelAndView;

public class SMTKnowledgeControl
{
	public ModelAndView getKnowledgeText(SMTAIServerRequest tranReq) throws Exception 
	{
		String id = tranReq.convParamToString("file_id", true);
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		SMTJsonWriter jsonWr = tranReq.newReturnJsonWriter(true, null, null);
			DBRecords recs = db.querySQL("SELECT file_text FROM doc_file_list WHERE file_id=?", new Object[] {
				id
			});
			
			if(recs.getRowCount() == 0)
				return tranReq.returnJsonState(false, "file id不存在", null);
			
			jsonWr.addKeyValue("value", recs.getRecord(0).getString(0));
			return tranReq.returnJson(jsonWr);
    	}
    	finally
    	{
    		db.close();
    	}
	}
	
	public ModelAndView listKnowledgeVector(SMTAIServerRequest tranReq) throws Exception 
	{
		String id = tranReq.convParamToString("file_id", true);
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		SMTJsonWriter jsonWr = tranReq.newReturnJsonWriter(true, null, null);
			DBRecords recs = db.querySQL("SELECT vector_text FROM doc_file_vector WHERE file_id=? ORDER BY vector_index", new Object[] {
				id
			});
			
			jsonWr.beginArray("values");
			for(DBRecord rec : recs.getRecords())
			{
				jsonWr.addKeyValue(null, rec.getString(0));
			}
			jsonWr.endArray();
			
			return tranReq.returnJson(jsonWr);
    	}
    	finally
    	{
    		db.close();
    	}
	}
	
	public ModelAndView checkDocvectorValidate(SMTAIServerRequest tranReq) throws Exception 
	{
		String id = tranReq.convParamToString("knowlg_id", true);
		String question = tranReq.convParamToString("question", true);
		
		DBRecords recs;
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
			recs = db.querySQL("SELECT * FROM ai_agent_knowlg WHERE knowlg_id=?", new Object[] {
				id
			});
			if(recs.getRowCount() == 0)
				return tranReq.returnJsonState(false, "未发现记录", null);	   		
    	}
    	finally
    	{
    		db.close();
    	}
    	
		DBRecord rec = recs.getRecord(0);
		SMTQwenAgent agent = (SMTQwenAgent) Class.forName(rec.getString("clz_name")).newInstance();
		agent.initInstance(rec);
		Map<String, SMTQwenAgent> mapId2Agent = new HashMap<>();
		mapId2Agent.put(agent.getAgentId(), agent);
		
		return new ModelAndView(new SMTCheckChatStreamView(tranReq, question, mapId2Agent));

	}
	
	public ModelAndView getDocvectorFileList(SMTAIServerRequest tranReq) throws Exception 
	{
		String id = tranReq.convParamToString("knowlg_id", true);
		
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		SMTJsonWriter jsonWr = tranReq.newReturnJsonWriter(true, null, null);
    		DBRecords recs = db.querySQL(
    			  " SELECT file_id, file_name, file_type, create_time FROM doc_file_list"
    			+ " WHERE file_id IN (SELECT knowlg_file_id FROM knowlg_file_list WHERE knowlg_id=?)"
    			  , new Object[] {id});
    		jsonWr.beginArray("values");
    		for(DBRecord rec : recs.getRecords())
    		{
    			jsonWr.beginMap(null);
    			{
    				jsonWr.addKeyValue("id", rec.getString("file_id"));
    				jsonWr.addKeyValue("name", rec.getString("file_name"));
    				jsonWr.addKeyValue("type", rec.getString("file_type"));
    				jsonWr.addKeyValue("time", rec.getString("create_time"));
    			}
    			jsonWr.endMap();
    		}
    		jsonWr.endArray();
    		
    		return tranReq.returnJson(jsonWr);
    	}
    	finally
    	{
    		db.close();
    	}
	} 
	
	public ModelAndView deleteDocvectorFile(SMTAIServerRequest tranReq) throws Exception 
	{
		String id = tranReq.convParamToString("knowlg_id", true);
		String fileId = tranReq.convParamToString("file_id", true);
		
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		db.executeSQL("DELETE FROM knowlg_file_list WHERE knowlg_id=? AND knowlg_file_id=?", new Object[] {
    			id,
    			fileId
    		});
    		
    		return tranReq.returnJsonState(true, null, null);
    	}
    	finally
    	{
    		db.close();
    	}
	} 
	
	public ModelAndView addDocvectorFile(SMTAIServerRequest tranReq) throws Exception 
	{
		String id = tranReq.convParamToString("knowlg_id", true);
		String fileId = tranReq.convParamToString("file_id", true);
		
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		db.executeSQL("DELETE FROM knowlg_file_list WHERE knowlg_id=? AND knowlg_file_id=?", new Object[] {
    			id,
    			fileId
    		});
    		
    		db.executeSQL("INSERT INTO knowlg_file_list(knowlg_id, knowlg_file_id)VALUES(?,?)", new Object[] {
    			id,
    			fileId
    		});
    		
    		return tranReq.returnJsonState(true, null, null);
    	}
    	finally
    	{
    		db.close();
    	}
	} 
	
	public ModelAndView publishDocvectorName(SMTAIServerRequest tranReq) throws Exception 
	{
		String id = tranReq.convParamToString("knowlg_id", true);
		String publish = tranReq.convParamToString("publish", true);
		
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		DBRecords recs = db.querySQL("SELECT is_publish FROM ai_agent_knowlg WHERE knowlg_id=?", new Object[] {
        		id
    		});
    		if(recs.getRowCount() == 0)
    			return tranReq.returnJsonState(false, "未发现知识库id", null);
    		
    		db.executeSQL("UPDATE ai_agent_knowlg SET is_publish=? WHERE knowlg_id=?", new Object[] {
        		publish,
        		id
    		});
    		
    		if(!publish.equals(recs.getRecord(0).getString(0)))
    			SMTAIServerApp.getApp().clearQwenAgentManager();
    		
    		return tranReq.returnJsonState(true, null, null);
    	}
    	finally
    	{
    		db.close();
    	}
	}     	
	
	public ModelAndView updateDocvectorGroup(SMTAIServerRequest tranReq) throws Exception 
	{
		String id = tranReq.convParamToString("knowlg_id", true);
		String groupId = tranReq.convParamToString("group_id", true);
		
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		db.executeSQL("UPDATE ai_agent_knowlg SET knowlg_title=? WHERE knowlg_id=?", new Object[] {
    			groupId,
    			id
    		});
    		
    		return tranReq.returnJsonState(true, null, null);
    	}
    	finally
    	{
    		db.close();
    	}
	}
	
	public ModelAndView updateDocvectorName(SMTAIServerRequest tranReq) throws Exception 
	{
		String id = tranReq.convParamToString("knowlg_id", true);
		String title = tranReq.convParamToString("title", true);
		String prompt = tranReq.convParamToString("prompt", true);
		int scope = tranReq.convParamToInteger("scope", true);
		
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		db.executeSQL("UPDATE ai_agent_knowlg SET knowlg_title=?, knowlg_prompt=?, vector_scope=? WHERE knowlg_id=?", new Object[] {
    			title,
    			prompt,
    			scope,
    			id
    		});
    		
    		return tranReq.returnJsonState(true, null, null);
    	}
    	finally
    	{
    		db.close();
    	}
	}     		
    		
	public ModelAndView deleteDocvectorName(SMTAIServerRequest tranReq) throws Exception 
	{
		String id = tranReq.convParamToString("knowlg_id", true);
		
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		db.executeSQL("BEGIN", null);
    		try
    		{
    			db.executeSQL("DELETE FROM ai_agent_knowlg WHERE knowlg_id=?", new Object[] {id});
    			db.executeSQL("DELETE FROM knowlg_file_list WHERE knowlg_id=?", new Object[] {id});
    			db.executeSQL("COMMIT", null);
    			return tranReq.returnJsonState(true, null, null);
    		}
    		catch(Exception ex)
    		{
    			db.rollbackDB(ex);
    			return null;
    		}
    	}
    	finally
    	{
    		db.close();
    	}
	}   		
	
	public ModelAndView addDocvectorName(SMTAIServerRequest tranReq) throws Exception 
	{
		String title = tranReq.convParamToString("title", true);
		String prompt = tranReq.convParamToString("prompt", true);
		String groupId = tranReq.convParamToString("group_id", true);
		int scope = tranReq.convParamToInteger("scope", true);
		
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		SMTJsonWriter jsonWr = tranReq.newReturnJsonWriter(true, null, null);
    		String id = "agent_" + SMTStatic.newUUID();
    		db.executeSQL(
    			"INSERT INTO ai_agent_knowlg(knowlg_id, knowlg_title, knowlg_prompt, is_publish, create_time, create_user,vector_scope, agent_group)VALUES(?,?,?,?,?,?,?,?)"
    			, new Object[] {
    				id,
    				title,
    				prompt,
    				"N",
    				new Date(),
    				tranReq.getLoginUserId(),
    				scope,
    				groupId
    			});
    		
    		jsonWr.addKeyValue("knowlg_id", id);
    		return tranReq.returnJson(jsonWr);
    	}
    	finally
    	{
    		db.close();
    	}
	}
    	
	
	public ModelAndView getDocvectorList(SMTAIServerRequest tranReq) throws Exception 
	{
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		SMTJsonWriter jsonWr = tranReq.newReturnJsonWriter(true, null, null);
    		
    		DBRecords recs = db.querySQL(
    			  " SELECT A.knowlg_id, A.knowlg_title, A.knowlg_prompt, A.is_publish, A.create_time, B.user_name, A.vector_scope, A.agent_group"
    			+ " FROM ai_agent_knowlg A LEFT JOIN sys_user_info B ON A.create_user=B.user_id"
    			, null);
    		
    		jsonWr.beginArray("values");
    		for(DBRecord rec : recs.getRecords())
    		{
    			jsonWr.beginMap(null);
    			{
    				jsonWr.addKeyValue("id", rec.getString("knowlg_id"));
    				jsonWr.addKeyValue("title", rec.getString("knowlg_title"));
    				jsonWr.addKeyValue("prompt", rec.getString("knowlg_prompt"));
    				jsonWr.addKeyValue("publish", rec.getString("is_publish"));
    				jsonWr.addKeyValue("create_time", rec.getString("create_time"));
    				jsonWr.addKeyValue("user_name", rec.getString("user_name"));
    				jsonWr.addKeyValue("scope", rec.getString("vector_scope"));
    				jsonWr.addKeyValue("group_id", rec.getString("agent_group"));
    			}
    			jsonWr.endMap();
    		}
    		
    		jsonWr.endArray();
    		
    		return tranReq.returnJson(jsonWr);
    	}
    	finally
    	{
    		db.close();
    	}
	}
	
	public ModelAndView listKnowledgeFile(SMTAIServerRequest tranReq) throws Exception 
	{
		String groupId = tranReq.convParamToString("group_id", true);
		
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		SMTJsonWriter jsonWr = tranReq.newReturnJsonWriter(true, null, null);
    		DBRecords recs = db.querySQL("SELECT file_id, file_name, file_type, create_time, length(file_text) as file_size FROM doc_file_list WHERE file_group=?", new Object[] {groupId});
    		jsonWr.beginArray("values");
    		for(DBRecord rec : recs.getRecords())
    		{
    			jsonWr.beginMap(null);
    			{
    				jsonWr.addKeyValue("id", rec.getString("file_id"));
    				jsonWr.addKeyValue("name", rec.getString("file_name"));
    				jsonWr.addKeyValue("type", rec.getString("file_type"));
    				jsonWr.addKeyValue("time", rec.getString("create_time"));
    				jsonWr.addKeyValue("size", rec.getInteger("file_size"));
    			}
    			jsonWr.endMap();
    		}
    		jsonWr.endArray();
    		
    		return tranReq.returnJson(jsonWr);
    	}
    	finally
    	{
    		db.close();
    	}		
	}
	
	public ModelAndView deleteKnowledgeFile(SMTAIServerRequest tranReq) throws Exception 
	{
		String fileId = tranReq.convParamToString("file_id", true);
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		db.executeSQL("BEGIN", null);
    		try
    		{
    			db.executeSQL("DELETE FROM doc_file_list WHERE file_id=?", new Object[] {fileId});
    			db.executeSQL("DELETE FROM doc_file_vector WHERE file_id=?", new Object[] {fileId}); 
    			db.executeSQL("DELETE FROM knowlg_file_list WHERE knowlg_file_id=?", new Object[] {fileId}); 
    			db.executeSQL("COMMIT", null);
    		}
    		catch(Exception ex)
    		{
    			db.rollbackDB(ex);
    		}
    		
    		return tranReq.returnJsonState(true, null, null);
    	}
    	finally
    	{
    		db.close();
    	}		
	}
	
    public ModelAndView addKnowledgeFile(SMTAIServerRequest tranReq, @RequestParam(value = "file", required = false) MultipartFile file) throws Exception 
    {
		String fileText = "";
		String fileName = file.getOriginalFilename();
		String groupId = tranReq.convParamToString("group_id", true);
		fileText = SMTStatic.readTextStream(file.getInputStream());

		String fileType = "";
		List<String> listBlock = null;

		// 先判断文件名是否为空，避免空指针异常
		if (fileName != null && fileName.contains(".")) {
		  fileType = fileName.substring(fileName.lastIndexOf(".")).toLowerCase(); // 统一转换为小写，避免大小写问题
		}

		// 用 Set 存储支持的文件类型，提高可读性
		Set<String> supportedTypes = new HashSet<>(Arrays.asList(".md", ".pdf", ".docx", ".doc"));
		  if (supportedTypes.contains(fileType)) {
		  fileText = SMTAIServerApp.fileTranslTxt(file);
		} else {
		  return tranReq.returnJsonState(false, "文件格式不支持", null);
		}
		listBlock = splitMarkdownToVectorBlock(fileText);
    	List<String> listVector = new ArrayList<>();
    	SMTLLMConnect llm = SMTAIServerApp.getApp().allocLLMConnect(null);
    	try
    	{
    		for(String text : listBlock)
    		{
    			listVector.add(llm.getVector(text));
    		}
    	}
    	finally
    	{
    		llm.close();
    	}
    	
    	
    	SMTDatabase db = SMTAIServerApp.getApp().allocDatabase();
    	try
    	{
    		db.executeSQL("BEGIN", null);
    		try
    		{
	    		String fileId = SMTStatic.newUUID();
	    		db.executeSQL("INSERT INTO doc_file_list(file_id, file_name, file_type, create_time, file_text, file_group)VALUES(?,?,?,?,?,?)", new Object[] {
	    			fileId,
	    			fileName,
	    			fileType,
	    			new Date(),
	    			fileText,
	    			groupId
	    		});
	    		
	    		for(int i = 0; i < listBlock.size(); i ++)
	    		{
	    			db.executeSQL("INSERT INTO doc_file_vector(file_id, vector_index, vector_text, vector_value)VALUES(?,?,?,?::vector)", new Object[] {
	    				fileId,
	    				i,
	    				listBlock.get(i),
	    				"[" + listVector.get(i) + "]"
	    			});
	    		}
	    		
	    		db.executeSQL("COMMIT", null);
	    		
	    		SMTJsonWriter jsonWr = tranReq.newReturnJsonWriter(true, null, null);
	    		jsonWr.addKeyValue("file_id", fileId);
	    		return tranReq.returnJson(jsonWr);
    		}
    		catch(Exception ex)
    		{
    			db.rollbackDB(ex);
    			return null;
    		}
    	}
    	finally
    	{
    		db.close();
    	}

    }
    
    private List<String> splitMarkdownToVectorBlock(String fileText) throws Exception
    {
        List<String> blocks = new ArrayList<>();
        String currentBlock = "";

        BufferedReader reader = new BufferedReader(new StringReader(fileText));
        try
        {
            String line;
            while ((line = reader.readLine())!= null) {
                // 忽略空行
                if (line.matches("\\s*?\\n")) {
                    continue;
                }

                // 当遇到标题或者是新的段落时，结束当前块
                if (line.matches("#+.+") || line.matches("\\s*?\\n\\s*?\\n")) {
                    if (!currentBlock.isEmpty()) {
                        blocks.add(currentBlock);
                        currentBlock = "";
                    }
                }

                // 添加行到当前块
                currentBlock += line + "\n";
            }

            // 添加最后一个块
            if (!currentBlock.isEmpty()) {
                blocks.add(currentBlock);
            }
        } 
        finally
        {
        	reader.close();
        }

        return blocks;
    }
}