#include "stdafx.h"
#include "AskAIPaletteSet.h"
#include "AIChatSettingDlg.h"
#include "sh_ai.h"

#define SHAI_CONFIG_FILE _T("\\shai.json")
namespace
{
	using namespace AIJSON;
	bool readConfig(SHAI::ChatConfig &chat)
	{
		CString sJsonPath(FilePath::GetSysDir() + SHAI_CONFIG_FILE);
		if (_taccess(sJsonPath, 0x0) != 0x0) return false;

		
		Document8 doc;
		if (!file2Json(sJsonPath, doc)) return false;

		CMIt8 itChat = doc.FindMember("chat_config");
		if (itChat == doc.MemberEnd()) return false;

		std::string s;
		if (parseValue<std::string>(itChat->value, "SystemPrompt", s))
			chat.set_system_prompt(s.c_str());

		int i;
		if (parseValue<int>(itChat->value, "MaxTokens", i))
			chat.set_max_tokens(i);

		if (parseValue<int>(itChat->value, "TopK", i))
			chat.set_top_k(i);

		float f;
		if (parseValue<float>(itChat->value, "TopP", f))
			chat.set_top_p(f);

		if (parseValue<float>(itChat->value, "FrequencyPenalty", f))
			chat.set_frequency_penalty(f);

		if (parseValue<float>(itChat->value, "Temperature", f))
			chat.set_temperature(f);

		return true;
	}

	bool writeConfig(const SHAI::ChatConfig &chat)
	{
		CString sJsonPath(FilePath::GetSysDir() + SHAI_CONFIG_FILE);
		if (_taccess(sJsonPath, 0x0) != 0x0)
		{
			CFile file(sJsonPath, CFile::modeCreate | CFile::modeWrite | CFile::typeBinary);
			if (file.m_hFile)
			{
				std::string s("{}");
				file.Write(s.c_str(), (UINT)(sizeof(char) * s.size()));
				file.Close();
			}
			else
				return false;
		}

		Document8 doc;
		if (!file2Json(sJsonPath, doc)) return false;

		MIt8 itChat = doc.FindMember("chat_config");
		if (itChat == doc.MemberEnd())
		{
			Value8 v;
			v.SetObject();
			addMember(v, "SystemPrompt", chat.system_prompt().c_str());
			addMember(v, "MaxTokens", chat.max_tokens());
			addMember(v, "TopK", chat.top_k());
			addMember(v, "TopP", chat.top_p());
			addMember(v, "FrequencyPenalty", chat.frequency_penalty());
			addMember(v, "Temperature", chat.temperature());
			addMember(doc, "chat_config", v);
		}
		else
		{
			if (itChat->value.HasMember("SystemPrompt"))
				itChat->value["SystemPrompt"].SetString(chat.system_prompt().c_str(), getJsonAllocator());

			if (itChat->value.HasMember("MaxTokens"))
				itChat->value["MaxTokens"].SetInt(chat.max_tokens());

			if (itChat->value.HasMember("TopK"))
				itChat->value["TopK"].SetInt(chat.top_k());

			if (itChat->value.HasMember("TopP"))
				itChat->value["TopP"].SetFloat(chat.top_p());

			if (itChat->value.HasMember("Temperature"))
				itChat->value["Temperature"].SetFloat(chat.temperature());

			if (itChat->value.HasMember("FrequencyPenalty"))
				itChat->value["FrequencyPenalty"].SetFloat(chat.frequency_penalty());
		}
		return json2File(sJsonPath, doc, true);
	}
}

IMPLEMENT_DYNAMIC(CAIChatSettingDlg, CDialog)

CAIChatSettingDlg::CAIChatSettingDlg(CWnd* pParent )
	: CDialog(IDD_AI_CHAT_SETTING, pParent)
{
	readConfig(m_chatConfig);
}

CAIChatSettingDlg::~CAIChatSettingDlg() { }

