local completion_provider = "";
local auto_completion = false;
local max_fim_lines = 10000;
local max_prefix_lines = 500;
local max_suffix_lines = 500;
local max_completion_tokens = 256;
local provider_configs = {};

local FLAGS_MULTI_API = 0x01;
local FLAGS_FETCH_MODEL = 0x02;
local FLAGS_FIXED_MODEL = 0x04;

local log_print = function(...)
	-- print(...);
end;

local request_base = {
	progress = "思考中...",
	timeout = {connect = 1, receive = 10},
	streaming = 16384;
	headers = {
		["Content-Type"] = "application/json";
	};
};

local error_code_map = {
	[-1] = "Network error!";
	[-2] = "Network timeout!";
};

local provider_base = {
	do_fim = function(self, prefix, suffix, content_type, fim_data_handler)
		log_print("FIM impl ", self.fim_api);
		return pinvoke(self[self.fim_api], self, prefix, suffix, content_type, fim_data_handler);
	end;

	shutdown = function(self)
		self._pending_req = nil;
	end;
};

local function _parse_error(response_code, response_data)
	local result = utils:fromjson(response_data);
	if result == nil then
		return {
			["code"] = response_code;
			["error"] = response_code;
			["message"] = response_data;
		};
	end;
	local msg = result.message or (result["error"] and result["error"]["message"]);
	return {
		["code"] = result.code or response_code;
		["error"] = result["type"];
		["message"] = msg;
	};
end;

local function save_env()
	edx.env.ai_assistant = {
		["configs"] = provider_configs;
		["completion"] = {
			["auto_completion"] = auto_completion;
			["provider"] = completion_provider;
			["prefix_lines"] = max_prefix_lines;
			["postfix_lines"] = max_suffix_lines;
			["max_lines"] = max_fim_lines;
			["max_tokens"] = max_completion_tokens;
		}
	};
	edx:save_env();
end;

local function load_env()
	local config = edx.env.ai_assistant or {
		["configs"] = {};
		["completion"] = {
			["auto_completion"] = auto_completion;
			["provider"] = completion_provider;
			["prefix_lines"] = max_prefix_lines;
			["postfix_lines"] = max_suffix_lines;
			["max_lines"] = max_fim_lines;
			["max_tokens"] = max_completion_tokens;
		}
	};
	provider_configs = config["configs"] or {};
	auto_completion = config["completion"]["auto_completion"] == true;
	completion_provider = config["completion"]["provider"] or "";
	max_prefix_lines = math.tointeger(config["completion"]["prefix_lines"] or 500);
	max_suffix_lines = math.tointeger(config["completion"]["postfix_lines"] or 500);
	max_fim_lines = math.tointeger(config["completion"]["max_lines"] or 10000);
	max_completion_tokens = math.tointeger(config["completion"]["max_tokens"] or 256);
end;

local function _make_open_ai_req(ctx, url, fim_data_handler, custom_impl)
	local req = {
		url,
		method = "post",
		on_error = function (self, code, err_code, err_msg, body)
			if ctx._pending_req ~= self then
				return -2;
			end;
			ctx._pending_req = nil;
			err_msg = error_code_map[err_code] or err_msg;
			log_print("[ERR][", err_code, "] - ", code, " ", err_msg);
			log_print("  DATA:\n", body);
			log_print("  URL:", self);
			return fim_data_handler(nil, "error", {
					["code"] = err_code;
					["message"] = err_msg;
			});
		end;
		on_start = function(self, code, content_length)
			if ctx._pending_req ~= self then
				return -2;
			end;
			self.code = code;
			self.content_length = content_length;
			log_print("[Start] - ", code, " ", content_length);
			self._data = {};
		end;
		on_data = function(self, offset, data)
			if ctx._pending_req ~= self then
				return -2;
			end;
			if self.code ~= 200 then
				table.insert(self._data, data);
				return;
			end;
			if #data < 6 then
				return;
			end;
			local line = data:sub(6);
			if line:match("%s*%[DONE%]") then
				return fim_data_handler(nil, true);
			end;
			local result = utils:fromjson(line);
			local choices = result["choices"];
			if choices == nil or choices[1] == nil then
				log_print("  >[", offset, "]", data);
				return;
			end;
			local content;
			if result["object"] == "text_completion" then
				content = choices[1]["text"];
			else
				content = choices[1]["delta"]["content"];
			end;
			return fim_data_handler(content);
		end;
		on_complete = function(self)
			if ctx._pending_req ~= self then
				return -2;
			end;
			ctx._pending_req = nil;
			log_print("\n[END]");
			if self.code ~= 200 then
				local response_body = table.concat(self._data, "");
				local result = _parse_error(self.code, response_body);
				log_print("Request URL:\n", self.url);
				log_print("Request Body:\n", self.body);
				log_print("ERROR!", result);
				log_print("Response Body:\n", response_body);
				return fim_data_handler(nil, "error", result);
			end;
		end;
	};
	return table.clone(req, (custom_impl or {}), request_base);
