import threading
import queue
import tiktoken

from langchain.chat_models import ChatOpenAI
from langchain.callbacks.manager import CallbackManager
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

from openai.error import InvalidRequestError

class ThreadedGenerator:
    def __init__(self):
        self.queue = queue.Queue()

    def __iter__(self):
        return self

    def __next__(self):
        item = self.queue.get()
        if item is StopIteration: raise item
        return item

    def send(self, data):
        self.queue.put(data)

    def close(self):
        self.queue.put(StopIteration)

class ChainStreamHandler(StreamingStdOutCallbackHandler):
    def __init__(self, gen):
        super().__init__()
        self.gen = gen

    def on_llm_new_token(self, token: str, **kwargs):
        self.gen.send(token)

def run(g, prompts, retries = 0):
    try:
        chat = ChatOpenAI(
            verbose=True,
            streaming=True,
            callback_manager=CallbackManager([ChainStreamHandler(g)]),
        )
        chat(prompts)
    except InvalidRequestError as e:
        if retries >= 10:
            raise e
        else:
            retries += 1
            del prompts[-1]
            run(g, prompts, retries)

def llm_thread(g, prompts):
    encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
    messages = []
    tokens = 0

    for prompt in prompts:
        tokens += len(encoding.encode(prompt.content))
        if tokens >= 4096:
            break
        
        messages.append(prompt)

    try:
        run(g, messages)
    finally:
        g.close()


def chat(prompts):
    g = ThreadedGenerator()
    threading.Thread(target=llm_thread, args=(g, prompts)).start()
    return g