import typing
import time
from dataclasses import dataclass, field

@dataclass
class UISnapshot:
 """Capture of restorable UI state."""
 ai_input: str
 project_system_prompt: str
 global_system_prompt: str
 base_system_prompt: str
 use_default_base_prompt: bool
 temperature: float
 top_p: float
 max_tokens: int
 auto_add_history: bool
 disc_entries: list[dict]
 files: list[dict]
 screenshots: list[str]

 def to_dict(self) -> dict:
  return {
   "ai_input": self.ai_input,
   "project_system_prompt": self.project_system_prompt,
   "global_system_prompt": self.global_system_prompt,
   "base_system_prompt": self.base_system_prompt,
   "use_default_base_prompt": self.use_default_base_prompt,
   "temperature": self.temperature,
   "top_p": self.top_p,
   "max_tokens": self.max_tokens,
   "auto_add_history": self.auto_add_history,
   "disc_entries": self.disc_entries,
   "files": self.files,
   "screenshots": self.screenshots
  }

 @classmethod
 def from_dict(cls, data: dict) -> "UISnapshot":
  return cls(
   ai_input=data.get("ai_input", ""),
   project_system_prompt=data.get("project_system_prompt", ""),
   global_system_prompt=data.get("global_system_prompt", ""),
   base_system_prompt=data.get("base_system_prompt", ""),
   use_default_base_prompt=data.get("use_default_base_prompt", True),
   temperature=data.get("temperature", 0.0),
   top_p=data.get("top_p", 1.0),
   max_tokens=data.get("max_tokens", 4096),
   auto_add_history=data.get("auto_add_history", False),
   disc_entries=data.get("disc_entries", []),
   files=data.get("files", []),
   screenshots=data.get("screenshots", [])
  )

@dataclass
class HistoryEntry:
 state: typing.Any
 description: str
 timestamp: float = field(default_factory=lambda: time.time())

class HistoryManager:
 def __init__(self, max_capacity: int = 100):
  self.max_capacity = max_capacity
  self._undo_stack: typing.List[HistoryEntry] = []
  self._redo_stack: typing.List[HistoryEntry] = []

 def push(self, state: typing.Any, description: str) -> None:
  """
  Pushes a new state to the undo stack and clears the redo stack.
  If the undo stack exceeds max_capacity, the oldest state is removed.
  """
  entry = HistoryEntry(state=state, description=description)
  self._undo_stack.append(entry)
  self._redo_stack.clear()
  if len(self._undo_stack) > self.max_capacity:
   self._undo_stack.pop(0)

 def undo(self, current_state: typing.Any, current_description: str = "Current State") -> typing.Optional[HistoryEntry]:
  """
  Undoes the last action by moving the current_state to the redo stack
  and returning the top of the undo stack.
  """
  if not self._undo_stack:
   return None
  
  redo_entry = HistoryEntry(state=current_state, description=current_description)
  self._redo_stack.append(redo_entry)
  return self._undo_stack.pop()

 def redo(self, current_state: typing.Any, current_description: str = "Current State") -> typing.Optional[HistoryEntry]:
  """
  Redoes the last undone action by moving the current_state to the undo stack
  and returning the top of the redo stack.
  """
  if not self._redo_stack:
   return None
  
  undo_entry = HistoryEntry(state=current_state, description=current_description)
  self._undo_stack.append(undo_entry)
  return self._redo_stack.pop()

 @property
 def can_undo(self) -> bool:
  return len(self._undo_stack) > 0

 @property
 def can_redo(self) -> bool:
  return len(self._redo_stack) > 0

 def get_history(self) -> typing.List[typing.Dict[str, typing.Any]]:
  """Returns a list of descriptions and timestamps for the undo stack."""
  return [
   {"description": e.description, "timestamp": e.timestamp}
   for e in self._undo_stack
  ]

 def jump_to_undo(self, index: int, current_state: typing.Any, current_description: str = "Before Jump") -> typing.Optional[HistoryEntry]:
  """
  Jumps to a specific state in the undo stack by moving subsequent states
  and the current_state to the redo stack.
  """
  if index < 0 or index >= len(self._undo_stack):
   return None
  
  # Move current state to redo
  self._redo_stack.append(HistoryEntry(state=current_state, description=current_description))
  
  # Move states between index and top of undo to redo
  while len(self._undo_stack) > index + 1:
   self._redo_stack.append(self._undo_stack.pop())
   
  return self._undo_stack.pop()