#!/usr/bin/env python3

import os
import argparse
from pathlib import Path
from collections.abc import Callable

from prompt_toolkit import Application
from prompt_toolkit.layout import Layout, HSplit, FloatContainer, Float, Dimension, WindowAlign
from prompt_toolkit.layout.menus import CompletionsMenu
from prompt_toolkit.widgets import TextArea, Frame
from prompt_toolkit.completion import PathCompleter
from prompt_toolkit.document import Document
from prompt_toolkit.lexers import Lexer
from prompt_toolkit.formatted_text import StyleAndTextTuples
from prompt_toolkit.key_binding import KeyPressEvent
from prompt_toolkit.layout.controls import FormattedTextControl
from prompt_toolkit.layout.containers import Window
from prompt_toolkit.styles import Style
from prompt_toolkit.key_binding import KeyBindings

from . import __doc__, __version__
from .utils import CallAction


def print_version_and_exit() -> None:
	'''show the version and exit'''
	print("symlink-editor %s" % __version__)
	exit(0)


# ---------- Coloring logic ----------

def colorize_path(text: str) -> StyleAndTextTuples:
	"""Return green for existing path prefix, red for non-existing suffix."""
	if not text:
		return [("class:nonexistent", "")]

	i = index_broken_part(text)

	existing = text[:i]
	broken = text[i:]

	out: 'StyleAndTextTuples' = []
	if existing:
		out.append(("class:existing", existing))
	if broken:
		out.append(("class:nonexistent", broken))

	return out

def index_broken_part(target: str) -> int:
	parts = Path(target).parts
	existing_len = 0
	for i in range(len(parts)):
		test = Path(*parts[: i + 1])
		if test.exists():
			existing_len = len(os.path.join(*parts[: i + 1]))
		else:
			break

	# Include trailing slash if it’s part of the text
	if target[existing_len:existing_len+1] == os.sep:
		existing_len += 1

	return existing_len



class PathLexer(Lexer):
	"""Color existing/non-existing path parts."""
	def lex_document(self, document: Document) -> 'Callable[[int], StyleAndTextTuples]':
		text = document.text

		def get_line(lineno: int) -> StyleAndTextTuples:
			return colorize_path(text)

		return get_line


# ---------- Main ----------

