from typing import Any, Dict
from workflow.task import Task, ExecutionContext, Status, TaskEvent


class Module(Task):
    """
    Leaf in the Composite Pattern.
    Performs actual computation.
    Users override ONLY run().
    """

    def __init__(self, name: str):
        super().__init__(name)
        self._in_ports: Dict[str, None] = {}
        self._out_ports: Dict[str, None] = {}

    # Ports (Module-only)
    def add_in_port(self, name: str):
        self._in_ports[name] = None

    def add_out_port(self, name: str):
        self._out_ports[name] = None

    # Execution
    def run(self, context: ExecutionContext):
        """
        Final execution method.
        Subclasses override THIS method only.
        No compute(), no run_node(), no extra execution methods.
        """

        # Logging + Events + Status
        self.log(f"Starting module: {self.name}")
        context.log(f"[{self.name}] start")
        self._emit(TaskEvent.BEFORE_RUN)
        self._set_status(Status.RUNNING)

        try:
            # Read inputs from ExecutionContex
            inputs: Dict[str, Any] = {}
            for port_name in self._in_ports.keys():
                key = f"{self.name}.{port_name}"
                inputs[port_name] = context.get_asset(key)

            result = self._user_run(context, inputs)

            # Write outputs
            if isinstance(result, dict):
                for out_name, value in result.items():
                    key = f"{self.name}.{out_name}"
                    context.put_asset(key, value)
            else:
                context.put_asset(f"{self.name}.result", result)

            # Metadata
            context.add_metadata(f"{self.name}.status", "completed")

            # Status + Events + Logging
            self._set_status(Status.COMPLETED)
            self._emit(TaskEvent.COMPLETED, result)
            self.log(f"Module completed: {self.name}")
            context.log(f"[{self.name}] completed")

            return result

        except Exception as e:
            self._set_status(Status.FAILED)
            self._emit(TaskEvent.ERROR, str(e))
            self.log(f"Module failed: {self.name} — {e}")
            context.log(f"[{self.name}] failed: {e}")
            raise

    # Internal wrapper for user logic
    def _user_run(self, context: ExecutionContext, inputs: Dict[str, Any]):
        """
        This is the ONLY method users override.
        """
        raise NotImplementedError(
            f"Module subclass '{self.name}' must implement run(context, inputs)"
        )