import json
import logging
from collections.abc import Callable
from dataclasses import dataclass
from typing import Any

import mqttapi as mqtt
from pubnub.enums import PNReconnectionPolicy
from pysnoo import ActivityState
from pysnoo import Device
from pysnoo import Snoo
from pysnoo import SnooAuthSession
from pysnoo import SnooPubNub as SnooPubNubBase
from pysnoo.models import EventType


# HACK: Avoid error on missing EventType
EventType._missing_ = lambda x: EventType.ACTIVITY

# Set global log level to debug
logging.getLogger().setLevel(logging.DEBUG)

# Use polling or listening
POLLING = True


# HACK: Subclass to modify original pubnub policy
class SnooPubNub(SnooPubNubBase):
    """Subclass of the original Snoo pubnub to alter the reconnect policy."""

    @staticmethod
    def _setup_pnconfig(access_token, uuid):
        """Generate Setup"""
        pnconfig = SnooPubNubBase._setup_pnconfig(access_token, uuid)
        pnconfig.reconnect_policy = PNReconnectionPolicy.EXPONENTIAL
        return pnconfig


def _serialize_device_info(device: Device) -> dict[str, Any]:
    return {
        "identifiers": [
            device.serial_number,
        ],
        "name": device.baby,
    }


@dataclass
class SnooActivityEntity:
    """Represents an entity that should be updated via Activity messages."""

    field: str
    component: str
    device_class: str | None = None
    name: str | None = None
    value_template: str | None = None
    command_callback: Callable | None = None

    @property
    def friendly_name(self) -> str:
        return self.field.replace("_", " ").capitalize()

    @property
    def config_topic(self) -> str:
        entity_name = self.field if self.field != "state" else "snoo_state"
        return f"homeassistant/{self.component}/{entity_name}/config"

    def device_topic(self, device: Device) -> str:
        return f"homeassistant/snoo_{device.serial_number}"

    def command_topic(self, device: Device) -> str | None:
        if self.component != "switch":
            return None

        return f"{self.device_topic(device)}/{self.field}/set"

    def discovery_message(self, device: Device) -> dict[str, Any]:
        payload = {
            "~": self.device_topic(device),
            "name": self.name or self.friendly_name,
            "device_class": self.device_class,
            "state_topic": "~/state",
            "unique_id": f"{device.serial_number}_{self.field}",
            "device": _serialize_device_info(device),
            "value_template": self.value_template
            or "{{ value_json." + self.field + " }}",
        }

        if self.component in {"binary_sensor", "switch"}:
            # Use "True" and "False" for binary sensor payloads
            payload.update(
                {
                    # Safety classes are actually True if unsafe, so flipping these
                    "payload_on": "True" if self.device_class != "safety" else "False",
                    "payload_off": "False" if self.device_class != "safety" else "True",
                }
            )

        if command_topic := self.command_topic(device):
            payload.update(
                {
                    "command_topic": command_topic,
                }
            )

        return payload


async def status_switch_callback(pubnub: SnooPubNub, data: dict[str, Any]):
    if data["state"] == "ON":
        await pubnub.publish_start()
    else:
        await pubnub.publish_goto_state("ONLINE")


async def hold_switch_callback(pubnub: SnooPubNub, data: dict[str, Any]):
    last_activity_state = (await pubnub.history(1))[0]
    current_state = last_activity_state.state_machine.state
    if not current_state.is_active_level():
        return
    if data["state"] == "ON":
        await pubnub.publish_goto_state(current_state, True)
    else:
        await pubnub.publish_goto_state(current_state, False)


ENTITIES = [
    SnooActivityEntity(
        "left_safety_clip",
        "binary_sensor",
        device_class="safety",
    ),
    SnooActivityEntity(
        "right_safety_clip",
        "binary_sensor",
        device_class="safety",
    ),
    SnooActivityEntity(
        "weaning",
        "binary_sensor",
        value_template="{{ value_json.state_machine.weaning }}",
    ),
    SnooActivityEntity(
        "is_active_session",
        "binary_sensor",
        device_class="running",
        name="Active Session",
        value_template="{{ value_json.state_machine.is_active_session }}",
    ),
    SnooActivityEntity(
        "state",
        "sensor",
        value_template="{{ value_json.state_machine.state }}",
    ),
    SnooActivityEntity(
        "status",
        "switch",
        value_template="{{ value_json.state_machine.state != 'ONLINE' }}",
        command_callback=status_switch_callback,
    ),
    SnooActivityEntity(
        "hold",
        "switch",
        value_template="{{ value_json.state_machine.hold }}",
        command_callback=hold_switch_callback,
    ),
]


def _get_token(token_file: str):
    """Read auth token from JSON file."""
    try:
        with open(token_file) as infile:
            return json.load(infile)
    except FileNotFoundError:
        pass
    except ValueError:
        pass