void CAIChatSettingDlg::DoDataExchange(CDataExchange* pDX)
{
	CDialog::DoDataExchange(pDX);
	CString sys(m_chatConfig.system_prompt_w().c_str());
	DDX_Text(pDX, IDC_EDIT_SYSTEM_PROMPT, sys);
	m_chatConfig.set_system_prompt(sys);

	int iMaxTokens = m_chatConfig.max_tokens();
	DDX_Text(pDX, IDC_EDIT_MAX_TOKENS, iMaxTokens);
	DDV_MinMaxInt(pDX, iMaxTokens, 1, 16384);
	m_chatConfig.set_max_tokens(iMaxTokens);

	float fTopP = m_chatConfig.top_p();
	DDX_Text(pDX, IDC_EDIT_TOP_P, fTopP);
	DDV_MinMaxFloat(pDX, fTopP, 0.1f, 1.0f);
	m_chatConfig.set_top_p(fTopP);

	float fTemperature = m_chatConfig.temperature();
	DDX_Text(pDX, IDC_EDIT_TEMPERATURE, fTemperature);
	DDV_MinMaxFloat(pDX, fTemperature, .0f, 2.0f);
	m_chatConfig.set_temperature(fTemperature);

	int iTopK = m_chatConfig.top_k();
	DDX_Text(pDX, IDC_EDIT_TOP_K, iTopK);
	DDV_MinMaxInt(pDX, iTopK, 0, 100);
	m_chatConfig.set_top_k(iTopK);

	float fFrequencyPenalty = m_chatConfig.frequency_penalty();
	DDX_Text(pDX, IDC_EDIT_FREQUENCY_PENALTY, fFrequencyPenalty);
	DDV_MinMaxFloat(pDX, fFrequencyPenalty, -2.0f, 2.0f);
	m_chatConfig.set_frequency_penalty(fFrequencyPenalty);

	DDX_Control(pDX, IDC_SLIDER_MAX_TOKENS, m_wndSlider_MaxTokens);
	DDX_Control(pDX, IDC_SLIDER_TEMPERATURE, m_wndSlider_Temperature);
	DDX_Control(pDX, IDC_SLIDER_TOP_P, m_wndSlider_TopP);
	DDX_Control(pDX, IDC_SLIDER_TOP_K, m_wndSlider_TopK);
	DDX_Control(pDX, IDC_SLIDER_FREQUENCY_PENALTY, m_wndSlider_FrequencyPenalty);
}

BOOL CAIChatSettingDlg::OnInitDialog()
{
	CDialog::OnInitDialog();

	initToolTip();
	initSlider();

	UpdateData(FALSE);
	return TRUE;
}

BOOL CAIChatSettingDlg::PreTranslateMessage(MSG *pMsg)
{
	m_wndToolTip_System_Prompt.RelayEvent(pMsg);
	m_wndToolTip_Max_Tokens.RelayEvent(pMsg);;
	m_wndToolTip_Temperature.RelayEvent(pMsg);;
	m_wndToolTip_Top_P.RelayEvent(pMsg);;
	m_wndToolTip_Top_K.RelayEvent(pMsg);;
	m_wndToolTip_Frequency_penalty.RelayEvent(pMsg);

	return CDialog::PreTranslateMessage(pMsg);
}

void CAIChatSettingDlg::initToolTip()
{
	m_wndToolTip_System_Prompt.Create(this);
	m_wndToolTip_System_Prompt.Activate(true);
	m_wndToolTip_System_Prompt.AddTool(GetDlgItem(IDC_STATIC_SYSTEM_PROMPT)
		, _T("在此描述系统提示"));

	m_wndToolTip_Max_Tokens.Create(this);
	m_wndToolTip_Max_Tokens.Activate(true);
	m_wndToolTip_Max_Tokens.AddTool(GetDlgItem(IDC_STATIC_MAX_TOKENS)
		, _T("单次交互所用的最大 Token 数, 会影响返回结果的长度。普通聊天建议 500-800；短文生成建议 800-2000；代码生成建议 2000-3600；长文生成建议切换模型到 4000 左右"));

	m_wndToolTip_Temperature.Create(this);
	m_wndToolTip_Temperature.Activate(true);
	m_wndToolTip_Temperature.AddTool(GetDlgItem(IDC_STATIC_TEMPERATURE)
		, _T("模型生成文本的随机程度。值越大，回复内容越赋有多样性、创造性、随机性；设为 0 根据事实回答。日常聊天建议设置为 0.7"));

	m_wndToolTip_Top_P.Create(this);
	m_wndToolTip_Top_P.Activate(true);
	m_wndToolTip_Top_P.AddTool(GetDlgItem(IDC_STATIC_TOP_P)
		, _T("指定模型输出的多样性。与温度类似，但更精确。"));

	m_wndToolTip_Top_K.Create(this);
	m_wndToolTip_Top_K.Activate(true);
	m_wndToolTip_Top_K.AddTool(GetDlgItem(IDC_STATIC_TOP_K)
		, _T("从前 k 个 token 中抽样。有助于加快生成过程，并能提高生成文本的质量。"));

	m_wndToolTip_Frequency_penalty.Create(this);
	m_wndToolTip_Frequency_penalty.Activate(true);
	m_wndToolTip_Frequency_penalty.AddTool(GetDlgItem(IDC_STATIC_FREQUENCY_PENALTY)
		, _T("通过惩罚已经频繁使用的单词来减少模型行中重复出现单词的可能性。"));

}

