from api.task import Task, Status, TaskEvent
from api.assets import Asset
from api.execution_context import ExecutionContext
from api.io_port import CompatibilityException, InPort, OutPort


class Module(Task):
    """
    Leaf node: atomic executable unit.
    Executes using its input ports and produces assets on output ports.
    """

    def add_in_port(self, port: InPort):
        port.parent_task = self
        self.in_ports.append(port)

    def add_out_port(self, port: OutPort):
        port.parent_task = self
        self.out_ports.append(port)

    # Hooks
    def before_run(self, context: ExecutionContext):
        pass

    def after_run(self, context: ExecutionContext):
        pass

    def on_error(self, error: Exception):
        pass

    # Validation
    def validate(self):
        super().validate()

    # Main execution entry point
    def run(self, context: ExecutionContext):
        self.before_run(context)

        if not self.check_inputs_ready(context):
            self._set_status(Status.PENDING)
            return

        self._set_status(Status.RUNNING)
        self._emit(TaskEvent.ON_START, {"module": self.name})

        try:
            # Gather input data
            consumed_data = {}
            for port in self.in_ports:
                if not port.is_ready(context):
                    raise CompatibilityException(
                        f"Input {port.full_name()} is not ready"
                    )
                asset = context.get(port)
                consumed_data[port.name] = asset.data

            # Execute module logic
            result = self.execute(consumed_data)

            if not isinstance(result, dict):
                raise ValueError(f"Module '{self.name}' must return dict")

            # Store outputs
            for out_port in self.out_ports:
                if out_port.name not in result:
                    raise ValueError(
                        f"Output '{out_port.name}' missing in module '{self.name}' result"
                    )

                asset = Asset(
                    data=result[out_port.name],
                    dtype=out_port.dtype
                )

                context.put(out_port, asset)

            self._set_status(Status.COMPLETED)
            self._emit(TaskEvent.ON_FINISH, {"module": self.name})
            self.after_run(context)

        except Exception as e:
            self._set_status(Status.FAILED)
            self.on_error(e)
            self._emit(TaskEvent.ON_ERROR, {"error": str(e)})
            raise e

    # Default execute
    def execute(self, inputs):
        return inputs