def _get_token_updater(token_file: str):
    def token_updater(token):
        with open(token_file, "w") as outfile:
            json.dump(token, outfile)

    return token_updater


def _device_to_pubnub(snoo: Snoo, device: Device) -> SnooPubNub:
    pubnub = SnooPubNub(
        snoo.auth.access_token,
        device.serial_number,
        f"pn-pysnoo-{device.serial_number}",
    )
    # TODO: Add some way of passing this config into pysnoo
    # Another temp option to subclass SnooPubNub
    # pubnub._pubnub.config.reconnect_policy = PNReconnectionPolicy.EXPONENTIAL
    return pubnub


class SnooMQTT(mqtt.Mqtt):
    def __init__(self, *args, **kwargs) -> None:
        self._snoo: Device | None = None
        self._devices: list[Device] = []
        self._pubnubs: dict[str, SnooPubNub] = {}
        super().__init__(*args, **kwargs)

    async def initialize(self) -> None:
        token_path = self._token_path
        token = _get_token(token_path)
        token_updater = _get_token_updater(token_path)

        async with SnooAuthSession(token, token_updater) as auth:
            if not auth.authorized:
                new_token = await auth.fetch_token(self._username, self._password)
                token_updater(new_token)
                self.log("got inital token")

            self._snoo = Snoo(auth)

        self._devices = await self._snoo.get_devices()
        if not self._devices:
            raise ValueError("No Snoo devices connected to account")

        # Publish discovery information
        self._publish_discovery(self._devices)

        # Subscribe to updates
        for device in self._devices:
            pubnub = _device_to_pubnub(self._snoo, device)
            self._pubnubs[device.serial_number] = pubnub

            if POLLING:
                await self._schedule_updates(device, pubnub)
            else:
                await self._subscribe_and_listen(device, pubnub)

        # Listen for home assistant status to republish discovery
        self.listen_event(
            self.birth_callback,
            "MQTT_MESSAGE",
            topic="homeassistant/status",
        )

    async def birth_callback(self, event_name, data, *args) -> None:
        """Callback listening for hass status messages.

        Should republish discovery and initial status for every device."""
        self.log("Read hass status event: %s, %s", event_name, data)

        self._publish_discovery(self._devices)
        for device in self._devices:
            pubnub = self._pubnubs[device.serial_number]
            for activity_state in await pubnub.history(1):
                self._create_activity_callback(device)(activity_state)

    async def _schedule_updates(self, device: Device, pubnub: SnooPubNub):
        """Schedules from pubnub history periodic updates."""
        cb = self._create_activity_callback(device)

        async def poll_history(*args):
            for activity_state in await pubnub.history(1):
                cb(activity_state)

        self.run_every(poll_history, "now", 30)

    async def _subscribe_and_listen(self, device: Device, pubnub: SnooPubNub):
        """Subscribes to pubnub activity and listens to new events."""
        cb = self._create_activity_callback(device)
        pubnub.add_listener(cb)

        # Publish first status message
        for activity_state in await pubnub.history(1):
            cb(activity_state)

        await pubnub.subscribe_and_await_connect()

    def _create_activity_callback(self, device: Device):
        """Creates an activity callback for a given device."""

        def activity_callback(activity_state: ActivityState):
            self._publish_object(
                f"homeassistant/snoo_{device.serial_number}/state",
                activity_state,
            )

        return activity_callback

    def _publish_discovery(self, devices: list[Device]):
        """Publishes discovery messages for the provided devices."""
        for device in devices:
            for entity in ENTITIES:
                self._publish_object(
                    entity.config_topic,
                    entity.discovery_message(device),
                )

                # See if we need to listen to commands
                if (
                    entity.command_callback is not None
                    and entity.command_topic is not None
                ):

                    async def cb(self, event_name, data, *args):
                        assert entity.command_callback
                        await entity.command_callback(
                            self._pubnubs[device.serial_number], json.loads(data)
                        )

                    self.listen_event(
                        cb, "MQTT_MESSAGE", topic=entity.command_topic(device)
                    )

    def _publish_object(self, topic, payload_object):
        """Publishes an object, serializing it's payload as JSON."""
        try:
            payload = json.dumps(payload_object.to_dict())
        except AttributeError:
            payload = json.dumps(payload_object)

        self.log("topic: %s payload: %s", topic, payload)
        self.mqtt_publish(topic, payload=payload)

    @property
    def _username(self):
        username = self.args.get("username")
        if not username:
            raise ValueError("Must provide a username")
        return username

    @property
    def _password(self):
        password = self.args.get("password")
        if not password:
            raise ValueError("Must provide a password")
        return password

    @property
    def _token_path(self):
        return self.args.get("token_path", "snoo_token.json")