import asyncio
import json
import logging
import secrets
import ssl
import uuid
from datetime import datetime as dt
from typing import Callable

import aiohttp
import aiomqtt
from pubnub.enums import PNReconnectionPolicy
from pubnub.pnconfiguration import PNConfiguration
from pubnub.pubnub_asyncio import PubNubAsyncio

from .containers import (
    AuthorizationInfo,
    BabyData,
    SnooData,
    SnooDevice,
    SnooStates,
)
from .exceptions import InvalidSnooAuth, SnooAuthException, SnooBabyError, SnooCommandException, SnooDeviceError
from .pubnub_async import SnooPubNub

_LOGGER = logging.getLogger(__name__)


class Snoo:
    def __init__(self, email: str, password: str, clientsession: aiohttp.ClientSession):
        self.email = email
        self.password = password
        self.session = clientsession
        self.aws_auth_url = "https://cognito-idp.us-east-1.amazonaws.com/"
        self.snoo_auth_url = "https://api-us-east-1-prod.happiestbaby.com/us/me/v10/pubnub/authorize"
        self.snoo_devices_url = "https://api-us-east-1-prod.happiestbaby.com/hds/me/v11/devices"
        self.snoo_data_url = "https://happiestbaby.pubnubapi.com"
        self.snoo_baby_url = "https://api-us-east-1-prod.happiestbaby.com/us/me/v10/babies"
        self.aws_auth_hdr = {
            "x-amz-target": "AWSCognitoIdentityProviderService.InitiateAuth",
            "accept-language": "US",
            "content-type": "application/x-amz-json-1.1",
            "accept-encoding": "gzip",
            "user-agent": "okhttp/4.12.0",
            "accept": "application/json",
        }
        self.aws_refresh_hdr = {
            "accept": "*/*",
            "content-type": "application/x-amz-json-1.1",
            "x-amz-target": "AWSCognitoIdentityProviderService.InitiateAuth",
            "accept-encoding": "br;q=1.0, gzip;q=0.9, deflate;q=0.8",
            "user-agent": "Happiest Baby/2.1.6 (com.happiestbaby.hbapp; build:88; iOS 18.3.0) Alamofire/5.9.1",
            "accept-language": "en-US;q=1.0",
            "content-length": "1895",
        }
        self.snoo_auth_hdr = {
            "accept-language": "US",
            "content-type": "application/json; charset=UTF-8",
            "accept-encoding": "gzip",
            "user-agent": "okhttp/4.12.0",
            "accept": "application/json",
        }

        self.aws_auth_data = {
            "AuthParameters": {
                "PASSWORD": self.password,
                "USERNAME": self.email,
            },
            "AuthFlow": "USER_PASSWORD_AUTH",
            "ClientId": "6kqofhc8hm394ielqdkvli0oea",
        }
        self.snoo_auth_data = {
            "advertiserId": "",
            "appVersion": "1.8.7",
            "device": "panther",
            "deviceHasGSM": True,
            "locale": "en",
            "os": "Android",
            "osVersion": "14",
            "platform": "Android",
            "timeZone": "America/New_York",
            "userCountry": "US",
            "vendorId": "eyqurgwYQSqmnExnzyiLO5",
        }

        self.aws_auth_data = json.dumps(self.aws_auth_data)
        self.snoo_auth_data = json.dumps(self.snoo_auth_data)
        self.tokens: AuthorizationInfo | None = None
        self.pubnub = None
        self.subscription_functions = {}
        self.data_map = {}
        self.pubnub_instances: dict[str, SnooPubNub] = {}
        self.reauth_task: asyncio.Task | None = None
        self._client_map: dict[str, aiomqtt.Client] = {}
        self._mqtt_tasks: dict[str, asyncio.Task] = {}
        self._mqtt_callbacks: dict[str, tuple[SnooDevice, Callable]] = {}
        self._client_cond = asyncio.Condition()

    async def refresh_tokens(self) -> int:
        """Refreshes AWS Cognito tokens and returns the new expiration time in seconds."""
        _LOGGER.info("AWS Cognito tokens expired, refreshing...")
        if not self.tokens or not self.tokens.aws_refresh:
            _LOGGER.error("No refresh token available. A full re-authentication is required.")
            raise SnooAuthException("Missing refresh token.")

        data = {
            "AuthParameters": {"REFRESH_TOKEN": self.tokens.aws_refresh},
            "AuthFlow": "REFRESH_TOKEN_AUTH",
            "ClientId": "6kqofhc8hm394ielqdkvli0oea",
        }
        r = await self.session.post(self.aws_auth_url, json=data, headers=self.aws_auth_hdr)
        resp = await r.json(content_type=None)

        if r.status >= 400:
            _LOGGER.error(f"Failed to refresh tokens. Status: {r.status}, Response: {resp}")
            raise InvalidSnooAuth(f"Token refresh failed: {resp.get('message', 'Unknown error')}")

        result = resp.get("AuthenticationResult")
        if not result:
            _LOGGER.error(f"Invalid response during token refresh: {resp}")
            raise SnooAuthException("Token refresh response missing 'AuthenticationResult'.")

        # Update tokens with the new ones from the response
        self.tokens = AuthorizationInfo(
            snoo=self.tokens.snoo,
            aws_access=result["AccessToken"],
            aws_id=result["IdToken"],
            aws_refresh=result.get("RefreshToken", self.tokens.aws_refresh),
        )
        _LOGGER.info("✅ Successfully refreshed AWS Cognito tokens.")
        return result.get("ExpiresIn", 3600)

    def check_tokens(self):
        if self.tokens is None:
            raise Exception("You need to authenticate before you continue")

    def generate_snoo_auth_headers(self, amz_token: str) -> dict[str, str]:
        hdrs = self.snoo_auth_hdr.copy()
        hdrs["authorization"] = f"Bearer {amz_token}"
        return hdrs

    def generate_snoo_data_url(self, device_id: str | float, snoo_token: str) -> str:
        if isinstance(device_id, float):
            device_id = str(int(device_id))
        req_uuid = uuid.uuid1()
        dev_uuid = uuid.uuid1()
        app_dev_id_len = 24
        n = app_dev_id_len * 3 // 4
        app_dev_id = secrets.token_urlsafe(n)
        url = f"https://happiestbaby.pubnubapi.com/v2/history/sub-key/sub-c-97bade2a-483d-11e6-8b3b-02ee2ddab7fe/channel/ActivityState.{device_id}?pnsdk=PubNub-Kotlin%2F7.4.0&l_pub=0.064&auth={snoo_token}&requestid={req_uuid}&include_token=true&count=1&include_meta=false&reverse=false&uuid=android_{app_dev_id}_{dev_uuid}"
        return url

    def generate_id(self) -> str:
        app_dev_id_len = 24
        n = app_dev_id_len * 3 // 4
        app_dev_id = secrets.token_urlsafe(n)
        return app_dev_id

    async def subscribe(self, device: SnooDevice, function: Callable):
        pnconfig = PNConfiguration()
        pnconfig.subscribe_key = "sub-c-97bade2a-483d-11e6-8b3b-02ee2ddab7fe"
        pnconfig.publish_key = "pub-c-699074b0-7664-4be2-abf8-dcbb9b6cd2bf"
        pnconfig.user_id = secrets.token_urlsafe(16)
        pnconfig.auth_key = self.tokens.snoo
        pnconfig.reconnect_policy = PNReconnectionPolicy.EXPONENTIAL
        self.pubnub = PubNubAsyncio(pnconfig)
        device_id = device.serialNumber

        if device_id not in self.pubnub_instances:
            self.pubnub_instances[device_id] = SnooPubNub(self.pubnub, device_id)
        pubnub_instance = self.pubnub_instances[device_id]
        unsub = pubnub_instance.subscribe(function)
        asyncio.create_task(pubnub_instance.run())
        return unsub

    async def disconnect(self):
        for pubnub_instance in self.pubnub_instances.values():
            if pubnub_instance.task:
                pubnub_instance.task.cancel()
                try:
                    await pubnub_instance.task
                except asyncio.CancelledError:
                    pass
        self.pubnub_instances = {}

        for task in self._mqtt_tasks.values():
            task.cancel()
        await asyncio.gather(*self._mqtt_tasks.values(), return_exceptions=True)
        self._mqtt_tasks = {}
        self._mqtt_callbacks = {}

        if self.reauth_task:
            self.reauth_task.cancel()
            self.reauth_task = None

    def publish_callback(self, result, status):
        if status.is_error():
            _LOGGER.warning(f"Message failed with {status.status_code}, {status.error_data.__dict__}")

    async def send_command(self, command: str, device: SnooDevice, **kwargs):
        ts = int(dt.now().timestamp() * 10_000_000)
        try:
            # Acquire the condition lock
            async with self._client_cond:
                try:
                    # Wait up to 30 seconds for the client to connect.
                    await asyncio.wait_for(
                        self._client_cond.wait_for(lambda: device.serialNumber in self._client_map), timeout=30.0
                    )
                except asyncio.TimeoutError:
                    _LOGGER.error(f"Timed out waiting for client for device {device.serialNumber} to connect.")
                    raise SnooCommandException(f"Client for device {device.serialNumber} is not connected.") from None

                # Once wait_for returns, we know the client exists and we hold the lock.
                await self._client_map[device.serialNumber].publish(
                    topic=f"{device.awsIoT.thingName}/state_machine/control",
                    payload=json.dumps({"ts": ts, "command": command, **kwargs}),
                )
        except Exception as e:
            raise SnooCommandException from e

    async def start_snoo(self, device: SnooDevice):
        await self.send_command("start_snoo", device)

    async def stop_snoo(self, device: SnooDevice):
        await self.send_command("go_to_state", device, **{"state": "ONLINE", "hold": "off"})

    async def set_level(self, device: SnooDevice, level: SnooStates, hold: bool = False):
        if hold:
            hold = "on"
        else:
            hold = "off"

        await self.send_command("go_to_state", device, **{"state": level.value, "hold": hold})

    async def set_sticky_white_noise(self, device: SnooDevice, on: bool):
        await self.send_command(
            "set_sticky_white_noise",
            device,
            **{"state": "on" if on else "off", "timeout_min": 15},
        )

    async def get_status(self, device: SnooDevice):
        await self.send_command("send_status", device)

    async def auth_amazon(self) -> dict:
        r = await self.session.post(self.aws_auth_url, data=self.aws_auth_data, headers=self.aws_auth_hdr)
        resp = await r.json(content_type=None)
        if "__type" in resp and resp["__type"] == "NotAuthorizedException":
            raise InvalidSnooAuth()
        result = resp["AuthenticationResult"]
        return result

    async def auth_snoo(self, id_token: str) -> dict:
        hdrs = self.generate_snoo_auth_headers(id_token)
        r = await self.session.post(self.snoo_auth_url, data=self.snoo_auth_data, headers=hdrs)
        return await r.json()

    async def authorize(self) -> AuthorizationInfo:
        try:
            amz = await self.auth_amazon()
            access = amz["AccessToken"]
            _id = amz["IdToken"]
            ref = amz["RefreshToken"]
            expires_in = amz["ExpiresIn"]

            snoo_token_data = await self.auth_snoo(_id)
            snoo_token = snoo_token_data["snoo"]["token"]

            self.tokens = AuthorizationInfo(snoo=snoo_token, aws_access=access, aws_id=_id, aws_refresh=ref)

            if self.reauth_task:
                self.reauth_task.cancel()

            # Schedule reauthorization with a 5-minute buffer before expiry
            reauth_delay = max(expires_in - 300, 0)
            self.reauth_task = asyncio.create_task(self.schedule_reauthorization(reauth_delay))
            _LOGGER.info(f"Authorization successful. Next token refresh scheduled in {reauth_delay} seconds.")

        except InvalidSnooAuth as ex:
            raise ex
        except Exception as ex:
            raise SnooAuthException from ex
        return self.tokens

    async def schedule_reauthorization(self, expiry_seconds: float):
        try:
            await asyncio.sleep(expiry_seconds)
            _LOGGER.info("Executing scheduled token refresh...")

            new_expires_in = await self.refresh_tokens()

            _LOGGER.info("Restarting MQTT subscriptions with new token...")
            # Cancel all existing MQTT tasks
            for task in self._mqtt_tasks.values():
                task.cancel()
            await asyncio.gather(*self._mqtt_tasks.values(), return_exceptions=True)
            self._mqtt_tasks.clear()

            # The `finally` block in `subscribe_mqtt` should clear the `_client_map`
            # as connections close.

            # Re-subscribe for all previously active subscriptions
            for device_sn, (device, function) in self._mqtt_callbacks.items():
                _LOGGER.info(f"Re-establishing MQTT subscription for device {device_sn}")
                self.start_subscribe(device, function)

            _LOGGER.info("✅ MQTT subscriptions restarted successfully.")

            # Schedule the *next* reauthorization
            reauth_delay = max(new_expires_in - 300, 0)
            self.reauth_task = asyncio.create_task(self.schedule_reauthorization(reauth_delay))
            _LOGGER.info(f"Next token refresh scheduled in {reauth_delay} seconds.")

        except asyncio.CancelledError:
            _LOGGER.info("Reauthorization task was cancelled.")
        except Exception:
            _LOGGER.exception("An unexpected error occurred during reauthorization.")

    async def get_devices(self) -> list[SnooDevice]:
        hdrs = self.generate_snoo_auth_headers(self.tokens.aws_id)
        try:
            r = await self.session.get(self.snoo_devices_url, headers=hdrs)
            resp = await r.json()
        except Exception as ex:
            raise SnooDeviceError from ex
        devs = [SnooDevice.from_dict(dev) for dev in resp["snoo"]]
        return devs

    async def get_babies(self) -> list[BabyData]:
        hdrs = self.generate_snoo_auth_headers(self.tokens.aws_id)
        try:
            r = await self.session.get(self.snoo_baby_url, headers=hdrs)
            resp = await r.json()
        except Exception as ex:
            raise SnooBabyError from ex
        babies = [BabyData.from_dict(baby) for baby in resp]
        return babies

    def start_subscribe(self, device: SnooDevice, function: Callable):
        if device.serialNumber in self._mqtt_tasks and not self._mqtt_tasks[device.serialNumber].done():
            _LOGGER.warning(f"Subscription task for device {device.serialNumber} is already running.")
            return

        # Store the device and callback function for re-subscription after re-auth
        self._mqtt_callbacks[device.serialNumber] = (device, function)

        self._mqtt_tasks[device.serialNumber] = asyncio.create_task(self.subscribe_mqtt(device, function))

    async def subscribe_mqtt(self, device: SnooDevice, function: Callable):
        host = device.awsIoT.clientEndpoint
        port = 443
        websocket_path = "/mqtt"
        token = self.tokens.aws_id
        headers = {"token": token}
        password = None
        client_id = f"HA_{uuid.uuid4()}"
        user_name = "?SDK=iOS&Version=2.40.1"

        logging.debug(f"Attempting to connect to wss://{host}:{port}{websocket_path}")

        # The default SSL context creation is a blocking I/O operation.
        # Run it in a separate thread to avoid blocking the Home Assistant event loop.
        ssl_context = await asyncio.to_thread(ssl.create_default_context)

        try:
            async with aiomqtt.Client(
                hostname=host,
                port=port,
                username=user_name,
                password=password,
                identifier=client_id,
                transport="websockets",
                websocket_path=websocket_path,
                websocket_headers=headers,
                tls_context=ssl_context,
                protocol=aiomqtt.ProtocolVersion.V31,
                timeout=10,
            ) as client:
                logging.info(f"✅ Successfully connected to MQTT broker for {device.serialNumber}!")

                # Acquire the lock, add the client to the map, and notify waiting tasks.
                async with self._client_cond:
                    self._client_map[device.serialNumber] = client
                    self._client_cond.notify_all()

                topic = f"{device.awsIoT.thingName}/state_machine/activity_state"
                await client.subscribe(topic)
                logging.info(f"Subscribed to topic: {topic}")

                async for message in client.messages:
                    logging.debug(f"Received message on topic '{message.topic}': {message.payload.decode()}")
                    function(SnooData.from_json(message.payload.decode()))

        except aiomqtt.MqttError as e:
            logging.error(f"MQTT connection for {device.serialNumber} failed: {e}")
        except Exception as e:
            logging.error(f"MQTT connection for {device.serialNumber} failed with an unexpected error: {e}")
        finally:
            # When the connection is lost, remove the client from the map.
            logging.info(f"MQTT connection closed for {device.serialNumber}.")
            async with self._client_cond:
                if device.serialNumber in self._client_map:
                    del self._client_map[device.serialNumber]