import json
import threading
import typing
from typing import Optional, Dict, Any, Callable, Union
import logging

import websocket
import websockets
from .events import EventType

class ListenConnection:
    def __init__(self, url: str, headers: Dict[str, str]):
        self.url = url
        self.headers = headers
        self._ws: Optional[websocket.WebSocket] = None
        self._handlers: Dict[EventType, list[Callable[[Any], None]]] = {
            EventType.OPEN: [],
            EventType.MESSAGE: [],
            EventType.CLOSE: [],
            EventType.ERROR: []
        }
        self._stop_event = threading.Event()

    def on(self, event: EventType, handler: Callable[[Any], None]):
        if event in self._handlers:
            self._handlers[event].append(handler)

    def _trigger(self, event: EventType, data: Any = None):
        for handler in self._handlers[event]:
            try:
                if data is not None:
                    handler(data)
                else:
                    handler(None) # Some handlers might expect an argument
            except Exception as e:
                print(f"Error in handler for {event}: {e}")

    def connect(self):
        # Header formatting for websocket-client: list of strings "Key: Value"
        header_list = [f"{k}: {v}" for k, v in self.headers.items()]
        self._ws = websocket.create_connection(self.url, header=header_list)
        self._trigger(EventType.OPEN, self)
        return self

    def close(self):
        self._stop_event.set()
        if self._ws:
            try:
                self._ws.send(json.dumps({"type": "Finalize"}))
                self._ws.send(json.dumps({"type": "CloseStream"}))
            except Exception:
                pass
            try:
                self._ws.close()
            except Exception:
                pass
        self._trigger(EventType.CLOSE, self)

    def send(self, data: Union[bytes, str]):
        if not self._ws:
            raise Exception("Connection not open")
        if isinstance(data, bytes):
            self._ws.send(data, opcode=websocket.ABNF.OPCODE_BINARY)
        else:
            self._ws.send(data)

    def start_listening(self):
        if not self._ws:
            raise Exception("Connection not open")
        
        while not self._stop_event.is_set():
            try:
                message = self._ws.recv()
                if message is None:
                    break
                
                # Parse message if it's JSON
                parsed_msg = message
                try:
                    if isinstance(message, (bytes, bytearray)):
                        text = message.decode("utf-8")
                    else:
                        text = message
                    parsed_msg = json.loads(text)
                except (json.JSONDecodeError, UnicodeDecodeError):
                    pass # Pass raw message if not JSON
                
                # Attach 'type' attribute if it's a dict, to match Deepgram usage
                # user example: msg_type = getattr(message, "type", "Unknown")
                # So the message should differ? 
                # Deepgram returns objects, but here we might just return dicts.
                # However, user example uses: getattr(message, "type", ...)
                # If message is a dict, getattr won't work.
                # I should probably wrap it in a SimpleNamespace or object if it's a dict.
                
                final_msg = parsed_msg
                if isinstance(parsed_msg, dict):
                     # Wrap in a simple class to support getattr like user example
                     class MessageObj:
                         def __init__(self, data):
                             self.__dict__.update(data)
                             self.raw = data
                         def __getattr__(self, name):
                             return self.__dict__.get(name)
                         def __repr__(self):
                             return str(self.raw)
                     
                     final_msg = MessageObj(parsed_msg)

                self._trigger(EventType.MESSAGE, final_msg)

            except Exception as e:
                # websocket._exceptions.WebSocketConnectionClosedException
                if self._stop_event.is_set():
                    break
                self._trigger(EventType.ERROR, e)
                break
        
        self.close()

    def __enter__(self):
        return self.connect()

    def __exit__(self, exc_type, exc_value, traceback):
        self.close()


class AsyncListenConnection:
    def __init__(self, url: str, headers: Dict[str, str]):
        self.url = url
        self.headers = headers
        self._ws = None
        self._handlers: Dict[EventType, list[Callable[[Any], None]]] = {
            EventType.OPEN: [],
            EventType.MESSAGE: [],
            EventType.CLOSE: [],
            EventType.ERROR: []
        }
        self._running = False

    def on(self, event: EventType, handler: Callable[[Any], None]):
        if event in self._handlers:
            self._handlers[event].append(handler)

    def _trigger(self, event: EventType, data: Any = None):
        for handler in self._handlers[event]:
            try:
                if data is not None:
                   # For async, handlers might be synchronous or asynchronous?
                   # Deepgram example shows synchronous lambda: lambda _: print(...)
                   handler(data)
                else:
                   handler(None)
            except Exception as e:
                print(f"Error in handler for {event}: {e}")

    async def connect(self):
        print(f">> connecting to {self.url} with {self.headers}")
        self._ws = await websockets.connect(self.url, additional_headers=self.headers)
        self._running = True
        self._trigger(EventType.OPEN, self)
        return self

    async def close(self):
        self._running = False
        if self._ws:
            try:
                await self._ws.send(json.dumps({"type": "Finalize"}))
                await self._ws.send(json.dumps({"type": "CloseStream"}))
            except Exception:
                pass
            try:
                await self._ws.close()
            except Exception:
                pass
        self._trigger(EventType.CLOSE, self)

    async def send(self, data: Union[bytes, str]):
        if not self._ws:
            raise Exception("Connection not open")
        await self._ws.send(data)

    async def start_listening(self):
        if not self._ws:
            raise Exception("Connection not open")
        
        try:
            async for message in self._ws:
                if not self._running:
                    break
                
                parsed_msg = message
                try:
                    text = message
                    if isinstance(message, bytes):
                        text = message.decode("utf-8")
                    parsed_msg = json.loads(text)
                except (json.JSONDecodeError, UnicodeDecodeError):
                    pass

                final_msg = parsed_msg
                if isinstance(parsed_msg, dict):
                     class MessageObj:
                         def __init__(self, data):
                             self.__dict__.update(data)
                             self.raw = data
                         def __getattr__(self, name):
                             return self.__dict__.get(name)
                         def __repr__(self):
                             return str(self.raw)
                     final_msg = MessageObj(parsed_msg)

                self._trigger(EventType.MESSAGE, final_msg)
        except Exception as e:
            self._trigger(EventType.ERROR, e)
        finally:
            if self._running:
                await self.close()

    async def __aenter__(self):
        return await self.connect()

    async def __aexit__(self, exc_type, exc_value, traceback):
        await self.close()