def main(largs: 'list[str]|None' = None) -> None:
	p = argparse.ArgumentParser(description=__doc__)
	p.add_argument('-v', '--version', action=CallAction, callback=print_version_and_exit)
	p.add_argument('-w', '--window', action='store_true', help="Give feedback in a window. This is useful if you want to integrate this into other TUI programs like yazi or ranger.")
	p.add_argument("link", help="path to the symlink to edit")
	args = p.parse_args(largs)

	if args.window:
		def print_or_display_in_window(msg: str) -> None:
			help_text = "Press [Enter] or [Ctrl+C] to quit"
			root_container = HSplit([
				Window(height=Dimension(weight=1), always_hide_cursor=True),
				Window(FormattedTextControl([("class:msg", msg)]), height=1, align=WindowAlign.CENTER),
				#Window(height=1),
				Window(FormattedTextControl([("class:help", help_text)]), height=1, align=WindowAlign.CENTER),
				Window(height=Dimension(weight=1)),
			])
			style = Style.from_dict({
				"help": "fg:#888888",
			})

			kb = KeyBindings()
			@kb.add("enter")
			@kb.add("escape")
			@kb.add("c-c")
			@kb.add("q")
			def _(event: KeyPressEvent) -> None:
				"""Exit."""
				event.app.exit()

			Application(
				layout = Layout(root_container),
				key_bindings = kb,
				style = style,
				full_screen = True,
			).run()
	else:
		def print_or_display_in_window(msg: str) -> None:
			print(msg)

	link_path = Path(args.link).absolute()
	os.chdir(link_path.parent)
	if not link_path.is_symlink():
		print_or_display_in_window(f"Error: {link_path} is not a symlink.")
		exit(1)

	target = os.readlink(link_path)

	# --- UI setup ---
	header = Window(
		content=FormattedTextControl(
			text=[("class:header", f"Link: {link_path}\n")]
		),
		height=1,
	)

	footer = Window(
		content=FormattedTextControl(
			text=[("class:footer", "Press [Enter] to accept or save, [Ctrl+C] to cancel")]
		),
		height=1,
	)

	textarea = TextArea(
		text=target,
		lexer=PathLexer(),
		completer=PathCompleter(),
		multiline=False,
		focus_on_click=True,
	)

	# Position cursor at start of non-existing part
	textarea.buffer.cursor_position = index_broken_part(target)

	frame = Frame(
		body=textarea,
		title="Target",
		style="class:frame",
	)

	def update_frame_style() -> None:
		text = textarea.text
		if text and Path(text).exists():
			style_dict['frame.border'] = style_dict['existing']
			style_dict['frame.label'] = style_dict['existing']
		else:
			style_dict['frame.border'] = style_dict['nonexistent']
			style_dict['frame.label'] = style_dict['nonexistent']

		app.style = Style.from_dict(style_dict)

	def open_menu() -> None:
		textarea.buffer.start_completion()

	textarea.buffer.on_text_changed.add_handler(lambda _: update_frame_style())
	textarea.buffer.on_text_changed.add_handler(lambda _: open_menu())

	# --- Layout ---

	float_container = FloatContainer(
		content = HSplit([frame, Window(height=Dimension(weight=1))]),
		floats = [
			Float(
				xcursor = True,  # position the menu where the cursor is instead of centering it
				ycursor = True,  # position the menu where the cursor is instead of centering it
				content = CompletionsMenu(
					scroll_offset = 1,  # scroll if the cursor gets closer than this number of lines to the edge
				),
			)
		],
	)

	root_container = HSplit([
		header,
		float_container,
		Window(height=1),
		footer,
	])

	style_dict = {
		"header": "fg:#888888",
		"footer": "fg:#888888",
		"existing": "fg:green",
		"nonexistent": "fg:red",
		"frame.border": "fg:ansiblue",
		"frame.label": "fg:ansiblue bold",
	}

	# --- Key bindings ---
	kb = KeyBindings()

	@kb.add("enter")
	def _(event: KeyPressEvent) -> None:
		"""Save and exit."""
		buff = event.app.current_buffer
		if buff.complete_state:
			# The completions menu is open
			if os.path.isdir(textarea.text):
				# Enter the currently selected directory
				buff.insert_text(os.path.sep)
			elif buff.complete_state.current_completion:
				buff.apply_completion(buff.complete_state.current_completion)
			return

		new_target = textarea.text
		event.app.exit(result=new_target)

	@kb.add("escape")
	@kb.add("c-c")
	def _(event: KeyPressEvent) -> None:
		"""Cancel without saving."""
		event.app.exit(result=None)

	# --- Application ---
	app: 'Application[str|None]' = Application(
		layout=Layout(root_container, focused_element=textarea),
		key_bindings=kb,
		style=Style.from_dict(style_dict),
		full_screen=True,
	)

	update_frame_style()
	new_target = app.run()

	# --- After exit ---
	if new_target is None:
		print_or_display_in_window("Cancelled.")
		return

	try:
		# Write to a tmp link first to not corrupt the real link in case something goes wrong.
		tmp = link_path.parent / (link_path.name + ".tmp_symlink")
		if tmp.exists() or tmp.is_symlink():
			tmp.unlink()
		os.symlink(new_target, tmp)
		tmp.replace(link_path)
		print_or_display_in_window(f"Updated {link_path} → {new_target}")
	except OSError as e:
		print_or_display_in_window(f"Error updating symlink: {e}")
		exit(1)


if __name__ == '__main__':
	main()
