local M = {}

local llm_server = "http://localhost:8080"
local llm_token = ""
local llm_buf = nil
local llm_messages = {}
local llm_response = ""
local llm_stream = nil
local llm_ctx = 0
local llm_params = {
	cache_prompt = true,
	temperature = 1.0,
	top_p = 0.95,
	top_k = 64,
	min_p = 0.0,
	presence_penalty = 0.0,
	repeat_penalty = 1.0,
	repeat_last_n = 256,
	n_predict = 2048,
	stop = {"<|im_start|>", "<|im_end|>", "<|endoftext|>"},
	chat_template_kwargs = {enable_thinking = false},
	stream = true
}

function M:setup(server, token)
	llm_server = server
	llm_token = token
end

local function symbols_to_string(symbols, indent)
	indent = indent or ""
	local lines = {}
	local kind_map = {
		[1] = "File", [2] = "Module", [5] = "Class",
		[6] = "Method", [12] = "Function", [13] = "Variable"
	}
	for _, symbol in ipairs(symbols) do
		local kind = kind_map[symbol.kind] or "Symbol"
		local range = symbol.range or (symbol.location and symbol.location.range)
		if range then
			local line = range.start.line
			table.insert(lines, string.format("%s[%s] %s %s (Line: %d)", indent, kind, symbol.name, symbol.detail, line+1))
		else
			table.insert(lines, string.format("%s[%s] %s %s", indent, kind, symbol.name, symbol.detail))
		end
		if symbol.children and #symbol.children > 0 then
			local children_str = symbols_to_string(symbol.children, indent .. "  ")
			table.insert(lines, children_str)
		end
	end
	return table.concat(lines, "\n") .. "\n"
end

local function get_file_outline()
	local status, params = pcall(vim.lsp.util.make_position_params(0, vim.bo.fileencoding))
	if not status then
		params= {
			textDocument = {
				uri = vim.uri_from_bufnr(0)
			}
		}
	end
	local clients = vim.lsp.get_clients({ bufnr = 0 })
	local client = clients[1]
	if not client then return "" end
	local response = client.request_sync("textDocument/documentSymbol", params, 1000, 0)
	if response and response.result then
		return "\n# File outline\n" .. symbols_to_string(response.result)
	end
	return ""
end

local function get_lsp_context()
	local context = ""
	context = context .. get_file_outline()
	local diagnostics = vim.diagnostic.get(0)
	if #diagnostics > 0 then
		context = context .. "\n# Errors\n"
		for _, diag in ipairs(diagnostics) do
			context = context .. string.format("Error at line %d: %s\n", diag.lnum +1, diag.message)
		end
	else
		context = context .. "\n# No errors found\n"
	end
	context = context .. "\n# User\n"
	return context
end

local function set_status(status)
	vim.b[llm_buf].llm_status = status
	vim.schedule(function()
		vim.cmd('redrawstatus!')
	end)
end

local function setup_chat()
	if llm_buf == nil then
		llm_buf = vim.api.nvim_create_buf(false, true)
		local win_opts = {split = "above", height = 20}
		local win = vim.api.nvim_open_win(llm_buf, false, win_opts)
		vim.bo[llm_buf].filetype =  "markdown"
		vim.cmd("syntax on")
		vim.bo[llm_buf].buftype = "nofile"
		vim.bo[llm_buf].bufhidden = "hide"
		vim.wo[win].scrolloff = 0
		vim.wo[win].statusline = " %{get(b:,'llm_status','Idle')} "
	end
	if #llm_messages == 0 then
		llm_messages = {{role = "system", content = "answer concisely"}}
		llm_response = ""
	end
	if llm_ctx == 0 then
		local cmd = "curl -s " .. llm_server .. "/props"
		local res = vim.fn.system(cmd)
		local success, data = pcall(vim.fn.json_decode, res)
		if success and data then
			llm_ctx = tonumber(data.default_generation_settings.n_ctx or 1)
		end
	end
end

