import os

from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import SemanticSplitterNodeParser
from llama_index.embeddings.ollama import OllamaEmbedding
import json

def generate_chunks_semantic(content):
    with open("./temp.txt", "wb") as file:
        file.write(content.encode("utf-8"))
        file.flush()
        file.close()

    embed_model = OllamaEmbedding(model_name="mxbai-embed-large")
    splitter = SemanticSplitterNodeParser(buffer_size=5, breakpoint_percentile_threshold=45, embed_model=embed_model)

    document = SimpleDirectoryReader(input_files=["./temp.txt"]).load_data()

    nodes = splitter.build_semantic_nodes_from_documents(document)

    os.remove("./temp.txt")

    output = []

    for i, node in enumerate(nodes):
        output.append(node.to_dict()["text"])

    return output

def generate_chunks_line_split(content):
    output = []

    sp = content.split(".")

    while len(sp) > 0:
        output.append("".join(sp[:10]))
        sp = sp[10:]

    return output


def generate_chunks(content):
    chunking_method = os.environ.get("CHUNKING_METHOD")
    print("CHUNKING_METHOD:", chunking_method)
    if chunking_method is None:
        chunking_method = "none"

    return {
        "semantic": lambda: generate_chunks_semantic(content),
        "lines": lambda: generate_chunks_line_split(content),
        "none": lambda: [content]
    }[chunking_method.lower().strip()]()