from typing import Dict, Any
from api.assets import Asset
from api.io_port import InPort, OutPort


class ExecutionContext:
    """
    Runtime data container for a workflow execution.
    Stores produced assets for OutPorts and allows InPorts to retrieve them.
    """

    def __init__(self, execution_id: str | None = None):
        self.execution_id = execution_id

        # Key: OutPort, Value: Asset
        self._data: Dict[OutPort, Asset] = {}

        # Arbitrary metadata (optional)
        self._metadata: Dict[str, Any] = {}

    # Store produced data
    def put(self, output_port: OutPort, asset: Asset):
        """
        Register an asset produced by an OutPort.
        Also updates output_port.produced_asset.
        """
        self._data[output_port] = asset
        output_port.produced_asset = asset

    # Retrieve data
    def get(self, port):
        """
        Retrieve asset for:
        - OutPort → return its asset
        - InPort → return asset from its connected OutPort
        """
        if isinstance(port, OutPort):
            return self._data.get(port)

        if isinstance(port, InPort):
            if port.connected_output is None:
                return None
            return self._data.get(port.connected_output)

        raise TypeError("ExecutionContext.get() expects InPort or OutPort")

    # Check if data exists
    def has_data(self, port) -> bool:
        if isinstance(port, OutPort):
            return port in self._data

        if isinstance(port, InPort):
            if port.connected_output is None:
                return False
            return port.connected_output in self._data

        return False

    # Metadata management
    def set_meta(self, key: str, value: Any):
        self._metadata[key] = value

    def get_meta(self, key: str, default=None):
        return self._metadata.get(key, default)

    # Debug snapshot
    def status_view(self):
        """
        Lightweight snapshot of current execution state.
        Safe for GUI display.
        """
        return {
            "execution_id": self.execution_id,
            "assets": {
                f"{port.parent_task.name}.{port.name}": {
                    "dtype": port.dtype.__class__.__name__,
                    "preserved": asset.preserved,
                }
                for port, asset in self._data.items()
            },
            "metadata": self._metadata.copy(),
        }