end;

local function _open_ai_fetch_models(self, response_cb)
	local url = self.url;
	if not url:match(".*/v%d$") then
		url = url .. "/v1";
	end;
	local req = {
		url.."/models",
		method = "get",
		timeout = {connect = 1, receive = 10},
		-- streaming = 16384;
		headers = {
			["Content-Type"] = "application/json";
		};
		on_error = function(self, code, err_code, err_msg, body)
			response_cb(code, nil);
		end;
		on_start = function(self, code, content_length)
			log_print("[Start] - ", code, " ", content_length);
			self.code = code;
			self._data = {};
		end;
		on_data = function(self, offset, data)
			table.insert(self._data, data);
		end;
		on_complete = function(self)
			local data = table.concat(self._data, "");
			if self.code ~= 200 then
				response_cb(self.code, data);
				return;
			end;
			local result = utils:fromjson(data);
			local models = {};
			for idx, m in ipairs(result["data"]) do
				table.insert(models, m.id);
			end;
			response_cb(self.code, models);
		end;
	};
	if type(self.key) == "string" and #self.key > 4 then
		req.headers["Authorization"] = "Bearer "..self.key;
	end;
	pinvoke(fetch, req);
end;

local function _open_ai_completions(self, prefix, suffix, fim_data_handler, api_url)
	local ctx = self;
	local req = _make_open_ai_req(
		ctx,
		self.url..api_url,
		fim_data_handler,
		{
			body = utils:tojson{
				["prompt"] = prefix;
				["suffix"] = suffix;
				["stream"] = true;
				["model"] = self.model;
				["max_tokens"] = max_completion_tokens;
				["temperature"] = 0;
			}
		}
	);
	ctx._pending_req = req;
	if type(self.key) == "string" and #self.key > 4 then
		req.headers["Authorization"] = "Bearer "..self.key;
	end;
	fetch(req);
end;

local function _open_ai_chat_completions(self, prefix, suffix, fim_data_handler, api_url)
	local ctx = self;
	local req = _make_open_ai_req(
		ctx,
		self.url..api_url,
		fim_data_handler,
		{
			body = utils:tojson{
				["messages"] = {
					{
						["role"] = "system";
						["content"] = "You are professional code assistant. \n"..
						"You can help user correcting/editing/completing their code.\n"..
						"Do not generate comments or examples.\n"..
						"Output the completed/modified code only.\n"..
						"";
					},
					{
						["role"] = "user";
						["content"] = prefix.."<fim_middle>"..suffix .. "\n/no_think";
					}
				};
				["stream"] = true;
				["model"] = self.model;
				["max_tokens"] = max_completion_tokens;
				["temperature"] = 0;
			}
		}
	);
	ctx._pending_req = req;
	if type(self.key) == "string" and #self.key > 4 then
		req.headers["Authorization"] = "Bearer "..self.key;
	end;
	fetch(req);
end;

local siliconflow = {
	["type"] = "SiliconFlow";
	["url"] = "https://api.siliconflow.cn/v1";
	["apis"] = {"/completions", "/chat/completions"};
	["models"] = {};
	["flags"] = FLAGS_FETCH_MODEL;
	["/chat/completions"] = function(self, prefix, suffix, content_type, fim_data_handler)
		return _open_ai_chat_completions(self, prefix, suffix, fim_data_handler, "/chat/completions");
	end;
	["/completions"] = function(self, prefix, suffix, content_type, fim_data_handler)
		return _open_ai_completions(self, prefix, suffix, fim_data_handler, "/completions");
	end;
};

