from typing import TypedDict, List

class ActProvision(TypedDict):
    number: str
    title: str
    sub_provs: List[str]
    history: List[str]

class Act(TypedDict):
    title: str
    front: str
    url: str
    provs: List[ActProvision]

from llama_index import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
)

# load the documents and create the index
documents = SimpleDirectoryReader("txt_data").load_data()
index = VectorStoreIndex.from_documents(documents)
# store it for later
index.storage_context.persist(persist_dir='./txt_storage')

query_engine = index.as_query_engine()
response = query_engine.query("When is it legal for a business to contract out of its obligations under the Fair Trading Act?")
print(response)

from typing import Any, Dict, List, Sequence
from llama_index.node_parser.interface import NodeParser
from llama_index.schema import BaseNode, TextNode, NodeRelationship
from llama_index.utils import get_tqdm_iterable

import tiktoken
import spacy

# Requirements:
# $ pip install llama-index tiktoken spacy 
# You will also need the spacy data:
# $ python -m spacy download en_core_web_sm 

# Example data
# ------------
# {
#   "summary": "This is a summary",
#   "sub_items": [
#       {
#           "title": "A title",
#           "contents": [
#               "Some list content",
#               "Some other list content"
#           ]
#       }
#   ]
# }

class JsonSplitter(NodeParser):
    max_split_tokens: int = 512

    @classmethod
    def class_name(cls) -> str:
        return "JsonSplitter"

    def _parse_nodes(self, nodes: Sequence[BaseNode], show_progress: bool = False, **kwargs: Any) -> List[BaseNode]:
        nodes_with_progress = get_tqdm_iterable(nodes, show_progress, "Parsing nodes")
        all_nodes: List[BaseNode] = []
        
        for node in nodes_with_progress:
            nodes = self.get_nodes_from_node(node)
            all_nodes.extend(nodes)
        
        return all_nodes

    def get_nodes_from_node(self, node: BaseNode) -> List[BaseNode]:
        encoder = tiktoken.Encoding = tiktoken.get_encoding("cl100k_base")
        nlp = spacy.load("en_core_web_sm")

        ret_nodes: List[BaseNode] = []

        your_data = json.loads(node.get_content())

        base_metadata = {}

        # Build a base node
        ret_nodes.append(self.build_text_node(
            text=your_data['summary'],
            parent_node=node,
            metadata={
                **base_metadata,
                "section": "root"
            }
        ))

        # Insert each sub item
        for sub_item in your_data['sub_items']:
            sub_item_metadata = {
                **base_metadata,
                "title": sub_item['title']
            }

            # Check if entire contents will fit into one chunk
            sub_prov_text = "\n".join(prov['contents'])
            if len(encoder.encode(sub_prov_text)) < self.max_split_tokens:
                # Insert the provision
                ret_nodes.append(self.build_text_node(
                    text=sub_prov_text,
                    parent_node=node,
                    metadata={**base_metadata, **prov_metadata},
                ))

                continue

            # Insert each sub-provision with respect to max_split_tokens
            for sub_prov in prov['contents']:
                token_count = len(encoder.encode(sub_prov))
                if token_count < self.max_split_tokens:
                    ret_nodes.append(self.build_text_node(
                        text=sub_prov,
                        parent_node=node,
                        metadata={**base_metadata, **prov_metadata},
                    ))
                else:
                    # Split into sentences and add one at a time till satisfied.
                    spacy_doc = nlp(sub_prov)
                    sentences = [sent.text.strip() for sent in spacy_doc.sents]
                    chunk = ""
                    for sent in sentences:
                        if len(encoder.encode(chunk + sent)) > self.max_split_tokens:
                            ret_nodes.append(self.build_text_node(
                                text=chunk,
                                parent_node=node,
                                metadata={**base_metadata, **prov_metadata},
                            ))
                            chunk = sent
                        else:
                            chunk += sent

                    if chunk != "":
                        ret_nodes.append(self.build_text_node(
                            text=chunk,
                            parent_node=node,
                            metadata={**base_metadata, **prov_metadata},
                        ))
                    
        return ret_nodes
    
    def build_text_node(self, text: str, parent_node: BaseNode, metadata: Dict[str,str] | None = None) -> TextNode:
        node = TextNode(
            text=text,
            excluded_embed_metadata_keys=(parent_node.excluded_embed_metadata_keys + self.exclude_metadata_fields),
            excluded_llm_metadata_keys=(parent_node.excluded_llm_metadata_keys + self.exclude_metadata_fields),
            relationships={NodeRelationship.SOURCE: parent_node.as_related_node_info()},
        )

        if metadata:
            node.metadata = {**node.metadata, **metadata}
        
        return node