from transformers import TextIteratorStreamer
from threading import Thread
from utils.tree_utils import get_docstrings, grab_before_comments

def combine_generation_kwargs(temperature, max_new_tokens, top_p, repetition_penalty):
    """
    Combines the generation kwargs into a single dict.
    """
    gen_kwargs = {}
    gen_kwargs["temperature"] = temperature
    gen_kwargs["max_new_tokens"] = max_new_tokens
    gen_kwargs["top_p"] = top_p
    gen_kwargs["repetition_penalty"] = repetition_penalty
    return gen_kwargs


def stream_generation(prompt:str, pipe, gen_kwargs:dict):
    """
    Text generation function
    Args:
        prompt (str): The context to start generation from.
        pipe (Pipeline): The pipeline to use for generation.
        gen_kwargs (dict): The generation kwargs.
    Returns:
        str: The generated text. (it iterates over time)
    """
    # Tokenize the model_context
    model_inputs = pipe.tokenizer(prompt, return_tensors="pt")

    # Start generation on a separate thread, so that we don't block the UI. The text is pulled from the streamer
    # in the main thread. Adds timeout to the streamer to handle exceptions in the generation thread.
    streamer = TextIteratorStreamer(pipe.tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15.0)
    generate_kwargs = dict(model_inputs, streamer=streamer, **gen_kwargs)
    t = Thread(target=pipe.model.generate, kwargs=generate_kwargs)
    t.start()

    # Pull the generated text from the streamer, and update the model output.
    model_output = ""
    for new_text in streamer:
        # print("step", end="")
        model_output += new_text
        yield model_output
    streamer.on_finalized_text("stream reached the end.")
    return model_output #is this ever reached?

def construct_model_context(func_node, prompt="") -> str:
    """
    Constructs the model context from a function node.
    """
    model_context = func_node.child_by_field_name("type").text.decode() + " " + func_node.child_by_field_name("declarator").text.decode() #func_start_idx:body_start_idx?
    docstring = get_docstrings(func_node) #might be empty?
    if docstring:
        model_context = model_context + "\n" + docstring
    model_context = grab_before_comments(func_node) + model_context #prepend comments
    if prompt != "":
        model_context = "//Title: " + prompt + "\n" + model_context #prepend user prompt/title
        model_context = "//Language: Shadertoy GLSL fragment shader\n" + model_context #prepend system prompt, language hint
    return model_context