"""Collapsible staged / unstaged / untracked file list with keyboard control."""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from pathlib import Path
from typing import Literal

from pygit2 import Repository
from pyratatui import (
	Constraint,
	Direction,
	Frame,
	Layout,
	List,
	ListItem,
	ListState,
	Rect,
	Style,
	TextPrompt,
	TextState,
)

from pygittools.tui.commit_draft import save_commit_draft
from pygittools.tui.git_colors import current_theme, focus_style, row_style_for_section
from pygittools.tui.keys import (
	KeyLike,
	is_backspace,
	is_down,
	is_enter,
	is_left,
	is_right,
	is_space,
	is_up,
	printable_char,
)
from pygittools.tui.worktree import (
	categorize_status,
	commit_staged,
	stage_path,
	stage_paths,
	unstage_path,
	unstage_paths,
)


class SectionId(Enum):
	STAGED = "staged"
	UNSTAGED = "unstaged"
	UNTRACKED = "untracked"


@dataclass(frozen=True, slots=True)
class _Row:
	kind: Literal["commit", "header", "file"]
	section: SectionId | None
	path: str | None
	label: str


@dataclass(frozen=True, slots=True)
class NavKeyResult:
	enter_input: bool = False
	layout_changed: bool = False
	committed: bool = False


@dataclass(frozen=True, slots=True)
class InputKeyResult:
	layout_changed: bool = False
	relinquish: bool = False
	committed: bool = False


_SECTION_ORDER: tuple[tuple[SectionId, str], ...] = (
	(SectionId.STAGED, "Staged"),
	(SectionId.UNSTAGED, "Unstaged"),
	(SectionId.UNTRACKED, "Untracked"),
)


