
from cag.framework.annotator import GenericAnnotator
from cag.framework.annotator.pipeline.pipeline_base import Pipeline
from cag.utils.config import *

from examples.graph_creation_example import AnyGraphCreator
from pyArango.query import AQLQuery

from pyArango.collection import Collection

class AnyAnnotator(GenericAnnotator):
    def __init__(self, conf: Config, params={'mode': 'run-1'}, filter_annotatable=True):
        super().__init__(query=f"""FOR dp IN {AnyGraphCreator._ANY_DATASET_NODE_NAME}
        RETURN dp
        """, params=params, conf=conf, filter_annotatable=filter_annotatable, run=True)

    def update_graph(self, timestamp, data: AQLQuery):
        for dp in data:
            pass


## Pipeline for specialized annotator

class SamplePipeline(Pipeline):

     def process_input(self) -> list:
        processed = []
        for txt_node in self.input:
            processed.append((txt_node.text, {"_key": txt_node._key}))

        return processed

     def init_and_run(self):
        # Extend the Default toml if needed, by either defining the path of a dictionary
         # Pipeline.extend_config(extened_dict)

        # Add the pipes defined in the Toml files
         self.add_annotation_pipe(name = "NamedEntityAnnotator", save_output= True, is_spacy=True)
         self.add_annotation_pipe(name = "EmpathAnnotator", save_output= True, is_spacy=True)
         self.add_annotation_pipe(name = "DummyDefaultAnnotator", save_output= False, is_spacy=False)

        # Get your Input in any way you like
         coll: Collection = self.database_config.db["TextNode"]
         docs = coll.fetchAll(limit=50)

        # Set the INPUT - this will automatically call process_input (make sure to implement it
         self.set_input(docs)

        # annotate the input
         self.set_spacy_language_model("en_core_web_md")
         self.annotate()
        # save annotations when "save_output" is set to True
         self.save()

def run_sample(config):
    sample_pipeline = SamplePipeline(config)
    sample_pipeline.init_and_run()
