# SPDX-License-Identifier: Apache-2.0
"""Opening the Blackhole p150 the way every published number of diffusion-planner-p150 was measured.

The implementation is the vendored ``ttaw.device`` (C01): dispatch on the idle ETH cores
(``ttnn.DispatchCoreConfig(ttnn.DispatchCoreType.ETH)``), which gives a 12x10 = 120-core compute grid on a p150
(11x10 with WORKER dispatch, the A/B switch). ETH dispatch needs ``patches/tt-metal-eth-dispatch.patch`` on tt-metal
44d6650 (the container image is built from a patched tree); if the ETH open fails, a ``RuntimeWarning`` is issued and
WORKER dispatch is used, unless ``allow_fallback=False``. Never hard-code the grid: use :func:`compute_grid` or
``device.compute_with_storage_grid_size()``.

This module binds it to the validated open parameters of this port, :data:`DEVICE_DEFAULTS` (the model class uses
the same dict), which the ``DIFFUSION_PLANNER_DISPATCH``, ``_NUM_CQS``, ``_L1_SMALL``, ``_TRACE_REGION`` and
``_WORKER_L1_SIZE`` variables and ``TT_DEVICE_ID`` override per process (SERVING.md section 3.4). Importing it has
no side effects (``ttnn`` is imported when a device is opened).
"""
from __future__ import annotations

import contextlib
import dataclasses
from typing import Any, Dict, Iterator, Optional

from .ttaw.device import DeviceConfig, close_device, compute_grid, core_grid, describe_device, full_core_range_set

__all__ = ["ENV_PREFIX", "DEVICE_DEFAULTS", "DeviceConfig", "device_config", "open_device", "device_session",
           "close_device", "describe_device", "compute_grid", "core_grid", "full_core_range_set"]

ENV_PREFIX = "DIFFUSION_PLANNER"

# Validated device-open parameters of this port (fill per model; the card's numbers are measured with them).
DEVICE_DEFAULTS: Dict[str, Any] = {
    "num_command_queues": 1,  # 1, or 2 when the input upload (CQ1) overlaps the trace (CQ0)
    "l1_small_size": 32768,  # L1_SMALL bytes per core (conv / pool config tensors)
    # DRAM bytes for the traces: measured 74.6 MB for the plan, 96.9 MB with the tests' debug variants (PORT_LOG 5.3)
    "trace_region_size": 192 << 20,
}


def device_config(**overrides: Any) -> DeviceConfig:
    """:data:`DEVICE_DEFAULTS` < the ``DIFFUSION_PLANNER_*`` / ``TT_DEVICE_ID`` environment < explicit non-None
    ``overrides``
    (``device_id``, ``dispatch``, ``num_command_queues``, ``l1_small_size``, ``trace_region_size``,
    ``worker_l1_size``, ``allow_fallback``): the same resolution as ``DiffusionPlanner.from_pretrained``."""
    config = DeviceConfig.from_env(ENV_PREFIX, **DEVICE_DEFAULTS)
    return dataclasses.replace(config, **{k: v for k, v in overrides.items() if v is not None})


def open_device(device_id: Optional[int] = None, *, dispatch: Optional[str] = None, **overrides: Any):
    """Open one chip like the published numbers: ETH dispatch (``dispatch="worker"`` is the A/B opt-in) and this
    port's sizes. Close it with :func:`close_device`, or use :func:`device_session`."""
    return device_config(device_id=device_id, dispatch=dispatch, **overrides).open()


@contextlib.contextmanager
def device_session(device_id: Optional[int] = None, **overrides: Any) -> Iterator[Any]:
    """``with device_session() as dev:`` opens with :func:`open_device` and always closes, also when the body
    raises."""
    device = open_device(device_id, **overrides)
    try:
        yield device
    finally:
        close_device(device)