from copy import deepcopy
import threading
from dataclasses import dataclass
from typing import Any, Callable, Literal, TypeVar, Union
from analyst_klondike.features.code.state import RunCodeState
from analyst_klondike.features.current.current_state import CurrentState
from analyst_klondike.features.data_context.data_state import DataState
from analyst_klondike.state.base_action import BaseAction


@dataclass
class Enable:
    run_code_button: bool
    open_file_button: bool


EventSenderType = Literal["local_file", "local_db", "remote_db", "no_events"]


@dataclass
class AppState:
    last_action_type: str
    user_email: str
    theme: str
    is_dark: bool
    is_editor_screen_ready: bool
    event_sender: EventSenderType
    data: DataState
    current: CurrentState
    run_code: RunCodeState
    is_teacher_mode: bool


INITIAL_STATE = AppState(
    user_email="",
    theme="",
    is_dark=True,
    is_editor_screen_ready=False,
    data=DataState(),
    current=CurrentState(),
    run_code=RunCodeState(),
    last_action_type="",
    event_sender="remote_db",
    is_teacher_mode=False
)

_state_callbacks: list[Callable[[AppState], None]] = []
_callback_lock = threading.RLock()

# ADD THIS FUNCTION


def register_state_callback(callback: Callable[[AppState], None]):
    """Allow UI to register for state updates"""
    with _callback_lock:
        if callback not in _state_callbacks:
            _state_callbacks.append(callback)


_state = INITIAL_STATE


def get_state() -> AppState:
    new_state = deepcopy(_state)
    return select(lambda _: new_state)


Reducers = Callable[[AppState, BaseAction], AppState]


def dispatch(action: BaseAction, *reducers: Reducers) -> None:
    copy_state = get_state()
    copy_state.last_action_type = action.type
    # pylint: disable=global-statement
    global _state
    _state = _pipe(copy_state, action, *reducers)

    # WITH THIS: (notify all registered callbacks)
    for callback in _state_callbacks.copy():
        callback(_state)


T = TypeVar('T')


def select(selector: Union[
        Callable[[AppState], T],
        Callable[[AppState, Any], T],
        Callable[[AppState, Any, Any], T]
    ],
        *params: Any) -> T:
    new_state = deepcopy(_state)
    return selector(new_state, *params)


def _pipe(state: AppState,
          action: BaseAction,
          *reducers: Callable[[AppState, BaseAction], AppState]) -> AppState:
    s = state
    for apply in reducers:
        s = apply(s, action)
    return s