void CAIChatSettingDlg::initSlider()
{
	m_wndSlider_MaxTokens.SetRange(1, 16384);
	m_wndSlider_MaxTokens.ModifyStyle(0, TBS_TOOLTIPS);
	m_wndSlider_MaxTokens.SetPos(m_chatConfig.max_tokens());

	int SCALE_FACTOR = 10;
	m_wndSlider_Temperature.SetRange((int)(0.0 * SCALE_FACTOR), (int)(2.0 * SCALE_FACTOR));
	m_wndSlider_Temperature.ModifyStyle(0, TBS_TOOLTIPS);
	m_wndSlider_Temperature.SetPos((int)(m_chatConfig.temperature() * SCALE_FACTOR));

	m_wndSlider_TopP.SetRange((int)(0.1 * SCALE_FACTOR), (int)(1.0 * SCALE_FACTOR));
	m_wndSlider_TopP.ModifyStyle(0, TBS_TOOLTIPS);
	m_wndSlider_TopP.SetPos((int)(m_chatConfig.top_p() * SCALE_FACTOR));

	m_wndSlider_TopK.SetRange(0, 100);
	m_wndSlider_TopK.ModifyStyle(0, TBS_TOOLTIPS);
	m_wndSlider_TopK.SetPos(m_chatConfig.top_k());

	m_wndSlider_FrequencyPenalty.SetRange(0, 40);
	m_wndSlider_FrequencyPenalty.ModifyStyle(0, TBS_TOOLTIPS);
	m_wndSlider_FrequencyPenalty.SetPos((int)((m_chatConfig.frequency_penalty() + 2.0f) * SCALE_FACTOR));
}


BEGIN_MESSAGE_MAP(CAIChatSettingDlg, CDialog)
	ON_BN_CLICKED(IDOK, &CAIChatSettingDlg::OnBnClickedOk)
	ON_WM_HSCROLL()
	ON_EN_KILLFOCUS(IDC_EDIT_MAX_TOKENS, &CAIChatSettingDlg::OnEnKillfocusEditMaxTokens)
	ON_EN_SETFOCUS(IDC_EDIT_FREQUENCY_PENALTY, &CAIChatSettingDlg::OnEnSetfocusEditFrequencyPenalty)
	ON_EN_KILLFOCUS(IDC_EDIT_TEMPERATURE, &CAIChatSettingDlg::OnEnKillfocusEditTemperature)
	ON_EN_KILLFOCUS(IDC_EDIT_TOP_P, &CAIChatSettingDlg::OnEnKillfocusEditTopP)
	ON_EN_KILLFOCUS(IDC_EDIT_TOP_K, &CAIChatSettingDlg::OnEnKillfocusEditTopK)
END_MESSAGE_MAP()


void CAIChatSettingDlg::OnBnClickedOk()
{
	if (!UpdateData(TRUE)) return;
	writeConfig(m_chatConfig);
	g_aiPaletteSet::instance().refreshAiCloud();
	CDialog::OnOK();
}

void CAIChatSettingDlg::OnHScroll(UINT nSBCode, UINT nPos, CScrollBar *pScrollBar)
{
	if (!pScrollBar) return;

	int SCALE_FACTOR = 10;

	if (pScrollBar->GetDlgCtrlID() == IDC_SLIDER_MAX_TOKENS)
	{
		nPos = m_wndSlider_MaxTokens.GetPos();
		m_chatConfig.set_max_tokens(nPos);
	}
	else if (pScrollBar->GetDlgCtrlID() == IDC_SLIDER_TEMPERATURE)
	{
		nPos = m_wndSlider_Temperature.GetPos();
		m_chatConfig.set_temperature((float)(nPos) / SCALE_FACTOR);
	}
	else if (pScrollBar->GetDlgCtrlID() == IDC_SLIDER_TOP_P)
	{

		nPos = m_wndSlider_TopP.GetPos();
		m_chatConfig.set_top_p((float)(nPos) / SCALE_FACTOR);
	}
	else if (pScrollBar->GetDlgCtrlID() == IDC_SLIDER_TOP_K)
	{

		nPos = m_wndSlider_TopK.GetPos();
		m_chatConfig.set_top_k(nPos);
	}
	else if (pScrollBar->GetDlgCtrlID() == IDC_SLIDER_FREQUENCY_PENALTY)
	{
		nPos = m_wndSlider_FrequencyPenalty.GetPos();
		m_chatConfig.set_frequency_penalty(-2.0f + nPos * 0.1f);
	}
	else
		return;

	UpdateData(FALSE);
	return;
}

void CAIChatSettingDlg::OnEnKillfocusEditMaxTokens()
{
	if (!UpdateData(TRUE)) return;
	m_wndSlider_MaxTokens.SetPos(m_chatConfig.max_tokens());
}

void CAIChatSettingDlg::OnEnSetfocusEditFrequencyPenalty()
{
	if (!UpdateData(TRUE)) return;
	m_wndSlider_FrequencyPenalty.SetPos((int)((m_chatConfig.frequency_penalty() + 2.0f) * 10));
}

void CAIChatSettingDlg::OnEnKillfocusEditTemperature()
{
	if (!UpdateData(TRUE)) return;
	m_wndSlider_Temperature.SetPos((int)(m_chatConfig.temperature() * 10));
}

void CAIChatSettingDlg::OnEnKillfocusEditTopP()
{
	if (!UpdateData(TRUE)) return;
	m_wndSlider_TopP.SetPos((int)(m_chatConfig.top_p() * 10));
}

void CAIChatSettingDlg::OnEnKillfocusEditTopK()
{
	if (!UpdateData(TRUE)) return;
	m_wndSlider_TopK.SetPos(m_chatConfig.top_k());
}

bool CAIChatSettingDlg::getChatConfig(SHAI::ChatConfig &config)
{
	return readConfig(config);
}