import json
import threading
import time
from .closable_sseclient import ClosableSSEClient
from .keep_auth_session import KeepAuthSession


class Stream:
    def __init__(self, url, stream_handler, build_headers, stream_id, is_async, token_refreshable, token_refresher, max_retries):
        self.build_headers = build_headers
        self.url = url
        self.stream_handler = stream_handler
        self.stream_id = stream_id
        self.sse = None
        self.thread = None
        self.token_refresher = token_refresher
        self.token_refreshable = token_refreshable
        self.max_retries = max_retries

        if is_async:
            self.start()
        else:
            self.start_stream()

    def make_session(self):
        """
        Return a custom session object to be passed to the ClosableSSEClient.
        """
        session = KeepAuthSession()
        return session

    def start(self):
        self.thread = threading.Thread(target=self.start_stream)
        self.thread.start()
        return self

    def start_stream(self):
        self.sse = ClosableSSEClient(self.url, session=self.make_session(), build_headers=self.build_headers, token_refreshable=self.token_refreshable, token_refresher=self.token_refresher, max_retries=self.max_retries)
        for msg in self.sse:
            if msg:
                msg_data = json.loads(msg.data)
                msg_data["event"] = msg.event
                if self.stream_id:
                    msg_data["stream_id"] = self.stream_id
                self.stream_handler(msg_data)

    def close(self):
        while not self.sse and not hasattr(self.sse, 'resp'):
            time.sleep(0.001)
        self.sse.running = False
        self.sse.close()
        self.thread.join()
        return self