class ChangesList:
	def __init__(self) -> None:
		self._collapsed: dict[SectionId, bool] = dict.fromkeys(SectionId, False)
		self._files: dict[SectionId, list[str]] = {section: [] for section in SectionId}
		self._rows: list[_Row] = []
		self._list_state = ListState()
		self._commit_state = TextState()
		self._error: str | None = None
		self._staged_count = 0

	@property
	def commit_message(self) -> str:
		return self._commit_state.value()

	def set_commit_message(self, message: str) -> None:
		self._commit_state.set_value(message)

	def persist_commit_draft(self, repo: Repository) -> None:
		save_commit_draft(Path(repo.path), self.commit_message)

	def is_at_top(self) -> bool:
		return (self._list_state.selected or 0) == 0

	def refresh(self, repo: Repository) -> None:
		staged, unstaged, untracked = categorize_status(repo)
		self._files = {
			SectionId.STAGED: staged,
			SectionId.UNSTAGED: unstaged,
			SectionId.UNTRACKED: untracked,
		}
		self._staged_count = len(staged)
		self._rebuild_rows()

	def _rebuild_rows(self) -> None:
		rows: list[_Row] = [
			_Row(kind="commit", section=None, path=None, label=self._commit_label()),
		]
		for section_id, title in _SECTION_ORDER:
			files = self._files[section_id]
			marker = ">" if self._collapsed[section_id] else "v"
			rows.append(
				_Row(
					kind="header",
					section=section_id,
					path=None,
					label=f"[{marker}] {title} ({len(files)})",
				),
			)
			if not self._collapsed[section_id]:
				for path in files:
					rows.append(_Row(kind="file", section=section_id, path=path, label=f"  {path}"))
		self._rows = rows
		selected = self._list_state.selected or 0
		self._list_state.select(min(selected, max(0, len(self._rows) - 1)))

	def draw(
		self,
		frame: Frame,
		area: Rect,
		*,
		highlight: bool = True,
		commit_input: bool = False,
	) -> None:
		if area.height <= 0 or area.width <= 0:
			return
		theme = current_theme()
		list_area = area
		if commit_input:
			chunks = (
				Layout()
				.direction(Direction.Vertical)
				.constraints([Constraint.length(1), Constraint.fill(1)])
				.split(area)
			)
			self._commit_state.focus()
			frame.render_text_prompt(TextPrompt("Commit: "), chunks[0], self._commit_state)
			list_area = chunks[1]

		items: list[ListItem] = []
		selected = self._list_state.selected or 0
		for index, row in enumerate(self._rows):
			if commit_input and row.kind == "commit":
				continue
			focused = highlight and index == selected and not commit_input
			if row.kind == "commit":
				style = focus_style(theme.commit) if focused else theme.commit
				items.append(ListItem(self._commit_label(), style))
				continue
			if row.kind == "header" and row.section is not None:
				style = row_style_for_section(
					row.section.value,
					focused=focused,
					highlight=highlight,
					is_header=True,
				)
			elif row.section is not None:
				style = row_style_for_section(
					row.section.value,
					focused=focused,
					highlight=highlight,
				)
			else:
				style = Style()
			items.append(ListItem(row.label, style))

		if not items:
			items.append(ListItem(" ", theme.normal))
		# Commit row is drawn as a TextPrompt above the list; keep ListState on
		# self._rows, but do not highlight a shifted list index while typing.
		list_state = self._list_state
		if commit_input:
			list_state = ListState()
			list_state.select(None)
		frame.render_stateful_list(
			List(items).highlight_style(theme.focus).highlight_symbol(""),
			list_area,
			list_state,
		)

	def handle_nav_key(self, key: KeyLike, repo: Repository) -> NavKeyResult:
		if self._error is not None:
			self._error = None

		if is_up(key):
			selected = self._list_state.selected or 0
			self._list_state.select(max(0, selected - 1))
			return NavKeyResult()
		if is_down(key):
			selected = self._list_state.selected or 0
			self._list_state.select(min(max(0, len(self._rows) - 1), selected + 1))
			return NavKeyResult()
		if not self._rows:
			return NavKeyResult()

		row = self._rows[self._list_state.selected or 0]
		if row.kind == "commit":
			if is_enter(key) and not self.commit_message.strip():
				return NavKeyResult(enter_input=True)
			if self._try_begin_input(key):
				return NavKeyResult(enter_input=True)
			return self._handle_commit_nav_key(key, repo)

		if (is_left(key) or is_right(key) or is_space(key)) and row.kind == "header" and row.section is not None:
			self._collapsed[row.section] = not self._collapsed[row.section]
			self._rebuild_rows()
			return NavKeyResult(layout_changed=True)
		if is_enter(key) and row.kind == "header" and row.section is not None:
			self._toggle_section(repo, row.section)
			self.refresh(repo)
			return NavKeyResult(layout_changed=True)
		if is_enter(key) and row.kind == "file" and row.path is not None and row.section is not None:
			self._toggle_stage(repo, row.section, row.path)
			self.refresh(repo)
			return NavKeyResult(layout_changed=True)
		return NavKeyResult()

	def handle_input_key(self, key: KeyLike, repo: Repository) -> InputKeyResult:
		if self._error is not None:
			self._error = None
		if is_enter(key):
			ok, error = commit_staged(repo, self.commit_message)
			if ok:
				self._commit_state.set_value("")
				self.persist_commit_draft(repo)
				self._error = None
				self.refresh(repo)
				return InputKeyResult(layout_changed=True, relinquish=True, committed=True)
			self._error = error
			return InputKeyResult()
		self._commit_state.focus()
		self._commit_state.handle_key(key)
		return InputKeyResult()

	def status_hint(self) -> str:
		if self._error:
			return self._error
		if not self._rows:
			return "Working tree clean"
		row = self._rows[self._list_state.selected or 0]
		if row.kind == "commit":
			if not self.commit_message.strip():
				return "Enter type message — j/k move"
			return "Type message — Enter commit — j/k move"
		if row.kind == "header":
			if row.section == SectionId.STAGED:
				if self._staged_count == 0:
					return "Enter stage all — Space collapse — j/k move"
				return "Enter unstage all — Space collapse — j/k move"
			return "Enter stage all — Space collapse — j/k move"
		if row.section == SectionId.STAGED:
			return "Enter unstage — j/k move"
		return "Enter stage — j/k move"

	def _try_begin_input(self, key: KeyLike) -> bool:
		if is_backspace(key):
			message = self.commit_message
			if not message:
				return False
			self._commit_state.set_value(message[:-1])
			return True
		char = printable_char(key)
		if char is not None:
			self._commit_state.set_value(self.commit_message + char)
			return True
		return False

	def _handle_commit_nav_key(self, key: KeyLike, repo: Repository) -> NavKeyResult:
		if is_enter(key):
			ok, error = commit_staged(repo, self.commit_message)
			if ok:
				self._commit_state.set_value("")
				self.persist_commit_draft(repo)
				self._error = None
				self.refresh(repo)
				return NavKeyResult(layout_changed=True, committed=True)
			self._error = error
		return NavKeyResult()

	def _commit_label(self) -> str:
		return f"Commit: {self.commit_message}"

	def _toggle_stage(self, repo: Repository, section: SectionId, path: str) -> None:
		if section == SectionId.STAGED:
			unstage_path(repo, path)
		else:
			stage_path(repo, path)

	def _toggle_section(self, repo: Repository, section: SectionId) -> None:
		paths = self._files[section]
		if section == SectionId.STAGED:
			if paths:
				unstage_paths(repo, paths)
			else:
				stage_paths(repo, self._files[SectionId.UNSTAGED] + self._files[SectionId.UNTRACKED])
		else:
			stage_paths(repo, paths)