import json
import logging
from multiprocessing import Queue
from typing import Any
import os
import asyncio

import aiohttp
from websockets.asyncio.client import connect
import websockets.exceptions

logger = logging.getLogger(__name__)
logging.getLogger("websockets.client").setLevel(logging.ERROR)

devices_template = """{% set devices = states | map(attribute='entity_id') | map('device_id') | unique | reject('eq',None) | list %}
{
	"devices": [
		{%- for device in devices %}
		{
			"id": "{{ device }}",
			"name": "{{ device_attr(device, 'name') }}",
			"manufacturer": "{{ device_attr(device, 'manufacturer') }}",
			"model": "{{ device_attr(device, 'model') }}",
			"identifiers": {{ device_attr(device, 'identifiers') | list | tojson }},
			"connections": {{ device_attr(device, 'connections') | list | tojson }},
			"configuration_url": "{{ device_attr(device, 'configuration_url') }}"
		}{% if not loop.last %},{% endif %}
		{%- endfor %}
	]
}"""


class HomeAssistantClient:
    def __init__(
        self,
        is_addon: bool,
        ha_url: str,
        session: aiohttp.ClientSession,
        db_queue: "Queue[tuple[str, dict[str, Any]]]",
    ):
        """Initialize the Home Assistant WebSocket client.

        Args:
            is_addon: Whether the client is running in an add-on
            ha_url: Home Assistant URL. Ignored if is_addon is True (add-on uses hardcoded URL)
            session: aiohttp.ClientSession
            db_queue: Queue to put the results in
        """
        self.supervisor_api_url: str | None = None
        if is_addon:
            self.core_api_url: str = "http://supervisor/core/api"
            self.supervisor_api_url: str = "http://supervisor/addons/self"
            self.ws_url: str = "ws://supervisor/core/websocket"
        else:
            protocol = "ws" if ha_url.startswith("http://") else "wss"
            host = ha_url.split("://")[1]
            self.core_api_url = f"{ha_url}/api"
            self.ws_url = f"{protocol}://{host}/api/websocket"
        self.message_id: int = 1
        self.websocket: Any | None = None
        self.aiohttp_session: aiohttp.ClientSession = session
        self.db_queue: "Queue[tuple[str, dict[str, Any]]]" = db_queue

    async def ws_connect(self) -> None:
        """Connect to Home Assistant and authenticate."""
        token = os.environ.get("SUPERVISOR_TOKEN")
        if not token:
            logger.warning("No token provided, cannot connect to Home Assistant")
            return

        self.websocket = await connect(self.ws_url)
        logger.debug(f"Connected to Home Assistant websocket: {self.ws_url}")

        # Wait for auth_required message
        auth_message = await self.ws_recv(timeout=5.0)
        if not auth_message:
            logger.error("No auth_message received from Home Assistant")
            return
        if auth_message["type"] != "auth_required":
            raise Exception(f"Expected auth_required, got {auth_message['type']}")

        # Send auth message
        await self.ws_send({"type": "auth", "access_token": token})

        # Wait for auth_ok message
        auth_result = await self.ws_recv(timeout=5.0)
        if not auth_result:
            logger.error("No auth_result received from Home Assistant")
            return
        if auth_result["type"] != "auth_ok":
            raise Exception(f"Expected auth_ok, got {auth_result['type']}")

        logger.info(f"Connected to Home Assistant {auth_result.get('ha_version', '')}")

    async def ws_send(self, data: dict[str, Any]) -> None:
        """Send JSON data to the websocket."""
        token = os.environ.get("SUPERVISOR_TOKEN")
        if self.websocket and token:
            try:
                await self.websocket.send(json.dumps(data))
            except websockets.exceptions.ConnectionClosed:
                logger.error("Websocket connection closed, attempting to reconnect")
                await self.ws_connect()
                await self.ws_send(data)
            except Exception as e:
                logger.error(f"Error sending data to websocket: {e}")
                return None

    async def ws_recv(self, timeout: float = 10.0) -> dict[str, Any] | None:
        """Receive JSON data from the websocket."""
        token = os.environ.get("SUPERVISOR_TOKEN")
        if self.websocket and token:
            try:
                raw_data = await asyncio.wait_for(
                    self.websocket.recv(), timeout=timeout
                )
                res = json.loads(raw_data)
                return res
            except asyncio.TimeoutError:
                logger.error(
                    f"Timeout waiting for websocket response (timeout: {timeout}s)"
                )
                return None
            except json.JSONDecodeError:
                logger.error("Error decoding JSON from websocket")
                return None
            except Exception as e:
                logger.error(f"Error receiving data from websocket: {e}")
                return None
        logger.error(
            "No websocket or token provided, cannot receive data from websocket"
        )
        return None

    async def ws_ping(self) -> bool:
        """Send ping to server and wait for pong."""
        msg_id = self.message_id
        self.message_id += 1

        await self.ws_send({"id": msg_id, "type": "ping"})

        # Wait for pong
        while True:
            result = await self.ws_recv(timeout=5.0)
            if not result:
                logger.error("No result received from websocket")
                return False
            if result.get("id") == msg_id and result.get("type") == "pong":
                return True
            else:
                logger.error(f"Unexpected response: {result}")
                return False

    async def dhcp(self) -> None:
        """Receive DHCP discovery data from Home Assistant."""
        msg_id = self.message_id
        self.message_id += 1

        await self.ws_send({"id": msg_id, "type": "dhcp/subscribe_discovery"})

        # Wait for result
        while True:
            result = await self.ws_recv(timeout=10.0)
            if not result:
                logger.error("No result received from websocket")
                break
            if result.get("id") == msg_id and result.get("type") == "event":
                try:
                    data = result.get("event", {}).get("add")
                    # logger.debug(f"Received DHCP event: {data}")
                    logger.info(
                        f"Received DHCP data from Home Assistant: {len(data)} devices"
                    )
                    self.db_queue.put(("ha_dhcp", data))
                except AttributeError:
                    logger.error(
                        "Error processing DHCP event: event {result} not formatted as expected"
                    )
                break

    async def rest_get_template(self, template: str) -> dict[str, Any] | None:
        logger.debug("Making request to Home Assistant API")
        token = os.environ.get("SUPERVISOR_TOKEN")
        if not token:
            logger.warning(
                "No token provided, cannot make request to Home Assistant API"
            )
            return None

        try:
            async with self.aiohttp_session.post(
                f"{self.core_api_url}/template",
                headers={"Authorization": f"Bearer {token}"},
                json={"template": template},
            ) as response:
                if response.status != 200:
                    logger.error(f"Error: API returned status code {response.status}")
                    response_text = await response.text()
                    logger.error(f"Response: {response_text}")
                    return None

                ha_data = await response.json(content_type="text/plain")
                return ha_data
        except json.JSONDecodeError as e:
            logger.warning(f"JSON decode error: {e}")
            return None
        except aiohttp.ClientResponseError as e:
            logger.warning(f"Client response error: {e}")
            return None
        except Exception as e:
            logger.warning(f"Error fetching Home Assistant data: {e}")
            return None

    async def supervisor_api_get_info(self) -> dict[str, Any] | None:
        """Get information about the Home Assistant supervisor."""
        token = os.environ.get("SUPERVISOR_TOKEN")
        if not token or not self.supervisor_api_url:
            logger.warning(
                "No token or supervisor API URL provided, cannot make request to Home Assistant supervisor API"
            )
            return None

        try:
            async with self.aiohttp_session.get(
                f"{self.supervisor_api_url}/info",
                headers={"Authorization": f"Bearer {token}"},
            ) as response:
                return (await response.json()).get("data", {})
        except Exception as e:
            logger.warning(f"Error fetching Home Assistant supervisor info: {e}")
            return None

    async def supervisor_api_post_config(
        self, data: dict[str, Any]
    ) -> dict[str, Any] | None:
        """Post data to the Home Assistant supervisor API."""
        token = os.environ.get("SUPERVISOR_TOKEN")
        if not token or not self.supervisor_api_url:
            logger.warning(
                "No token or supervisor API URL provided, cannot make request to Home Assistant supervisor API"
            )
            return None

        try:
            async with self.aiohttp_session.post(
                f"{self.supervisor_api_url}/options",
                headers={"Authorization": f"Bearer {token}"},
                json=data,
            ) as response:
                return (await response.json()).get("data", {})
        except Exception as e:
            logger.warning(f"Error posting data to Home Assistant supervisor API: {e}")
            return None

    async def supervisor_api_fix_config(self) -> dict[str, Any] | None:
        """Fix the Home Assistant configuration."""
        return await self.supervisor_api_post_config(
            {
                # "auto_update": True, # https://github.com/synergylabs/iot-transparency/issues/128
                "watchdog": True,
                "boot": "auto",
            }
        )

    async def get_and_save_devices(
        self,
    ) -> None:
        d = await self.rest_get_template(devices_template)
        if d:
            logger.info(
                f"Received REST API devices data: {len(d.get('devices', []))} devices"
            )
            self.db_queue.put(("ha_devices", d.get("devices", [])))
        else:
            logger.error("No response received from Home Assistant API")

    async def close(self) -> None:
        """Close the connection."""
        if self.websocket:
            await self.websocket.close()
            self.websocket = None