import random
import task_handler


def main():
    handler = task_handler.get_handler()
    config = handler.get_configuration()

    text = config["input_text"].lower()
    order = config["order"]
    tasks = config["tasks"]

    words = []
    word = ""
    for char in text:
        if char == " " or char == "\n":
            if word != "":
                words.append(word)
                word = ""
        else:
            word += char
    if word != "":
        words.append(word)

    model = build_model(words, order)
    results = []

    # First: print 10 lines starting with 'imagine' and 2 next words
    imagine_prefix = ("imagine",)
    if imagine_prefix in model:
        for _ in range(10):
            generated = generate_text(model, order, 3, imagine_prefix)
            results.append(" ".join(generated))
    else:
        for _ in range(10):
            results.append("imagine ??? ???")

    # Then: process the rest of the tasks in order from config
    for i, task in enumerate(tasks):
        count = task["generate_n_words"]
        prefix_str = task.get("prefix")

        prefix = None
        if prefix_str:
            prefix_words = []
            temp = ""
            for char in prefix_str.lower():
                if char == " ":
                    if temp != "":
                        prefix_words.append(temp)
                        temp = ""
                else:
                    temp += char
            if temp != "":
                prefix_words.append(temp)
            prefix = tuple(prefix_words)

        # skip the first task ("Imagine") since we manually handled it
        if i == 0:
            continue

        result = generate_text(model, order, count, prefix)
        results.append(" ".join(result))

    handler.write_output("\n".join(results).rstrip())


def build_model(words, order):
    model = {}
    for i in range(len(words) - order):
        key = tuple(words[i:i + order])
        value = words[i + order]
        if key not in model:
            model[key] = []
        model[key].append(value)
    return model


def generate_text(model, order, n_words, prefix=None):
    if prefix is None or len(prefix) != order or prefix not in model:
        prefix = random.choice(list(model.keys()))

    output = list(prefix)

    while len(output) < n_words:
        if prefix not in model:
            prefix = random.choice(list(model.keys()))
            output.extend(prefix)
            output = output[-n_words:]
            continue

        next_words = model[prefix]
        next_word = next_words[0] if len(next_words) == 1 else random.choice(next_words)
        output.append(next_word)
        prefix = tuple(output[-order:])

    return output[:n_words]