from __future__ import annotations

from contextlib import contextmanager
from contextvars import Context, ContextVar, copy_context
from dataclasses import dataclass, replace
from typing import Optional, Final

from comfy_execution.graph_types import FrozenTopologicalSort
from .cli_args_types import Configuration, LatentPreviewMethod
from .component_model import cvpickle
from .component_model.executor_types import ExecutorToClientProgress
from .component_model.folder_path_types import FolderNames
from .distributed.server_stub import ServerStub
from .nodes.package_typing import ExportedNodes, exported_nodes_view
from .progress_types import AbstractProgressRegistry, ProgressRegistryStub


@dataclass(frozen=True)
class ExecutionContext:
    # at any time
    server: ExecutorToClientProgress
    folder_names_and_paths: FolderNames
    custom_nodes: ExportedNodes
    configuration: Optional[Configuration] = None

    # during prompt execution
    progress_registry: Optional[AbstractProgressRegistry] = None
    preview_method_override: Optional[LatentPreviewMethod] = None

    # during node execution
    node_id: Optional[str] = None
    task_id: Optional[str] = None
    list_index: Optional[int] = None
    inference_mode: bool = True
    execution_list: Optional[FrozenTopologicalSort] = None
    executed: Optional[frozenset] = None

    @property
    def inputs(self) -> dict:
        return self.execution_list.dynprompt.get_node(self.node_id)['inputs']

    def __iter__(self):
        """
        Provides tuple-like unpacking behavior, similar to a NamedTuple.
        Yields task_id, node_id, and list_index.
        """
        yield self.task_id
        yield self.node_id
        yield self.list_index


comfyui_execution_context: Final[ContextVar] = ContextVar("comfyui_execution_context", default=ExecutionContext(server=ServerStub(), folder_names_and_paths=FolderNames(is_root=True), custom_nodes=ExportedNodes(), progress_registry=ProgressRegistryStub(), configuration=Configuration(), preview_method_override=None))  # noqa: B039  intentional: deep-copied per-frame in cvpickle to avoid sharing
# enables context var propagation across process boundaries for process pool executors
cvpickle.register_contextvar(comfyui_execution_context, __name__)


def current_execution_context() -> ExecutionContext:
    return comfyui_execution_context.get()


def copy_process_context() -> Context:
    """Copy context variables without request-local, non-picklable state.

    Process peers need Configuration, folder paths, custom-node exports, and
    trace context. The live server/progress objects and execution graph remain
    owned by the parent process and may contain generators, locks, or sockets.
    """

    source = copy_context()
    target = Context()
    for context_variable, value in source.items():
        if context_variable is comfyui_execution_context:
            value = replace(
                value,
                server=ServerStub(),
                progress_registry=ProgressRegistryStub(),
                execution_list=None,
                executed=None,
            )
        target.run(context_variable.set, value)
    return target


@contextmanager
def _new_execution_context(ctx: ExecutionContext):
    token = comfyui_execution_context.set(ctx)
    try:
        yield ctx
    finally:
        comfyui_execution_context.reset(token)


def _config_fields_changed(old_config: Optional[Configuration], new_config: Configuration, fields: frozenset[str]) -> bool:
    """Check if any of the specified configuration fields have changed."""
    if old_config is None:
        return True
    for field in fields:
        if getattr(old_config, field, None) != getattr(new_config, field, None):
            return True
    return False


@contextmanager
def context_configuration(configuration: Optional[Configuration] = None):
    """
    Context manager for setting configuration.

    If path-related configuration fields (defined in AFFECTS_PATHS) differ
    from the current context, folder_names_and_paths will be reinitialized
    with the new configuration.

    See component_model.configuration.AFFECTS_PATHS for the list of fields
    that trigger folder path reinitialization.
    """
    from .component_model.configuration import AFFECTS_PATHS

    current_ctx = current_execution_context()
    if configuration is None:
        from .cli_args import cli_args_configuration
        configuration = cli_args_configuration()

    # Check if paths changed and reinitialize folder_names_and_paths if needed
    if _config_fields_changed(current_ctx.configuration, configuration, AFFECTS_PATHS):
        from .cmd.folder_paths import init_default_paths
        new_folder_names = FolderNames(is_root=True)
        init_default_paths(new_folder_names, configuration, replace_existing=True)
        new_ctx = replace(current_ctx, configuration=configuration, folder_names_and_paths=new_folder_names)
    else:
        new_ctx = replace(current_ctx, configuration=configuration)

    with _new_execution_context(new_ctx):
        yield new_ctx


@contextmanager
def context_folder_names_and_paths(folder_names_and_paths: FolderNames):
    current_ctx = current_execution_context()
    new_ctx = replace(current_ctx, folder_names_and_paths=folder_names_and_paths)
    with _new_execution_context(new_ctx):
        yield new_ctx


@contextmanager
def context_execute_prompt(server: ExecutorToClientProgress, prompt_id: str, progress_registry: AbstractProgressRegistry, inference_mode: bool = True, preview_method_override=None):
    current_ctx = current_execution_context()
    new_ctx = replace(current_ctx, server=server, task_id=prompt_id, inference_mode=inference_mode, progress_registry=progress_registry, preview_method_override=preview_method_override)
    with _new_execution_context(new_ctx):
        yield new_ctx


@contextmanager
def context_execute_node(node_id: str):
    current_ctx = current_execution_context()
    new_ctx = replace(current_ctx, node_id=node_id)
    with _new_execution_context(new_ctx):
        yield new_ctx


@contextmanager
def context_add_custom_nodes(exported_nodes: ExportedNodes):
    """
    Adds custom nodes to the execution context
    :param exported_nodes: an object that represents a gathering of custom node export symbols
    :return: a context
    """
    current_ctx = current_execution_context()
    if len(exported_nodes) == 0:
        yield current_ctx

    if len(current_ctx.custom_nodes) == 0:
        merged_custom_nodes = exported_nodes
    else:
        merged_custom_nodes = exported_nodes_view(current_ctx.custom_nodes, exported_nodes)

    new_ctx = replace(current_ctx, custom_nodes=merged_custom_nodes)
    with _new_execution_context(new_ctx):
        yield new_ctx


@contextmanager
def context_set_node_and_prompt(prompt_id: str, node_id: str, list_index: Optional[int] = None):
    """
    A context manager to set the prompt_id (task_id), node_id, and optional list_index for the current execution.
    This is useful for fine-grained context setting within a node's execution, especially for batch processing.

    Replaces the @guill code upstream
    """
    current_ctx = current_execution_context()
    new_ctx = replace(current_ctx, task_id=prompt_id, node_id=node_id, list_index=list_index)
    with _new_execution_context(new_ctx):
        yield new_ctx


@contextmanager
def context_set_execution_list_and_inputs(execution_list: FrozenTopologicalSort, executed: frozenset):
    current_ctx = current_execution_context()
    new_ctx = replace(current_ctx, execution_list=execution_list, executed=executed)
    with _new_execution_context(new_ctx):
        yield new_ctx
