import logging
import sys
import typing

from luma.core.error import DeviceNotFoundError
from luma.core.interface.serial import i2c
from luma.oled.device import device, sh1106
from PIL import Image

log = logging.getLogger(__name__)


class I2C:
    def __init__(self) -> None:
        self.__device: device | None = None

    def __enter__(self) -> typing.Self:
        self.__open_device()
        return self

    def __exit__(self, *_) -> None:
        assert self.__device is not None
        self.__device.cleanup()

    def __open_device(self) -> None:
        # Pin definitions: I2C Bus 1 (SDA=GPIO2, SCL=GPIO3)
        # Default SSD1306 I2C address is 0x3C
        serial_interface = i2c(port=1, address=0x3C)

        try:
            # Initialize the device. If using SH1106, change ssd1306 to sh1106
            self.__device = sh1106(serial_interface, width=128, height=64, rotate=0)
        except DeviceNotFoundError:
            log.critical(
                "Fatal: OLED not found at 0x3C. Check wiring and run i2cdetect -y 1"
            )
            sys.exit(1)
        except OSError as e:
            if e.errno == 121:
                log.critical(
                    "Fatal: Remote I/O error. I2C bus is locked or physically disconnected."
                )
            else:
                log.critical(f"Fatal OS Error: {e}")
            sys.exit(1)

    def update_display(self, image: Image.Image) -> None:
        assert self.__device is not None
        try:
            self.__device.display(image)
        except OSError as e:
            if e.errno == 121:
                log.error("Runtime Error: I2C bus dropped. Check physical connections.")
            raise