local doubao = {
	["type"] = "Doubao";
	["url"] = "https://ark.cn-beijing.volces.com/api/v3";
	["apis"] = {"/chat/completions"};
	["models"] = {};
	["flags"] = 0;
	["/chat/completions"] = function(self, prefix, suffix, content_type, fim_data_handler)
		return _open_ai_chat_completions(self, prefix, suffix, fim_data_handler, "/chat/completions");
	end;
};

local llama_cpp = {
	["type"] = "llama.cpp";
	["url"] = "http://localhost:8012";
	["apis"] = {"/infill", "/chat/completions"};
	["models"] = {"llama.cpp"};
	["flags"] = FLAGS_FIXED_MODEL+FLAGS_MULTI_API;
	["/chat/completions"] = function(self, prefix, suffix, content_type, fim_data_handler)
		return _open_ai_chat_completions(self, prefix, suffix, fim_data_handler, "/v1/chat/completions");
	end;
	["/infill"] = function(self, prefix, suffix, content_type, fim_data_handler)
		local ctx = self;
		local req = _make_open_ai_req(
			ctx,
			self.url.."/infill",
			fim_data_handler,
			{
				on_data = function(self, offset, data)
					if ctx._pending_req ~= self then
						return -2;
					end;
					if self.code ~= 200 then
						table.insert(self._data, data);
						return;
					end;
					if #data < 6 then
						return;
					end;
					local line = data:sub(6);
					if line:match("%s*%[DONE%]") then
						return fim_data_handler(nil, true);
					end;
					local result = utils:fromjson(line);
					local content = result["content"];
					return fim_data_handler(content);
				end;
				body = utils:tojson{
					input_prefix = prefix;
					input_suffix = suffix;
					input_extra = {},
					prompt = "";
					stream = true;
					n_predict = 256;
					top_k = 40;
					top_p= 0.99;
					samplers =  {"top_k", "top_p", "infill"};
					cache_prompt = true;
					t_max_prompt_ms = 1000;
					t_max_predict_ms = 2000;
				}
			}
		);
		ctx._pending_req = req;
		if type(self.key) == "string" and #self.key > 4 then
			req.headers["Authorization"] = "Bearer "..self.key;
		end;
		fetch(req);
	end;
};

local ollama = {
	["type"] = "Ollama";
	["url"] = "http://localhost:11434";
	["apis"] = {"/completions", "/chat/completions"};
	["models"] = {};
	["flags"] = FLAGS_FETCH_MODEL + FLAGS_MULTI_API;
	["/chat/completions"] = function(self, prefix, suffix, content_type, fim_data_handler)
		return _open_ai_chat_completions(self, prefix, suffix, fim_data_handler, "/v1/chat/completions");
	end;
	["/completions"] = function(self, prefix, suffix, content_type, fim_data_handler)
		return _open_ai_completions(self, prefix, suffix, fim_data_handler, "/v1/completions");
	end;
};

local openai = {
	["type"] = "OpenAI";
	["url"] = "https://api.openai.com/v1";
	["apis"] = {"/completions", "/chat/completions"};
	["models"] = {};
	["flags"] = FLAGS_FETCH_MODEL + FLAGS_MULTI_API;
	["/chat/completions"] = function(self, prefix, suffix, content_type, fim_data_handler)
		return _open_ai_chat_completions(self, prefix, suffix, fim_data_handler, "/chat/completions");
	end;
	["/completions"] = function(self, prefix, suffix, content_type, fim_data_handler)
		return _open_ai_completions(self, prefix, suffix, fim_data_handler, "/completions");
	end;
};

local api_providers = {
	llama_cpp;
	ollama;
	siliconflow;
	doubao;
	openai;
};

local function make_context()
	for idx, info in ipairs(provider_configs) do
		if info["name"] == completion_provider then
			for idx, provider in ipairs(api_providers) do
				if provider["type"] == info["type"] then
					return table.clone({}, provider, provider_base, info);
				end;
			end;
		end;
	end;
	return nil;
end;

