extends Control

var _llama: GodotLlama
var _streaming_active := false
var _model_name := "Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf"
var _model_path
var _pending_user_text := ""
var _chat_history: Array[Dictionary] = []
var _saved_state_blob := PackedByteArray()

const DEFAULT_SYSTEM_PROMPT := "CHATBOT_SYSTEM_PROMPT"
const DEFAULT_MAX_HISTORY_TURNS := 0


var DEFAULT_N_CTX := 40096
const DEFAULT_THREADS := 4
const temperature:= 0.1
const top_p:= 0.95
const top_k:= 20
const min_p:= 0.05
const repeat_penalty:= 1.15
const frequency_penalty:= 0.1
const presence_penalty:= 0.0
const penalty_last_n:= 64
const max_tokens := 150


@onready var _prompt_edit: TextEdit = $MarginContainer/Scroll/VBox/PromptEdit
@onready var _generate_button: Button = $MarginContainer/Scroll/VBox/GenerateRow/GenerateButton
@onready var _output_text: RichTextLabel = $MarginContainer/Scroll/VBox/OutputText
@onready var _status_label: Label = $MarginContainer/Scroll/VBox/StatusLabel
@onready var _close_button: Button = $MarginContainer/Scroll/VBox/GenerateRow2/CloseButton
@onready var main: Control = $"."
@onready var option_button: OptionButton = $MarginContainer/Scroll/VBox/OptionButton
@onready var option_label: Label = $MarginContainer/Scroll/VBox/Label
@onready var helper = $VBoxContainer/SubViewportContainer/SubViewport/base_node/helper

var document_id

signal closing_chatbot


func _ready() -> void:
	helper.play_idle()
	print('local chatbot ready')
	_generate_button.pressed.connect(_on_generate_pressed)
	_close_button.pressed.connect(_on_close_pressed)
	_prompt_edit.connect("input_event", _on_prompt_edit_gui_input)
	for i in range(option_button.item_count):
		option_button.set_item_text(i, tr(option_button.get_item_text(i)))
	_llama = GodotLlama.new()
	_llama.context.token_generated.connect(_on_token_generated)
	_llama.context.generation_finished.connect(_on_generation_finished)
	_llama.context.generation_error.connect(_on_generation_error)
	document_id = option_button.selected

	print(DEFAULT_N_CTX)
	_set_status("ready")
	var exe_path = OS.get_executable_path()
	var exe_dir = exe_path.get_base_dir()
	_model_path = exe_dir.path_join("model/%s" % _model_name)
	_on_load_model_pressed()
	

func _unhandled_input(event: InputEvent) -> void:
	if not main.visible: return # Ignorieren wenn zu
	if event.is_action_pressed("drop"):
		get_viewport().set_input_as_handled() # Event stoppen
		_on_close_pressed()
		
	if event.is_action_pressed("advance_dialogue"):
		get_viewport().set_input_as_handled() # Event stoppen
		_on_generate_pressed()

func _on_TextEdit_gui_input(event):
	if event is InputEventKey and event.pressed:
		if event.keycode == KEY_ENTER and not event.shift_pressed:
			_on_generate_pressed()
			accept_event()  # verhindert neue Zeile

func _on_close_pressed() -> void:
	print("close clicked")
	main.visible = false
	Input.set_mouse_mode(Input.MOUSE_MODE_CAPTURED)
	Global.chatbot_opened = false
	TextToSpeech._stop_dialogue()


func _on_load_model_pressed() -> void:
	if _model_path.is_empty():
		_set_status("select a model first")
		return

	_set_status("loading model...")
	var err := _llama.load_model(_model_path)
	if err != OK:
		_set_status("model load failed: %s" % error_string(err))
		return

	_set_status("model loaded, creating context...")
	_create_context_with_defaults()

func _create_context_with_defaults() -> void:
	var params := {
		"n_ctx": DEFAULT_N_CTX,
		"threads": DEFAULT_THREADS,
		"threads_batch": DEFAULT_THREADS,
	}
	var err := _llama.create_context(params)
	if err != OK:
		_set_status("create context failed: %s" % error_string(err))
		return
	_set_status("ready")


func _on_generate_pressed() -> void:
	$VBoxContainer/SubViewportContainer/SubViewport/AnimationPlayer.play("walk_helper")
	helper.play_think_walk()
	print("generating...")
	if not _llama.context.is_initialized():
		_set_status("create context first")
		return
	document_id = option_button.selected
	var user_text := _prompt_edit.text.strip_edges()
	if user_text.is_empty():
		_set_status("prompt is empty")
		return

	# The demo rebuilds the full conversation prompt each turn, so reset context first.
	_llama.context.reset()

	var prompt := _build_chat_prompt(user_text)
	var max_tokens = max_tokens
	var params := {
		"temperature": temperature,
		"top_p": top_p,
		"top_k": top_k,
		"min_p": min_p,
		"repeat_penalty": repeat_penalty,
		"frequency_penalty": frequency_penalty,
		"presence_penalty": presence_penalty,
		"penalty_last_n": penalty_last_n,
	}
	var stop_sequences := _parse_stop_sequences("<|eot_id|>")
	if not stop_sequences.is_empty():
		params["stop_sequences"] = stop_sequences

	_output_text.clear()
	_pending_user_text = user_text
	print("real")
	_set_status("generating...")
	var result := _llama.generate(prompt, max_tokens, params)
	var clean_result := _clean_assistant_output(result)
	_output_text.text = clean_result
	var line = {"text": clean_result}
	TextToSpeech._on_dialogue(line)
	#_append_chat_turn("user", user_text)
	#_append_chat_turn("assistant", clean_result)
	_set_status("generation finished")