local function fillBar(current, total)
	if type(current) ~= "number" then
		return "current is not a number!"
	elseif type(total) ~= "number" then
		return "total is not a number!"
	end
	local width = 10
	local percentage = math.floor((current * 100) / total)
	local filled = math.floor((current * width) / total)
	filled = math.min(filled, width)
	local bar = "[" .. string.rep("#", filled) .. string.rep("-", width - filled) .. "]"
	return bar .. " " .. percentage .. "%"
end

local function scroll_to_bottom()
	local win_ids = vim.fn.win_findbuf(llm_buf)
	local last_line = vim.api.nvim_buf_line_count(llm_buf)
	for _, win_id in ipairs(win_ids) do
		vim.api.nvim_win_set_cursor(win_id, {last_line, 0})
		vim.api.nvim_win_call(win_id, function()
			vim.cmd("normal! zt")
		end)
	end
end

local function append_to_llm_buffer(content)
	local parts = vim.split(content, "\n", { trimempty = false })
	for i, part in ipairs(parts) do
		if i == 1 then
			local last_line_idx = vim.api.nvim_buf_line_count(llm_buf) - 1
			local last_line_text = vim.api.nvim_buf_get_lines(llm_buf, last_line_idx, -1, false)[1] or ""
			vim.api.nvim_buf_set_lines(llm_buf, last_line_idx, -1, false, { last_line_text .. part })
		else
			vim.api.nvim_buf_set_lines(llm_buf, -1, -1, false, { part })
		end
	end
end

local function append_chunk_to_llm_buffer(raw_chunk)
	if not raw_chunk or raw_chunk == "" then return end
	for line in raw_chunk:gmatch("[^\r\n]+") do
		local json_str = line:gsub("^data: ", "")
		if json_str ~= "" and json_str ~= "[DONE]" then
			local ok, decoded = pcall(vim.fn.json_decode, json_str)
			if ok and decoded.choices then
				if decoded.choices[1].delta.content and type(decoded.choices[1].delta.content) == "string" then
					local last_line_idx = vim.api.nvim_buf_line_count(llm_buf) - 1
					local last_line_text = vim.api.nvim_buf_get_lines(llm_buf, last_line_idx, -1, false)[1] or ""
					local content = decoded.choices[1].delta.content
					append_to_llm_buffer(content)
					llm_response = llm_response .. content
				elseif decoded.choices[1].finish_reason == "stop" then
					local used = tonumber(decoded.timings.prompt_n + decoded.timings.predicted_n + decoded.timings.cache_n)
					local bar = fillBar(used, llm_ctx)
					set_status(decoded.model .. " (" .. decoded.timings.predicted_ms .. " ms) " .. bar)
				elseif decoded.choices[1].finish_reason == "length" then
					set_status("Aborted")
				end
			elseif decoded.error then
				vim.schedule(function()
					vim.notify(decoded.error.message, vim.log.levels.ERROR)
				end)
			end
		end
	end
end

local function append_chunk_to_current_buffer(raw_chunk)
	if not raw_chunk or raw_chunk == "" then return end
	for line in raw_chunk:gmatch("[^\r\n]+") do
		local json_str = line:gsub("^data: ", "")
		if json_str ~= "" and json_str ~= "[DONE]" then
			local ok, decoded = pcall(vim.fn.json_decode, json_str)
			if ok and decoded.choices then
				if decoded.choices[1].delta.content and type(decoded.choices[1].delta.content) == "string" then
					local content = decoded.choices[1].delta.content
					local parts = vim.split(content, "\n", { trimempty = false })
					vim.api.nvim_put(parts, "c", true, true)
				end
			elseif decoded.error then
				vim.schedule(function()
					vim.notify(decoded.error.message, vim.log.levels.ERROR)
				end)
			end
		end
	end
end

function M:abort()
	if llm_stream then
		llm_stream:kill(15)
		llm_stream = nil
	end
end

