from __future__ import annotations

import json
import logging
from typing import Callable
import websocket
import time

class HomeAssistantAPI:
    def __init__(self, token:str, initialize: Callable[[HomeAssistantAPI], None]) -> None:
        self.token = token
        self.msg_id = 1
        self.ws = None
        self.subscriptions = {}
        self.init_callback = initialize

    def handle_message(self, ws: websocket.WebSocket, msg: str) -> None:
        if self.ws is None:
            self.ws = ws

        message: object = json.loads(msg)
        
        if message["type"] == "auth_required":
            response = {
                "type": "auth",
                "access_token": self.token
            }
            logging.debug(response)
            ws.send(json.dumps(response))
            return
        elif message["type"] == "auth_invalid":
            logging.info("Auth failed")
            ws.close()
            return None
        elif message["type"] == "auth_ok":
            logging.debug("Authenticated")
            self.init_callback(self)
            self.init_callback = None
            return
        elif message["type"] == "event":
            event = message["event"]
            if event["event_type"] in self.subscriptions.keys():
                self.subscriptions[event["event_type"]](event)
                
        else:
            print("Received", message)
    
    def subscribe_event(self, event_type: str, callback: Callable[[object], None]):

        if self.ws is None:
            logging.debug("Websocket not set")
            return

        if event_type in self.subscriptions.keys():
            logging.warning(f"Already subscribed to {event_type}")
            return
        
        logging.info(f"Subscribe to {event_type}")
        self.subscriptions[event_type] = callback
        response = {
            "id": self.msg_id,
            "type": "subscribe_events",
            "event_type": event_type
        }
        self.msg_id += 1
        self.ws.send(json.dumps(response))