inplace_suggestion = {
	reinit = function(self)
		local old_ctx = nil;
		if self._provider then
			old_ctx = self._provider;
			self._provider:shutdown();
			self._provider = nil;
		end;

		local ctx = make_context();
		if ctx == nil then
			if old_ctx then
				print("Shutdown AI assistant provider:", old_ctx["name"], "/", old_ctx["type"]);
			end;
			return;
		end;

		self._provider = ctx;
		if old_ctx == nil then
			print("Create AI assistant provider:", ctx["name"], "/", ctx["type"]);
		elseif old_ctx["name"] == ctx["name"] and old_ctx["type"] == ctx["type"]  then
			print("Restart AI assistant provider:", ctx["name"], "/", ctx["type"]);
		else
			print("Shutdown AI assistant provider:", old_ctx["name"], "/", old_ctx["type"]);
			print("Create AI assistant provider:", ctx["name"], "/", ctx["type"]);
		end;
		if self._provider == nil then
			return;
		end;
		return true;
	end;
	init = function(self)
		load_env();
		pinvoke(self.reinit, self);
	end;
	do_fim = function(self, doc, mode)
		if mode == "auto" and not auto_completion then
			return;
		end;
		local provider = self._provider;
		if provider == nil then
			log_print("DO_FIM: no provider!");
			return;
		end;
		if doc.doc_type ~= doc_type.text then
			log_print("DO_FIM: not a text doc!");
			return;
		end;
		if doc.is_selection_valid then
			return;
		end;
		local line_count = doc.line_count;
		if line_count > max_fim_lines then
			return;
		end;
		local doc_name = doc.file_name;
		local doc_id = doc.id;
		local cursor_line = doc.cursor_line;
		local cursor_col = doc.cursor_column;
		local ext_text = doc.get_text(cursor_line, cursor_col, 32768);
		if ext_text ~= "" and not ext_text:match("^[%s`~!@#$%%^&*()_+|=%\\{}%[%]:;\"\'<>,.?/]+$") then
			-- symbol name was found
			return;
		end;
		local begin_line = 0;
		if cursor_line > begin_line + max_prefix_lines then
			begin_line = cursor_line - max_prefix_lines;
		end;
		if line_count > cursor_line + max_suffix_lines then
			line_count = cursor_line + max_suffix_lines;
		end;
		log_print("DO_FIM["..doc_name.."]: ext_text=", ext_text);
		log_print("PREFIX:", begin_line, ",", 0, ",", cursor_line, ",", cursor_col);
		
		local prefix = doc.get_text_range(begin_line, 0, cursor_line, cursor_col);
		log_print("SUFFIX:", cursor_line, ",", cursor_col, ",", line_count, ",", 2048);
		local suffix = doc.get_text_range(cursor_line, cursor_col, line_count, 2048);
		if prefix == "" and suffix == "" then
			log_print("[EMPTY FILE]");
			return;
		end;
		-- log_print("----prefix---");
		-- log_print(prefix);
		-- log_print("----suffix---");
		-- log_print(suffix);
		local should_append_new_line = false;
		local ctx_lines = {
			(doc:get_text_range(cursor_line - 1, 0, cursor_line - 1, 2048) or ""):match("^(.-)%s*$");
			(doc:get_text_range(cursor_line, 0, cursor_line, 2048) or ""):match("^(%s*.-)%s*$");
		};
		if cursor_col == 0 and ctx_lines[2] ~= "" then
			should_append_new_line = true;
			table.insert(ctx_lines, 2, "");
		else
			table.insert(ctx_lines, (doc:get_text_range(cursor_line + 1, 0, cursor_line + 1, 2048) or ""):match("^(%s*.-)%s*$"));
		end;
		log_print(ctx_lines);
		local fim_ctx = {
			["len"] = 0;
			["lines"] = 0;
			["data"] = {};
			["line_idx"] = 0;
			add_fim_content = function(self, text, is_done)
				if doc.id ~= doc_id or doc.file_name ~= doc_name then
					log_print("document changed!");
					return -2;
				end;
				if  cursor_line ~= doc.cursor_line or cursor_col ~= doc.cursor_column then
					log_print("document cursor moved!");
					return -2;
				end;
				if text == "<think>" then
					self._trim_thinking_content = true;
					return;
				end;
				if self._trim_thinking_content then
					if text == "</think>" then
						self._trim_thinking_content = nil;
						self._trim_thinking_gap = true;
					end;
					return;
				end;
				if text ~= nil then
					if self._trim_thinking_gap then
						if string.match(text,"\n+") then
							return;
						end;
						self._trim_thinking_gap = nil;
					end;

					table.insert(self.data, text);
					self.len = self.len + #text;
					if not text:match("\n+") then
						return;
					end;
					text = table.concat(self.data, "");
					if self.lines == 0 and string.match(text,"^%s+$") then
						return;
					end;
					self.lines = self.lines + 1;
				else
					if #self.data == 0 then
						return;
					end;
					is_done = true;
					text = table.concat(self.data, "");
				end;
				-- edx:output(text);

				if self.lines > 0 and self.len/self.lines < 3 then
					-- empty lines
					if not (is_done and self.lines <= 1) then
						return;
					end;
				end;

				if self._trim_markdown_border and is_done then
					if text:sub(1, 3) == "```" then
						-- trim markdown code block tail
						log_print("is_done ", self._prev_line);
						if self._prev_line or self._pending_new_line then
							text = self._prev_line;
							self._trim_markdown_border = false;
							self.line_idx = -9999;
						else
							return;
						end;
					end;
				end;
				local cur_line_idx = self.line_idx;
				local cur_lines = self.lines;
				local cur_lens = self.lens;
				self.line_idx = self.line_idx + self.lines;
				self.data = {};
				self.len = 0;
				self.lines = 0;
				if cur_line_idx == 0 then
					if text:match("^\n*```") then
						-- has markdown code block border
						self._trim_markdown_border = true;
						return;
					end;
				end;

				if cur_line_idx == 0 or (self._trim_markdown_border and cur_line_idx == 1) then
					-- log_print("trim first line", {
					-- 		ctx_lines[1] ~= "";
					-- 		(ctx_lines[2] == "" or string.match(ctx_lines[2], "%s+") ~= nil);
					-- 		ctx_lines[1] == text:sub(1, #ctx_lines[1]);
					-- });
					-- is first line
					if ctx_lines[1] ~= "" and
						(ctx_lines[2] == "" or string.match(ctx_lines[2], "%s+") ~= nil) and
						ctx_lines[1] == text:sub(1, #ctx_lines[1]) then
						-- match the previous line. discard the input
						self.line_idx = cur_line_idx;
						return;
					end;
					local first_line_prefix = ctx_lines[2]:sub(1, cursor_col);
					text = text:gsub("\t", "    ");
					local prefix_matched = false;
					for idx = 1, #first_line_prefix, 1 do
						local s = first_line_prefix:sub(idx);
						s = s:gsub("\t", "    ");
						if text:sub(1, #s) == s then
							-- prefix matched
							text = text:sub(#s + 1);
							prefix_matched = true;
							break;
						end;
					end;
					if not prefix_matched then
						local text_content = text:match("^(.-)%s+$");
						if text_content and first_line_prefix:sub(1, #text_content) == text_content then
							-- prefix matched
							text = text:sub(#text_content + 1);
						end;
					end;
				end;

				if self._trim_markdown_border then
					local prev_line = self._prev_line;
					self._prev_line = text;
					text = prev_line;
					if text == nil then
						return;
					end;
				elseif is_done then
					-- the last line
					if ctx_lines[3] ~= "" and ctx_lines[3] == text:sub(1, #ctx_lines[3]) then
						-- discard the last line which matcheds the next suffix line
						if should_append_new_line and self._pending_new_line then
							text = "";
						else
							return;
						end;
					end;
					local first_line_suffix = ctx_lines[2]:sub(cursor_col + 1);
					if #first_line_suffix > 0 and first_line_suffix == text:sub(1, #first_line_suffix) then
						if self._pending_new_line then
							text = "";
						else
							return;
						end;
					end;
				end;
				
				local fim_blocks = self._fim_blocks;
				local fim_block_id = 0;
				if fim_blocks == nil then
					self._fim_blocks = {};
					fim_blocks = self._fim_blocks;
					fim_block_id = doc:create_marked_text(cursor_line, cursor_col, cursor_line, cursor_col, 0);
					table.insert(fim_blocks, fim_block_id);
					doc:marked_text_set_tag(fim_block_id, 1);
				else
					fim_block_id = fim_blocks[#fim_blocks];
				end;
				local next_pending_new_line = nil;
				log_print("Text[", text, "]", self._pending_new_line);
				if text ~= "" and text:sub(#text) == "\n" then
					text = text:match("^(.-)%s*$");
					next_pending_new_line = true;
				end;
				local ret = 0
				if self._pending_new_line then
					ret = doc:marked_text_append_text(fim_block_id, "\n"..text);
				else
					ret = doc:marked_text_append_text(fim_block_id, text);
				end;
				self._pending_new_line = next_pending_new_line;
				if ret < 0 then
					log_print("Cancel FIM operation!");
					return -2;
				end;
				return 0;
			end;
		};
		local content_type = content_type_name[doc.content_type];
		if doc_name == nil then
			doc_name = "noname";
		end;
		if content_type == nil then
			content_type = "text";
		end;
		prefix = "file name: "..doc_name.."\n```"..content_type.."\n"..prefix;
		suffix = suffix .. "\n```\n";
		provider:do_fim(prefix, suffix, content_type, function(text, mode, info)
			if mode == "error" then
				print("FIM error! [", info["code"], "] - ", info["message"]);
				return -2;
			end;
			return pinvoke(fim_ctx.add_fim_content, fim_ctx, text);
		end);
	end;
};

function ai_assistant:on_init(page)
	edx:reload_env();
	load_env();
	local providers = {};
	for idx, info in ipairs(api_providers) do
		table.insert(providers, {
				["type"] = info["type"];
				["url"] = info["url"];
				["apis"] = info["apis"];
				["models"] = info["models"];
				["flags"] = info["flags"];
		});
	end;
	self:init{
		["providers"] = providers;
		["configs"] = provider_configs;
		["completion"] = {
			["auto_completion"] = auto_completion;
			["provider"] = completion_provider;
			["prefix_lines"] = math.tointeger(max_prefix_lines);
			["postfix_lines"] = math.tointeger(max_suffix_lines);
			["max_lines"] = math.tointeger(max_fim_lines);
			["max_tokens"] = math.tointeger(max_completion_tokens);
		}
	};
end;

local function init_event_handler(event)
	ai_assistant[event] = function(self, page)
		local func = self[event.. "_"..page];
		if func == nil then
			return 0;
		end;
		return pinvoke(func, self);
	end;
end;

init_event_handler("on_dismissed");
init_event_handler("on_apply");
init_event_handler("on_cancel");

function ai_assistant:on_dismissed_providers()
	-- print("on_dismissed_providers", self.provider_configs);
end;

function ai_assistant:on_apply_config()
	local config = self.completion_config;
	auto_completion = config["auto_completion"];
	max_prefix_lines = math.tointeger(config["prefix_lines"]);
	max_suffix_lines = math.tointeger(config["postfix_lines"]);
	max_fim_lines = math.tointeger(config["max_lines"]);
	max_completion_tokens = math.tointeger(config["max_tokens"]);
	completion_provider = config["provider"];
	save_env();
	inplace_suggestion:reinit();
end;

function ai_assistant:on_apply_providers()
	local old_ctx = make_context();
	provider_configs = self.provider_configs;
	local provider_list = {};
	for idx, info in ipairs(provider_configs) do
		table.insert(provider_list, info["name"]);
	end;
	self:update_provider_list(provider_list);
	save_env();
	local new_ctx = make_context();
	if new_ctx == nil or old_ctx == nil then
		inplace_suggestion:reinit();
		return;
	end;
	for key, val in pairs(new_ctx) do
		if type(val) == "string" then
			if old_ctx[key] ~= val then
				-- provider changed
				inplace_suggestion:reinit();
				break;
			end;
		end;
	end;
end;

function ai_assistant:on_cancel_providers()
	local providers = {};
	for idx, info in ipairs(api_providers) do
		table.insert(providers, {
				["type"] = info["type"];
				["url"] = info["url"];
				["apis"] = info["apis"];
				["models"] = info["models"];
				["flags"] = info["flags"];
		});
	end;
	self:init{
		["providers"] = providers;
		["configs"] = provider_configs;
	};
end;

function ai_assistant:on_fetch_models(info)
	local model = table.clone({}, info, {fetch_models = _open_ai_fetch_models});
	model:fetch_models(
		function(code, models)
			if code ~= 200 then
				-- error!
				ai_assistant:popup_tip(LANG("notify/error/ai_assistant/fetch_model_error").."\n"..models, 3000);
				return;
			end;
			ai_assistant:update_provider_models(info["type"], models);
		end
	);
end;
