"""HamiltonSTAR composition class."""

from __future__ import annotations

from typing import Optional

from pylabrobot.arms import Arm
from pylabrobot.liquid_handling import LiquidHandler
from pylabrobot.liquid_handling.backends.hamilton.star import STARBackend
from pylabrobot.resources import Coordinate, Resource
from pylabrobot.resources.hamilton.hamilton_decks import HamiltonSTARDeck


class HamiltonSTAR(Resource):
  """Hamilton STAR liquid handling system.

  Composes a STARBackend with capability sub-frontends:
    - self.lh: LiquidHandler (pipetting, tip handling, 96-head)
    - self.iswap: Arm (iSWAP/core gripper resource movement), None if not installed

  Usage::

    star = HamiltonSTAR(deck=STARLetDeck())
    await star.setup()
    await star.lh.aspirate(plate["A1:H1"], vols=50)
    await star.iswap.move_plate(plate, plt_car[2])
    await star.stop()
  """

  def __init__(
    self,
    deck: HamiltonSTARDeck,
    backend: Optional[STARBackend] = None,
    name: str = "star",
  ):
    super().__init__(
      name=name,
      size_x=deck._size_x,
      size_y=deck._size_y,
      size_z=deck._size_z,
      category="liquid_handler",
    )
    self.backend = backend or STARBackend()
    self.deck = deck

    # Resource tree: STAR owns the deck
    self.location = Coordinate.zero()
    super().assign_child_resource(deck, location=deck.location or Coordinate.zero())

    # Capability sub-frontends (iswap created in setup() if installed)
    self.lh = LiquidHandler(backend=self.backend, deck=self.deck)
    self.iswap: Optional[Arm] = None

  async def setup(self, **backend_kwargs):
    self.backend.set_deck(self.deck)
    await self.backend.setup(**backend_kwargs)
    await self.lh.setup()

    if self.backend.iswap_installed:
      self.iswap = Arm(backend=self.backend, deck=self.deck)
      await self.iswap.setup()

  async def stop(self):
    await self.lh.stop()
    if self.iswap is not None:
      await self.iswap.stop()
    await self.backend.stop()

  async def __aenter__(self):
    await self.setup()
    return self

  async def __aexit__(self, *args):
    await self.stop()

  def assign_child_resource(self, resource, location=None, reassign=True):
    raise NotImplementedError(
      "Cannot assign child resource to HamiltonSTAR. Use star.deck.assign_child_resource() instead."
    )
