from __future__ import annotations

import copy
import importlib
import sys
import traceback

from dataclasses import dataclass, field
from typing      import Any


@dataclass
class HotModule:
 name:               str
 file_path:          str
 state_keys:         list[str] = field(default_factory=list)
 delegation_targets: list[str] = field(default_factory=list)

class HotReloader:
 HOT_MODULES:    dict[str, HotModule] = {}
 last_error:     str | None = None
 is_error_state: bool = False

 @classmethod
 def register(cls, module: HotModule) -> None:
  if module.name in cls.HOT_MODULES:
   raise ValueError(f"Module {module.name} already registered")
  cls.HOT_MODULES[module.name] = module

 @classmethod
 def capture_state(cls, app: Any, state_keys: list[str]) -> dict[str, Any]:
  return {key: copy.deepcopy(getattr(app, key, None)) for key in state_keys if hasattr(app, key)}

 @classmethod
 def restore_state(cls, app: Any, state: dict[str, Any]) -> None:
  for key, value in state.items():
   setattr(app, key, value)

 @classmethod
 def reload(cls, module_name: str, app: Any) -> bool:
  if module_name not in cls.HOT_MODULES:
   cls.last_error     = f"Module {module_name} not registered"
   cls.is_error_state = True
   return False

  hm = cls.HOT_MODULES[module_name]
  state = cls.capture_state(app, hm.state_keys)

  try:
   if module_name in sys.modules:
    old_module = sys.modules[module_name]
    importlib.reload(old_module)
   else:
    importlib.import_module(module_name)
   cls.last_error     = None
   cls.is_error_state = False
   return True
  except Exception:
   cls.restore_state(app, state)
   cls.last_error     = traceback.format_exc()
   cls.is_error_state = True
   return False

 @classmethod
 def reload_all(cls, app: Any) -> bool:
  success = True
  for name in cls.HOT_MODULES:
   if not cls.reload(name, app): success = False
  return success