func _on_clear_memory_pressed() -> void:
	_chat_history.clear()
	_pending_user_text = ""
	_set_status("memory cleared")

func _on_clear_kv_pressed() -> void:
	if not _llama.context.is_initialized():
		_set_status("context not initialized")
		return
	_llama.context.clear_kv_cache()
	_set_status("kv cache cleared")

func _on_save_state_memory_pressed() -> void:
	if not _llama.context.is_initialized():
		_set_status("context not initialized")
		return
	_saved_state_blob = _llama.context.save_state()
	if _saved_state_blob.is_empty():
		_set_status("save state blob failed")
		return
	_set_status("state blob saved (%d bytes)" % _saved_state_blob.size())

func _on_load_state_memory_pressed() -> void:
	if not _llama.context.is_initialized():
		_set_status("context not initialized")
		return
	if _saved_state_blob.is_empty():
		_set_status("no saved state blob")
		return
	var err := _llama.context.load_state(_saved_state_blob)
	if err != OK:
		_set_status("load state blob failed: %s" % error_string(err))
		return
	_set_status("state blob loaded")

func _on_token_generated(token_text: String, _token_id: int) -> void:
	if _streaming_active:
		_output_text.text += token_text

func _on_generation_finished(_full_text: String) -> void:
	$VBoxContainer/SubViewportContainer/SubViewport/AnimationPlayer.stop()
	helper.play_idle()
	if _streaming_active:
		var clean_text := _clean_assistant_output(_full_text)
		_output_text.text = clean_text
		#_append_chat_turn("user", _pending_user_text)
		#_append_chat_turn("assistant", clean_text)
		_streaming_active = false
	_set_status("generation finished")

func _on_generation_error(message: String) -> void:
	_streaming_active = false
	_set_status("generation error: %s" % message)

func _set_status(message: String) -> void:
	_status_label.text = "Status: %s" % message

func _load_txt(path: String) -> String:
	var file := FileAccess.open(path, FileAccess.READ)

	if file == null:
		_set_status("file open failed: %s" % error_string(FileAccess.get_open_error()))
		return ""
	var content := file.get_as_text()
	file.close()
	return content

func _build_chat_prompt(user_text: String) -> String:
	var parts := PackedStringArray()
	var clean_system := tr(DEFAULT_SYSTEM_PROMPT).strip_edges()

	# Llama 3 Header
	parts.append("<|begin_of_text|>")

	# system prompt
	var system_combined := ""
	if not clean_system.is_empty():
		system_combined += clean_system
	if not system_combined.is_empty():
		parts.append("<|start_header_id|>system<|end_header_id|>\n\n")
		parts.append(system_combined)
		parts.append("<|eot_id|>")
	var document
	document = _load_txt("res://addons/godot_llama/demo/documents/%s.md" % [document_id]).strip_edges()

	# user-message with document
	parts.append("<|start_header_id|>user<|end_header_id|>\n\n")
	if not document.is_empty():
		parts.append("Dokument:\n%s\n\nFrage: " % document)
	parts.append(user_text.strip_edges())
	parts.append("<|eot_id|>")
	parts.append("<|start_header_id|>assistant<|end_header_id|>\n\n")
	
	print("".join(parts))
	
	return "".join(parts)
	
func _append_chat_turn(role: String, content: String) -> void:
	var clean := content.strip_edges()
	if clean.is_empty():
		return
	_chat_history.append({
		"role": role,
		"content": clean,
	})
	_chat_history.clear()
	return

func _clean_assistant_output(text: String) -> String:
	var out := text
	var end_idx := out.find("<|")
	if end_idx >= 0:
		out = out.substr(0, end_idx)
	out = out.replace("<|start_header_id|>assistant<|end_header_id|>", "")
	out = out.replace("<|begin_of_text|>", "")
	out = out.replace("<|end_of_text|>", "")
	return out.strip_edges()

func _parse_stop_sequences(raw: String) -> PackedStringArray:
	var stops := PackedStringArray()
	var parts := raw.split(",")
	for part in parts:
		var clean := String(part).strip_edges()
		if not clean.is_empty():
			stops.append(clean)
	return stops


func _on_close_button_pressed() -> void:
	main.visible = false
	Input.set_mouse_mode(Input.MOUSE_MODE_CAPTURED)
	TextToSpeech._stop_dialogue()
	Global.chatbot_opened = false
	TextToSpeech._stop_dialogue()
	closing_chatbot.emit()
	
func _on_prompt_edit_gui_input(event: InputEvent) -> void:
	if event is InputEventKey and event.pressed:
		print(event.keycode)
		if event.keycode == KEY_ENTER:
			if event.shift_pressed:
				_prompt_edit.insert_text_at_caret("\n")
			else: 
				_on_generate_pressed()
			accept_event()