function M:infill()
	M:abort()

	local endpoint = llm_server .. '/v1/chat/completions'
	local bufnr = vim.api.nvim_get_current_buf()
	local cursor = vim.api.nvim_win_get_cursor(0)
	local row = cursor[1] - 1
	local col = cursor[2] + 1

	local system = "You are a " .. vim.bo.filetype .. " completion engine. Return ONLY the raw code that fits between the prefix and suffix."

	local prefix_lines = vim.api.nvim_buf_get_lines(bufnr, 0, row, false)
	local current_line = vim.api.nvim_buf_get_lines(bufnr, row, row + 1, false)[1]
	local prefix = table.concat(prefix_lines, "\n")

	local suffix_lines = vim.api.nvim_buf_get_lines(bufnr, row + 1, -1, false)
	local suffix_tail = current_line:sub(col + 1)
	local suffix = suffix_tail .. "\n" .. table.concat(suffix_lines, "\n")

	local params = vim.deepcopy(llm_params)
	-- infill does not preserve history. Instead, the current file is provided as context
	params["messages"] = {
		{role = "system", content = system},
		{role = "user", content = "<|fim_prefix|>" .. prefix .. "<|fim_suffix|>" .. suffix .. "<|fim_middle|>"},
		{role = "assistant", content = current_line:sub(1, col)}
	}
	--params["chat_template_kwargs"] = {enable_thinking = true}
	params["stop"] = { "<|file_separator|>", "<|endoftext|>"}
	local data = vim.fn.json_encode(params)
	local cmd = {
		"curl",
		"-N",
		"-s",
		"-X", "POST",
		endpoint,
		"-H", "Content-Type: application/json",
		"-H", "Authorization: " .. llm_token,
		"-d", data
	}
	llm_stream = vim.system(cmd, {
		stdout = function(err, chunk)
			if chunk then
				vim.schedule(function()
					append_chunk_to_current_buffer(chunk)
				end)
			end
		end,
	}, function(complete)
		llm_stream = nil
	end)
end

local function prompt(prompt)
	setup_chat()

	set_status("Thinking...")
	local endpoint = llm_server .. '/v1/chat/completions'

	if #llm_response ~= 0 then
		-- Append previous response to chat history
		table.insert(llm_messages, {role = "assistant", content = llm_response})
	end

	prompt = get_lsp_context() .. prompt
	table.insert(llm_messages, {role = "user", content = prompt})
	append_to_llm_buffer("\n" .. prompt .. "\n\n# LLM\n")
	local params = vim.deepcopy(llm_params)
	params["messages"] = llm_messages
	local data = vim.fn.json_encode(params)

	llm_response = ""
	local cmd = {
		"curl",
		"-N",
		"-s",
		"-X", "POST",
		endpoint,
		"-H", "Content-Type: application/json",
		"-H", "Authorization: " .. llm_token,
		"-d", data
	}
	scroll_to_bottom()
	vim.system(cmd, {
		stdout = function(err, chunk)
			if chunk then
				vim.schedule(function()
					append_chunk_to_llm_buffer(chunk)
				end)
			end
		end,
	})
end

function M:visual()
	local start_line = vim.fn.line("'<")
	local end_line = vim.fn.line("'>")
	if end_line < start_line then
		start_line, end_line = end_line, start_line
	end

	local code_lines = vim.fn.getline(start_line, end_line)
	local user_input = vim.fn.input("> ")
	if user_input == "" then
		user_input = "Is this code correct? If yes, what does it do?"
	end
	local filename = vim.fn.fnamemodify(vim.api.nvim_buf_get_name(0), ":t")
	context = string.format("File: %s, Location: lines %d to %d", filename, start_line, end_line)
	local code_block = context .. "\n```" .. vim.bo.filetype .. "\n" .. table.concat(code_lines, "\n") .. "\n```"
	local final_text = user_input .. "\n\n" .. code_block
	prompt(final_text)
end

function M:input()
	local input = vim.fn.input("> ")
	if #input < 1 then
		return
	end
	prompt(input)
end

return M