from cog.database import Cog
from cog.database import in_nodes, out_nodes, hash_predicate, parse_tripple
from cog.core import cog_hash, Record
import json
import logging
from logging.config import dictConfig
from . import config as cfg
from cog.view import graph_template, script_part1, script_part2, graph_lib_src
import os
import shutil
from os import listdir
import time
import random
from math import isclose
import warnings
import heapq
import array
import math

# Optional simsimd for SIMD-optimized similarity
try:
    import simsimd
    _HAS_SIMSIMD = True
except ImportError:
    _HAS_SIMSIMD = False

NOTAG = "NOTAG"

# Sort direction constants for order()
ASC = "asc"
DESC = "desc"


class Vertex(object):

    def __init__(self, _id):
        self.id = _id
        self.tags = {}
        self.edges = set()
        self._path = None

    def set_edge(self, edge):
        self.edges.add(edge)
        return self

    def get_dict(self):
        return {k: v for k, v in self.__dict__.items() if k != '_path'}

    def __str__(self):
        return json.dumps(self.get_dict())


CHARS = u'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'


class BlankNode(object):
    ID_PREFIX = "_id_"

    def __init__(self, label=None):
        if not label:
            label = str(time.time_ns()) + ''.join(random.choices(CHARS, k=4))
            self.id = "_:{}".format(label)
        else:
            self.id = "_:{}{}".format(BlankNode.ID_PREFIX, label)

    def __str__(self):
        return self.id

    @classmethod
    def is_id(cls, label):
        return label.startswith("_:" + BlankNode.ID_PREFIX)


class Graph:
    """
    Creates a graph object.
    
    Args:
        graph_name: Name of the graph
        cog_home: Home directory name for the database
        cog_path_prefix: Root directory location for Cog db
        enable_caching: Enable in-memory caching for faster reads
        flush_interval: Number of writes before auto-flush per store.
                       1 = flush every write (safest, default)
                       0 = manual flush only (fastest, use sync())
                       N>1 = flush every N writes with async background threads
    """

    def __init__(self, graph_name, cog_home="cog_home", cog_path_prefix=None, enable_caching=True,
                 flush_interval=1):
        """
        :param graph_name:
        :param cog_home: Home directory name, for most use cases use default.
        :param cog_path_prefix: sets the root directory location for Cog db. Default: '/tmp' set in cog.Config. Change this to current directory when running in an IPython environment.
        :param flush_interval: Number of writes before auto-flush. 1 = every write (safest).
        """

        self.config = cfg
        self.config.COG_HOME = cog_home

        if cog_path_prefix:
            self.config.COG_PATH_PREFIX = cog_path_prefix

        self.graph_name = graph_name

        if enable_caching:
            self.cache = {}
        else:
            self.cache = None

        dictConfig(self.config.logging_config)
        self.logger = logging.getLogger("torque")

        self.logger.debug(f"Torque init on graph: {graph_name} (flush_interval={flush_interval})")

        self.cog = Cog(self.cache, flush_interval=flush_interval)
        self.cog.create_or_load_namespace(self.graph_name)

        self.all_predicates = self.cog.list_tables()
        self.views_dir = self.config.cog_views_dir()

        if not os.path.exists(self.views_dir):
            os.mkdir(self.views_dir)
        self.logger.debug("predicates: " + str(self.all_predicates))

        self.last_visited_vertices = None
        self._predicate_reverse_lookup_cache = {}  # hash -> human-readable predicate name
        # Hydrate predicate names from persisted edge set for reopened graphs
        try:
            self.cog.use_namespace(self.graph_name)
            for pred_hash in self.all_predicates:
                edge_record = self.cog.use_table(self.config.GRAPH_EDGE_SET_TABLE_NAME).get(pred_hash)
                if edge_record is not None:
                    self._predicate_reverse_lookup_cache[pred_hash] = edge_record.value
        except Exception:
            pass  # Edge set table may not exist yet for new graphs
        self._server_port = None  # Port this graph is being served on

    # === Network Methods ===
    
    def serve(self, port=8080, host="0.0.0.0", blocking=False, writable=False, share=False):
        """
        Start HTTP server for this graph instance.
        
        Multiple graphs can be served on the same port - each graph is accessible
        at /{graph_name}/ path. The index page at / lists all available graphs.
        
        Args:
            port: HTTP port to listen on (default 8080)
            host: Bind address (default "0.0.0.0" for all interfaces)
            blocking: If True, blocks forever (for dedicated servers)
            writable: If True, allows write operations via API
            share: If True, connect to CogDB relay
        
        Returns:
            self for method chaining
        
        Example:
            # Serve single graph
            g.serve(port=8080)
            
            # Serve multiple graphs on same port
            g1.serve(port=8080)
            g2.serve(port=8080)  # Both accessible at /{graph_name}/
            
            # Allow remote writes
            g.serve(port=8080, writable=True)
            
            # Share graph publicly
            g.serve(port=8080, share=True)
        """
        from cog.server import get_or_create_server
        
        if self._server_port is not None:
            raise RuntimeError(f"Graph '{self.graph_name}' already being served. Call stop() first.")
        
        # Get or create shared server on this port
        server, is_new = get_or_create_server(port, host)
        
        # Check if this graph name is already registered
        if server.has_graph(self.graph_name):
            raise RuntimeError(f"Graph '{self.graph_name}' already registered on port {port}")
        
        # Register this graph
        server.register_graph(self, writable=writable)
        self._server_port = port
        
        # Start share if requested
        if share:
            from cog.share import start_share
            share_info = start_share(port, local_host=host)
            share_url = share_info.wait()
            self._share_url = share_url
            self.logger.info(f"Share connected: {share_url}")
        
        # Start server if it's new
        if is_new:
            server.start(blocking=blocking)
            
            # If blocking mode exited, clean up
            if blocking:
                self._server_port = None
        
        return self
    
    def stop(self):
        """
        Stop serving this graph.
        
        If this is the last graph on the server, the server shuts down.
        
        Returns:
            self for method chaining
        """
        from cog.server import unregister_from_server
        
        if self._server_port is not None:
            unregister_from_server(self._server_port, self.graph_name)
            self._server_port = None
        return self
    
    def share_url(self):
        """
        Get the public share URL for this graph.
        
        Only available after calling serve(share=True).
        
        Returns:
            str: The share URL (e.g., "https://abc123.s.cogdb.io/") or None if not sharing
        
        Example:
            g.serve(port=8080, share=True)
            print(g.share_url())  # https://abc123.s.cogdb.io/
        """
        return getattr(self, '_share_url', None)
    
    @classmethod
    def connect(cls, url, timeout=30):
        """
        Connect to a remote CogDB server.
        
        Args:
            url: HTTP(S) URL including graph name path
                 (e.g., "http://localhost:8080/my_graph")
            timeout: Request timeout in seconds (default 30)
        
        Returns:
            RemoteGraph instance that can be used like a local Graph
        
        Example:
            remote = Graph.connect("http://192.168.1.5:8080/social")
            remote.v("alice").out("knows").all()
            
            # Via ngrok
            remote = Graph.connect("https://abc123.ngrok.io/my_graph")
        """
        from cog.remote import RemoteGraph
        return RemoteGraph(url, timeout=timeout)

    def sync(self):
        """
        Force flush all pending writes to disk.
        Blocks until all flushes are complete.
        
        Use this when flush_interval > 1 or when you need to ensure 
        data durability at a specific point.
        """
        self.cog.sync()

    def refresh(self):
        self.cog.refresh_all()

    def updatej(self, json_object):
        self.put_json(json_object, True)

    def putj(self, json_object, update=False):
        """
        Shorthand for put_json
        :param update:
        :param json_object:
        :return: None
        """
        self.put_json(json_object, update)

    def put_json(self, json_object, update=False):
        """
        Experimental Feature
        ====================

        Inserts a JSON object into the graph. Each object (including the root object) in this JSON object will be
        identified by a BlankNode with a unique label. For example: {"name" : "bob", "location" : { "city" :
        "Toronto", "country" : "Canada"} } will be transformed into the following triples:

        _:1654006783197959000lIxa, name, bob
        _:1654006783197959000lIxa, location, _:1654006783844002000kAgC
        _:1654006783844002000kAgC, city, toronto
        _:1654006783844002000kAgC, country, canada

        """
        if isinstance(json_object, str):
            json_object = json.loads(json_object)
        self._traverse_json(json_object, update)

    def _traverse_json(self, jsn, update=False):
        new_edge_created = set()

        def traverse(json_obj, subject, predicate=None, update_object=False, in_list=False):

            if type(json_obj) is dict:
                # every object has an id
                if "_id" in json_obj:
                    if in_list and update_object:
                        raise Exception("Updating a sub object or list item with an _id is not supported.")
                    child_id = str(BlankNode(json_obj["_id"]))
                else:
                    # if _id is not present generate one.
                    child_id = str(BlankNode())
                if predicate:
                    # this is to skip the first iteration where predicate is None.
                    # For items inside a list, always add (don't replace) even during updates
                    effective_update = update_object and not in_list
                    self.put(subject, predicate, child_id, effective_update)
                for a in json_obj:
                    # Properties of a dict are NOT in a list - reset in_list to False
                    traverse(json_obj[a], child_id, a, update_object, in_list=False)

            elif type(json_obj) is list:
                # create a new blank node for each list.
                list_id = str(BlankNode())
                self.put(subject, predicate, list_id, update_object)
                # new_edge_created.add((str(subject), str(predicate)))

                # traverse the list - mark items as being in a list
                for obj in json_obj:
                    traverse(obj, list_id, predicate, update_object, in_list=True)

            else:
                # For items inside a list, always add (don't replace) even during updates
                effective_update = update_object and not in_list
                if (str(subject), str(predicate)) in new_edge_created:
                    self.put(subject, predicate, json_obj, effective_update)
                else:
                    self.put(subject, predicate, json_obj, effective_update)
                    new_edge_created.add((str(subject), str(predicate)))

        if "_id" in jsn:
            traverse(jsn, str(BlankNode(jsn["_id"])), update_object=update)
        else:
            traverse(jsn, str(BlankNode()), update_object=update)

    def load_triples(self, graph_data_path, graph_name=None):
        """
        Loads triples from a file (one triple per line) into a graph.

        :param graph_data_path:
        :param graph_name:
        :return: None
        """

        graph_name = self.graph_name if graph_name is None else graph_name
        self.cog.load_triples(graph_data_path, graph_name)
        self.all_predicates = self.cog.list_tables()
        # Rebuild _predicate_reverse_lookup_cache by parsing the triples file
        with open(graph_data_path) as f:
            for line in f:
                _, predicate, _, _ = parse_tripple(line)
                self._predicate_reverse_lookup_cache[hash_predicate(predicate)] = predicate
        return None

    def load_csv(self, csv_path, id_column_name, graph_name=None):
        """
        Loads a CSV file to a graph. One column must be designated as ID column. This method is intended for loading
        simple CSV data, for more complex ones that require additional logic, convert the CSV to triples using custom
        logic.

        :param csv_path:
        :param id_column_name:
        :param graph_name:
        :return: None
        """

        if id_column_name is None:
            raise Exception("id_column_name must not be None")
        graph_name = self.graph_name if graph_name is None else graph_name
        self.cog.load_csv(csv_path, id_column_name, graph_name)
        self.all_predicates = self.cog.list_tables()
        # Rebuild _predicate_reverse_lookup_cache from CSV column headers
        import csv
        with open(csv_path) as csv_file:
            reader = csv.DictReader(csv_file)
            if reader.fieldnames:
                for col in reader.fieldnames:
                    self._predicate_reverse_lookup_cache[hash_predicate(col)] = col

    def close(self):
        self.logger.info("closing graph: " + self.graph_name)
        self.cog.close()

    def put(self, vertex1, predicate, vertex2, update=False, create_new_edge=False):
        self._predicate_reverse_lookup_cache[hash_predicate(predicate)] = predicate
        self.cog.use_namespace(self.graph_name)
        if update:
            if create_new_edge:
                self.cog.put_new_edge(vertex1, predicate, vertex2)
            else:
                self.cog.update_edge(vertex1, predicate, vertex2)
        else:
            self.cog.put_node(vertex1, predicate, vertex2)
        self.all_predicates = self.cog.list_tables()
        return self

    def put_batch(self, triples):
        """
        Insert multiple triples efficiently using batch mode.
        Significantly faster than calling put() in a loop for large datasets.
        
        :param triples: List of (vertex1, predicate, vertex2) tuples
        :return: self for method chaining
        
        Example:
            g.put_batch([
                ("alice", "follows", "bob"),
                ("bob", "follows", "charlie"),
                ("charlie", "follows", "alice")
            ])
        """
        self.cog.use_namespace(self.graph_name)
        self.cog.begin_batch()
        try:
            for v1, pred, v2 in triples:
                self._predicate_reverse_lookup_cache[hash_predicate(pred)] = pred
                self.cog.put_node(v1, pred, v2)
        finally:
            self.cog.end_batch()
        self.all_predicates = self.cog.list_tables()
        return self

    def delete(self, vertex1, predicate, vertex2):
        """
        Removes a specific triple/edge from the graph.
        
        :param vertex1: Source vertex
        :param predicate: Edge predicate/relationship
        :param vertex2: Target vertex
        :return: self for method chaining
        
        Example:
            g.put("alice", "knows", "bob")
            g.delete("alice", "knows", "bob")
        """
        self.cog.delete_edge(vertex1, predicate, vertex2)
        return self

    def drop(self, *args):
        """
        Deletes the entire graph and its persistent storage from disk.
        
        WARNING: This is a destructive operation that cannot be undone.
        The graph object becomes unusable after this call.
        
        :return: None
        
        Example:
            g.drop()  # Deletes entire graph from disk
        """
        if len(args) > 0:
            raise DeprecationWarning(
                "drop(s, p, o) is deprecated. Use delete(s, p, o) for edges. "
                "Use drop() with no arguments to delete the entire graph."
            )
        
        # Clear the cache
        if self.cache is not None:
            self.cache.clear()
        
        # Unregister from server if currently served
        self.stop()
        
        # Close the graph
        self.close()
        
        # Delete the entire graph directory using the same method Cog uses
        graph_path = self.config.cog_data_dir(self.graph_name)
        if os.path.exists(graph_path):
            shutil.rmtree(graph_path)

    def truncate(self):
        """
        Wipes all triples but keeps the graph structure/directory intact.
        
        Useful for resetting state without needing to re-initialize.
        The graph remains usable after this call.
        
        :return: self for method chaining
        
        Example:
            g.put("alice", "knows", "bob")
            g.truncate()  # Graph is now empty but still usable
            g.put("new", "data", "here")  # Works fine
        """
        # Get the graph directory path using the same method Cog uses
        # This correctly handles CUSTOM_COG_DB_PATH if set
        graph_path = self.config.cog_data_dir(self.graph_name)
        
        # Save flush_interval before closing (access while cog is still open)
        flush_interval = self.cog.flush_interval
        
        # Close current connections
        self.cog.close()
        
        try:
            # Delete all contents but keep the directory
            if os.path.exists(graph_path):
                for item in os.listdir(graph_path):
                    item_path = os.path.join(graph_path, item)
                    if os.path.isdir(item_path):
                        shutil.rmtree(item_path)
                    else:
                        os.remove(item_path)
            
            # Clear the cache to prevent stale data
            if self.cache is not None:
                self.cache.clear()
        finally:
            # Re-initialize the graph to ensure the object remains usable,
            # even if some files could not be deleted.
            self.cog = Cog(self.cache, flush_interval=flush_interval)
            self.cog.create_or_load_namespace(self.graph_name)
            self.all_predicates = self.cog.list_tables()
        
        return self


    def update(self, vertex1, predicate, vertex2):
        self.updatej(vertex1, predicate, vertex2)
        return self

    def v(self, vertex=None, func=None):
        if func:
            warnings.warn("The use of func is deprecated, please use filter instead.", DeprecationWarning)
        if vertex is not None:
            if isinstance(vertex, list):
                self.last_visited_vertices = [Vertex(v) for v in vertex]
            else:
                self.last_visited_vertices = [Vertex(vertex)]
        else:
            self.last_visited_vertices = []
            self.cog.use_namespace(self.graph_name).use_table(self.config.GRAPH_NODE_SET_TABLE_NAME)
            for r in self.cog.scanner():
                if func is not None and not func(r.key):
                    continue
                self.last_visited_vertices.append(Vertex(r.key))
        return self

    def out(self, predicates=None, func=None):
        """
        Traverse forward through edges.
        :param func:
        :param predicates: A string or a List of strings.
        :return: self for method chaining.
        """

        if func:
            warnings.warn("The use of func is deprecated, please use filter instead.", DeprecationWarning)
            assert callable(func), "func must be a lambda. Example: func = lambda d: int(d) > 5"
            assert not isinstance(predicates, list), "func cannot be used with a list of predicates"

        if predicates is not None:
            if not isinstance(predicates, list):
                predicates = [predicates]
            predicates = list(map(hash_predicate, predicates))
        else:
            predicates = self.all_predicates

        self.logger.debug("OUT: predicates: " + str(predicates))
        self.__hop("out", predicates=predicates, func=func)
        return self

    def inc(self, predicates=None, func=None):
        """
        Traverse backward through edges.
        :param predicates: List of predicates
        :return: self for method chaining.
        """

        if func:
            warnings.warn("The use of func is deprecated, please use filter instead.", DeprecationWarning)
            assert callable(func), "func must be a lambda. Example: func = lambda d: int(d) > 5"
            assert not isinstance(predicates, list), "func cannot be used with a list of predicates"

        if predicates is not None:
            if not isinstance(predicates, list):
                predicates = [predicates]
            predicates = list(map(hash_predicate, predicates))
        else:
            predicates = self.all_predicates

        self.__hop("in", predicates, func=func)
        return self

    def __adjacent_vertices(self, vertex, predicates, direction='out'):
        self.cog.use_namespace(self.graph_name)
        adjacent_vertices = []
        for predicate in predicates:
            if direction == 'out':
                out_record = self.cog.use_table(predicate).get(out_nodes(vertex.id))
                if out_record is not None:
                    for v_adj in out_record.value:
                        adjacent_vertices.append(Vertex(v_adj).set_edge(predicate))
            elif direction == 'in':
                in_record = self.cog.use_table(predicate).get(in_nodes(vertex.id))
                if in_record is not None:
                    for v_adj in in_record.value:
                        adjacent_vertices.append(Vertex(v_adj).set_edge(predicate))

        return adjacent_vertices

    def has(self, predicates, vertex):
        """
        Filters all outgoing edges from a vertex that matches a list of predicates.
        :param predicates: List of predicates
        :param vertex: Vertex ID
        :return: self for method chaining.
        """

        if predicates is not None:
            if not isinstance(predicates, list):
                predicates = [predicates]
            predicates = list(map(hash_predicate, predicates))

        has_vertices = []
        for lv in self.last_visited_vertices:
            adj_vertices = self.__adjacent_vertices(lv, predicates)
            for av in adj_vertices:
                if av.id == vertex:
                    has_vertices.append(lv)

        self.last_visited_vertices = has_vertices
        return self

    def hasr(self, predicates, vertex):
        """
        'Has' in reverse. Filters all incoming edges from a vertex that matches a list of predicates.
        :param predicates: List of predicates
        :param vertex: Vertex ID
        :return: self for method chaining.
        """

        if predicates is not None:
            if not isinstance(predicates, list):
                predicates = [predicates]
            predicates = list(map(hash_predicate, predicates))

        has_vertices = []
        for lv in self.last_visited_vertices:
            adj_vertices = self.__adjacent_vertices(lv, predicates, 'in')
            for av in adj_vertices:
                if av.id == vertex:
                    has_vertices.append(lv)

        self.last_visited_vertices = has_vertices
        return self

    def scan(self, limit=10, scan_type='v'):
        """
        Scan vertices or edges in the current graph namespace and return their IDs.
        :param limit: Maximum number of items to return from the scan
        :param scan_type: use 'v' to scan the vertex set or 'e' to scan the edge set
        :return: A dictionary containing a list of scanned item(vertex) IDs, e.g., `{'result': [{'id': '...'}]}`.
        """
        assert type(scan_type) is str, "Scan type must be either 'v' for vertices or 'e' for edges."
        if scan_type == 'e':
            self.cog.use_namespace(self.graph_name).use_table(self.config.GRAPH_EDGE_SET_TABLE_NAME)
        else:
            self.cog.use_namespace(self.graph_name).use_table(self.config.GRAPH_NODE_SET_TABLE_NAME)
        result = []
        for i, r in enumerate(self.cog.scanner()):
            if i < limit:
                if scan_type == 'v':
                    v = Vertex(r.key)
                else:
                    v = Vertex(r.value)
                result.append({"id": v.id})
            else:
                break
        return {"result": result}

    def __hop(self, direction, predicates=None, func=None):
        self.logger.debug("__hop : direction: " + str(direction) + " predicates: " + str(
            predicates) + " graph name: " + self.graph_name)
        self.cog.use_namespace(self.graph_name)
        self.logger.debug("hopping from vertices: " + str(map(lambda x: x.id, self.last_visited_vertices)))
        self.logger.debug("direction: " + str(direction) + " predicates: " + str(self.all_predicates))
        traverse_vertex = []
        for predicate in predicates:
            self.logger.debug("__hop predicate: " + predicate + " of " + str(predicates))
            for v in self.last_visited_vertices:
                if direction == "out":
                    record = self.cog.use_table(predicate).get(out_nodes(v.id))
                else:
                    record = self.cog.use_table(predicate).get(in_nodes(v.id))
                if record is not None:
                    if record.value_type == "s":
                        v_adjacent = str(record.value)
                        if func is not None and not func(v_adjacent):
                            continue
                        v_adjacent_obj = Vertex(v_adjacent).set_edge(predicate)
                        v_adjacent_obj.tags.update(v.tags)
                        parent_path = v._path or [{'vertex': v.id}]
                        v_adjacent_obj._path = parent_path + [
                            {'edge': self._predicate_reverse_lookup_cache.get(predicate, predicate)},
                            {'vertex': v_adjacent}
                        ]
                        traverse_vertex.append(v_adjacent_obj)
                    elif record.value_type == "l":
                        for v_adjacent in record.value:
                            self.logger.debug("record v: " + str(record.value) + " type: " + str(record.value_type))
                            if func is not None and not func(v_adjacent):
                                continue
                            v_adjacent_obj = Vertex(v_adjacent).set_edge(predicate)
                            v_adjacent_obj.tags.update(v.tags)
                            parent_path = v._path or [{'vertex': v.id}]
                            v_adjacent_obj._path = parent_path + [
                                {'edge': self._predicate_reverse_lookup_cache.get(predicate, predicate)},
                                {'vertex': v_adjacent}
                            ]
                            traverse_vertex.append(v_adjacent_obj)
        self.last_visited_vertices = traverse_vertex

    def filter(self, func):
        """
            Applies a filter function to the vertices and removes any vertices that do not pass the filter.
        """
        self.last_visited_vertices = [v for v in self.last_visited_vertices if func(v.id)]
        return self

    def both(self, predicates=None):
        """
        Traverse edges in both directions (out + in).
        :param predicates: A string or list of predicate strings to follow.
        :return: self for method chaining.
        """
        if predicates is not None:
            if not isinstance(predicates, list):
                predicates = [predicates]
            predicates = list(map(hash_predicate, predicates))
        else:
            predicates = self.all_predicates

        self.cog.use_namespace(self.graph_name)
        traverse_vertex = []

        for predicate in predicates:
            for v in self.last_visited_vertices:
                # Outgoing edges
                out_record = self.cog.use_table(predicate).get(out_nodes(v.id))
                if out_record is not None:
                    if out_record.value_type == "s":
                        v_adjacent = str(out_record.value)
                        v_adj = Vertex(v_adjacent).set_edge(predicate)
                        v_adj.tags.update(v.tags)
                        parent_path = v._path or [{'vertex': v.id}]
                        v_adj._path = parent_path + [
                            {'edge': self._predicate_reverse_lookup_cache.get(predicate, predicate)},
                            {'vertex': v_adjacent}
                        ]
                        traverse_vertex.append(v_adj)
                    elif out_record.value_type == "l":
                        for v_adjacent in out_record.value:
                            v_adj = Vertex(v_adjacent).set_edge(predicate)
                            v_adj.tags.update(v.tags)
                            parent_path = v._path or [{'vertex': v.id}]
                            v_adj._path = parent_path + [
                                {'edge': self._predicate_reverse_lookup_cache.get(predicate, predicate)},
                                {'vertex': v_adjacent}
                            ]
                            traverse_vertex.append(v_adj)

                # Incoming edges
                in_record = self.cog.use_table(predicate).get(in_nodes(v.id))
                if in_record is not None:
                    if in_record.value_type == "s":
                        v_adjacent = str(in_record.value)
                        v_adj = Vertex(v_adjacent).set_edge(predicate)
                        v_adj.tags.update(v.tags)
                        parent_path = v._path or [{'vertex': v.id}]
                        v_adj._path = parent_path + [
                            {'edge': self._predicate_reverse_lookup_cache.get(predicate, predicate)},
                            {'vertex': v_adjacent}
                        ]
                        traverse_vertex.append(v_adj)
                    elif in_record.value_type == "l":
                        for v_adjacent in in_record.value:
                            v_adj = Vertex(v_adjacent).set_edge(predicate)
                            v_adj.tags.update(v.tags)
                            parent_path = v._path or [{'vertex': v.id}]
                            v_adj._path = parent_path + [
                                {'edge': self._predicate_reverse_lookup_cache.get(predicate, predicate)},
                                {'vertex': v_adjacent}
                            ]
                            traverse_vertex.append(v_adj)

        self.last_visited_vertices = traverse_vertex
        return self

    def is_(self, *nodes):
        """
        Filter paths to only those currently at the specified node(s).
        :param nodes: One or more node IDs to filter to.
        :return: self for method chaining.
        """
        if len(nodes) == 1 and isinstance(nodes[0], list):
            node_set = set(nodes[0])
        else:
            node_set = set(nodes)
        self.last_visited_vertices = [v for v in self.last_visited_vertices if v.id in node_set]
        return self

    def unique(self):
        """
        Remove duplicate vertices from the result set.
        :return: self for method chaining.
        """
        seen = set()
        unique_vertices = []
        for v in self.last_visited_vertices:
            if v.id not in seen:
                seen.add(v.id)
                unique_vertices.append(v)
        self.last_visited_vertices = unique_vertices
        return self

    def limit(self, n):
        """
        Limit results to the first N vertices.
        :param n: Maximum number of vertices to return.
        :return: self for method chaining.
        """
        self.last_visited_vertices = self.last_visited_vertices[:n]
        return self

    def skip(self, n):
        """
        Skip the first N vertices in the result set.
        :param n: Number of vertices to skip.
        :return: self for method chaining.
        """
        self.last_visited_vertices = self.last_visited_vertices[n:]
        return self

    def order(self, direction="asc"):
        '''
        Sort vertices by their id.
        :param direction: Sort direction - "asc" (default) or "desc".
        :return: self for method chaining.

        Example:
            g.v("Person").out("created_at").order().all()        # ascending (default)
            g.v("Person").out("created_at").order(desc).all()    # descending
        '''
        reverse = (direction == "desc")
        self.last_visited_vertices = sorted(self.last_visited_vertices, key=lambda v: v.id, reverse=reverse)
        return self

    def back(self, tag):
        """
        Return to vertices saved at the given tag, preserving all constraints.
        :param tag: A previous tag in the query to jump back to.
        :return: self for method chaining.
        """
        vertices = []
        for v in self.last_visited_vertices:
            if tag in v.tags:
                tagged_vertex = Vertex(v.tags[tag])
                tagged_vertex.tags = v.tags.copy()
                tagged_vertex.edges = v.edges.copy()
                tagged_vertex._path = v._path
                vertices.append(tagged_vertex)
        self.last_visited_vertices = vertices
        return self

    def __get_adjacent(self, vertex, predicates, direction):
        """Get adjacent vertices based on direction: 'out', 'inc', or 'both'."""
        adjacent = []
        if direction in ("out", "both"):
            adjacent.extend(self.__adjacent_vertices(vertex, predicates, 'out'))
        if direction in ("inc", "both"):
            adjacent.extend(self.__adjacent_vertices(vertex, predicates, 'in'))
        return adjacent

    def bfs(self, predicates=None, max_depth=None, min_depth=0,
            direction="out", until=None, unique=True):
        """
        Traverse the graph breadth-first from current vertices.

        BFS explores level-by-level, visiting all neighbors at the current depth
        before moving deeper. Guarantees shortest path in unweighted graphs.

        :param predicates: Edge type(s) to follow: str, list, or None (all edges)
        :param max_depth: Maximum traversal depth (None = unlimited)
        :param min_depth: Minimum depth to include in results (default 0)
        :param direction: Traversal direction: "out", "inc", or "both"
        :param until: Stop condition lambda: func(vertex_id) -> bool
        :param unique: If True, visit each vertex only once (prevents cycles)
        :return: self for method chaining

        Example:
            g.v("alice").bfs(predicates="follows", max_depth=2).all()
            g.v("alice").bfs(max_depth=3, min_depth=2).all()  # depths 2-3 only
            g.v("alice").bfs(until=lambda v: v == "target").all()
        """
        from collections import deque

        # Normalize predicates
        if predicates is not None:
            if not isinstance(predicates, list):
                predicates = [predicates]
            predicates = list(map(hash_predicate, predicates))
        else:
            predicates = self.all_predicates

        result_vertices = []
        visited = set()
        queue = deque()  # (vertex, depth)

        # Initialize with current vertices at depth 0
        for v in self.last_visited_vertices:
            queue.append((v, 0))
            if unique:
                visited.add(v.id)

        while queue:
            current, depth = queue.popleft()

            if until and until(current.id):
                if depth >= min_depth:
                    result_vertex = Vertex(current.id)
                    result_vertex.tags = current.tags.copy()
                    result_vertex.edges = current.edges.copy()
                    result_vertex._path = current._path
                    result_vertices.append(result_vertex)
                continue

            if depth > 0 and depth >= min_depth:
                if max_depth is None or depth <= max_depth:
                    result_vertex = Vertex(current.id)
                    result_vertex.tags = current.tags.copy()
                    result_vertex.edges = current.edges.copy()
                    result_vertex._path = current._path
                    result_vertices.append(result_vertex)

            # Stop exploring if at max depth
            if max_depth is not None and depth >= max_depth:
                continue

            adjacent = self.__get_adjacent(current, predicates, direction)
            for adj in adjacent:
                if unique:
                    if adj.id in visited:
                        continue
                    visited.add(adj.id)
                adj.tags = current.tags.copy()
                # Build path for neighbor from parent's path
                parent_path = current._path or [{'vertex': current.id}]
                edge_hash = next(iter(adj.edges)) if adj.edges else None
                edge_name = self._predicate_reverse_lookup_cache.get(edge_hash, edge_hash) if edge_hash else None
                adj._path = list(parent_path) + ([{'edge': edge_name}] if edge_name else []) + [{'vertex': adj.id}]
                queue.append((adj, depth + 1))

        self.last_visited_vertices = result_vertices
        return self

    def dfs(self, predicates=None, max_depth=None, min_depth=0,
            direction="out", until=None, unique=True):
        """
        Traverse the graph depth-first from current vertices.

        DFS explores as deep as possible along each branch before backtracking.
        More memory-efficient than BFS for deep graphs.

        :param predicates: Edge type(s) to follow: str, list, or None (all edges)
        :param max_depth: Maximum traversal depth (None = unlimited)
        :param min_depth: Minimum depth to include in results (default 0)
        :param direction: Traversal direction: "out", "inc", or "both"
        :param until: Stop condition lambda: func(vertex_id) -> bool
        :param unique: If True, visit each vertex only once (prevents cycles)
        :return: self for method chaining

        Example:
            g.v("alice").dfs(predicates="follows", max_depth=3).all()
            g.v("alice").dfs(direction="both", max_depth=2).all()
        """
        # Normalize predicates
        if predicates is not None:
            if not isinstance(predicates, list):
                predicates = [predicates]
            predicates = list(map(hash_predicate, predicates))
        else:
            predicates = self.all_predicates

        result_vertices = []
        visited = set()
        stack = []  # (vertex, depth)

        # Initialize with current vertices at depth 0
        for v in self.last_visited_vertices:
            stack.append((v, 0))
            if unique:
                visited.add(v.id)

        while stack:
            current, depth = stack.pop()  # LIFO for DFS

            if until and until(current.id):
                if depth >= min_depth:
                    result_vertex = Vertex(current.id)
                    result_vertex.tags = current.tags.copy()
                    result_vertex.edges = current.edges.copy()
                    result_vertex._path = current._path
                    result_vertices.append(result_vertex)
                continue

            if depth > 0 and depth >= min_depth:
                if max_depth is None or depth <= max_depth:
                    result_vertex = Vertex(current.id)
                    result_vertex.tags = current.tags.copy()
                    result_vertex.edges = current.edges.copy()
                    result_vertex._path = current._path
                    result_vertices.append(result_vertex)

            # Stop exploring if at max depth
            if max_depth is not None and depth >= max_depth:
                continue

            adjacent = self.__get_adjacent(current, predicates, direction)
            for adj in adjacent:
                if unique:
                    if adj.id in visited:
                        continue
                    visited.add(adj.id)
                adj.tags = current.tags.copy()
                # Build path for neighbor from parent's path
                parent_path = current._path or [{'vertex': current.id}]
                edge_hash = next(iter(adj.edges)) if adj.edges else None
                edge_name = self._predicate_reverse_lookup_cache.get(edge_hash, edge_hash) if edge_hash else None
                adj._path = list(parent_path) + ([{'edge': edge_name}] if edge_name else []) + [{'vertex': adj.id}]
                stack.append((adj, depth + 1))

        self.last_visited_vertices = result_vertices
        return self

    def sim(self, word, operator, threshold, strict=False):
        """
            Applies cosine similarity filter to the vertices and removes any vertices that do not pass the filter.

            Parameters:
            -----------
            word: str
                The word to compare to the other vertices.
            operator: str
                The comparison operator to use. One of "==", ">", "<", ">=", "<=", or "in".
            threshold: float or list of 2 floats
                The threshold value(s) to use for the comparison. If operator is "==", ">", "<", ">=", or "<=", threshold should be a float. If operator is "in", threshold should be a list of 2 floats.
            strict: bool, optional
                If True, raises an exception if a word embedding is not found for either word. If False, assigns a similarity of 0.0 to any word embedding that is not found.

            Returns:
            --------
            self: GraphTraversal
                Returns self to allow for method chaining.

            Raises:
            -------
            ValueError:
                If operator is not a valid comparison operator or if threshold is not a valid threshold value for the given operator.
                If strict is True and a word embedding is not found for either word.
    """
        if not isinstance(threshold, (float, int, list)):
            raise ValueError("Invalid threshold value: {}".format(threshold))

        if operator == 'in':
            if not isinstance(threshold, list) or len(threshold) != 2:
                raise ValueError("Invalid threshold value: {}".format(threshold))
            if not all(isinstance(t, (float, int)) for t in threshold):
                raise ValueError("Invalid threshold value: {}".format(threshold))

        filtered_vertices = []
        for v in self.last_visited_vertices:
            similarity = self.__cosine_similarity(word, v.id)
            if not similarity:
                # similarity is None if a word embedding is not found for either word.
                if strict:
                    raise ValueError("Missing word embedding for either '{}' or '{}'".format(word, v.id))
                else:
                    # Treat vertices without word embeddings as if they have no similarity to any other vertex.
                    similarity = 0.0
            if operator == '=':
                if isclose(similarity, threshold):
                    filtered_vertices.append(v)
            elif operator == '>':
                if similarity > threshold:
                    filtered_vertices.append(v)
            elif operator == '<':
                if similarity < threshold:
                    filtered_vertices.append(v)
            elif operator == '>=':
                if similarity >= threshold:
                    filtered_vertices.append(v)
            elif operator == '<=':
                if similarity <= threshold:
                    filtered_vertices.append(v)
            elif operator == 'in':
                if not threshold[0] <= similarity <= threshold[1]:
                    continue
                filtered_vertices.append(v)
            else:
                raise ValueError("Invalid operator: {}".format(operator))
        self.last_visited_vertices = filtered_vertices
        return self

    def __cosine_distance(self, x, y):
        """Compute cosine distance (1 - similarity) with simsimd or pure Python fallback."""
        if _HAS_SIMSIMD:
            return simsimd.cosine(x, y)
        else:
            # Pure Python fallback for Pyodide/environments without simsimd
            dot = sum(a * b for a, b in zip(x, y))
            norm_x = math.sqrt(sum(a * a for a in x))
            norm_y = math.sqrt(sum(b * b for b in y))
            if norm_x == 0 or norm_y == 0:
                return 1.0  # Max distance if either vector is zero
            return 1.0 - (dot / (norm_x * norm_y))

    def __cosine_similarity(self, word1, word2):
        """Compute cosine similarity using SIMD-optimized simsimd library or pure Python fallback."""
        x_list = self.get_embedding(word1)
        y_list = self.get_embedding(word2)

        if x_list is None or y_list is None:
            return None
        
        # Use python array for buffer protocol (compatible with simsimd)
        x = array.array('f', x_list)
        y = array.array('f', y_list)

        # cosine distance = 1 - similarity, so we convert
        distance = self.__cosine_distance(x, y)
        return 1.0 - float(distance)

    def tag(self, tag_names):
        """
        Saves vertices with tag name(s). Used to capture vertices while traversing a graph.
        
        Tag names are validated to prevent XSS when used in graph views.
        
        :param tag_names: A string or list of strings.
        :return: self for method chaining.
        """
        if not isinstance(tag_names, list):
            tag_names = [tag_names]
        for v in self.last_visited_vertices:
            for tag_name in tag_names:
                # Validate tag name to prevent XSS in view output
                if not isinstance(tag_name, str):
                    raise TypeError("Tag names must be strings")
                v.tags[tag_name] = v.id
        return self

    def count(self):
        return len(self.last_visited_vertices)

    def all(self, options=None):
        """
        Returns all the vertices that are resultant of the graph query. Options 'e' would include the edges that were traversed.
        https://github.com/cayleygraph/cayley/blob/master/docs/GizmoAPI.md
        :return:
        """
        result = []
        show_edge = True if options is not None and 'e' in options else False
        for v in self.last_visited_vertices:
            item = {"id": v.id}
            if show_edge and v.edges:
                item['edges'] = [
                    self.cog.use_namespace(self.graph_name).use_table(self.config.GRAPH_EDGE_SET_TABLE_NAME).get(
                        edge).value for edge in v.edges]
            # item['edge'] = self.cog.use_namespace(self.graph_name).use_table(self.config.GRAPH_EDGE_SET_TABLE_NAME).get(item['edge']).value
            item.update(v.tags)

            result.append(item)
        res = {"result": result}
        return res

    def graph(self):
        """
        Returns graph structure ready for D3.js / vis.js visualization.
        Deduplicates nodes and edges from all traversal paths.

        Returns:
            dict with 'nodes' (list of {id}) and 'links' (list of {source, target, label})

        Example:
            g.v("bob").out("follows").tag("from").out("works_at").tag("to").graph()
            # {'nodes': [{'id': 'bob'}, ...], 'links': [{'source': 'bob', 'target': 'fred', 'label': 'follows'}, ...]}
        """
        nodes = {}
        links = {}

        for v in self.last_visited_vertices:
            path = v._path or []
            for i, step in enumerate(path):
                if 'vertex' in step:
                    nodes[step['vertex']] = {'id': step['vertex']}
                if 'edge' in step and i > 0 and i < len(path) - 1:
                    src = path[i - 1]['vertex']
                    tgt = path[i + 1]['vertex']
                    key = (src, step['edge'], tgt)
                    links[key] = {
                        'source': src,
                        'target': tgt,
                        'label': step['edge']
                    }

        return {
            'nodes': list(nodes.values()),
            'links': list(links.values())
        }

    def view(self, view_name,
             js_src="https://cdnjs.cloudflare.com/ajax/libs/vis-network/9.1.2/dist/vis-network.min.js"):
        """
            Returns html view of the resulting graph from a query.
            :return:
        """
        assert view_name is not None, "a view name is required to create a view, it can be any string."
        result = self.graph()
        # Escape HTML special characters to prevent XSS when embedding in script tag
        # Replace < with \u003c to prevent </script> injection
        safe_json = json.dumps(result).replace('<', '\\u003c').replace('>', '\\u003e')
        view_html = script_part1 + graph_lib_src.format(js_src=js_src) + graph_template.format(
            plot_data_insert=safe_json) + script_part2
        view = self.views_dir + "/{view_name}.html".format(view_name=view_name)
        view = View(view, view_html)
        view.persist()
        return view

    def getv(self, view_name):
        view = self.views_dir + "/{view_name}.html".format(view_name=view_name)
        assert os.path.isfile(view), "view not found, create a view by calling .view()"
        with open(view, 'r') as f:
            view_html = f.read()
        view = View(view, view_html)
        return view

    def lsv(self):
        return [f.split(".")[0] for f in listdir(self.views_dir)]

    def get_new_graph_instance(self):
        return Graph(self.graph_name, self.config.COG_HOME, self.config.COG_PATH_PREFIX)

    def put_embedding(self, word, embedding):
        """
        Saves a word embedding.
        """

        assert isinstance(word, str), "word must be a string"
        self.cog.use_namespace(self.graph_name).use_table(self.config.EMBEDDING_SET_TABLE_NAME).put(Record(
            word, embedding))

    def get_embedding(self, word):
        """
        Returns a word embedding.
        """
        assert isinstance(word, str), "word must be a string"
        record = self.cog.use_namespace(self.graph_name).use_table(self.config.EMBEDDING_SET_TABLE_NAME).get(
            word)
        if record is None:
            return None
        return record.value

    def delete_embedding(self, word):
        """
        Deletes a word embedding.
        """
        assert isinstance(word, str), "word must be a string"
        self.cog.use_namespace(self.graph_name).use_table(self.config.EMBEDDING_SET_TABLE_NAME).delete(
            word)

    def put_embeddings_batch(self, word_embedding_pairs):
        """
        Bulk insert multiple embeddings efficiently.
        
        :param word_embedding_pairs: List of (word, embedding) tuples
        :return: self for method chaining
        
        Example:
            g.put_embeddings_batch([
                ("apple", [0.1, 0.2, ...]),
                ("orange", [0.3, 0.4, ...]),
            ])
        """
        self.cog.use_namespace(self.graph_name)
        self.cog.begin_batch()
        try:
            for word, embedding in word_embedding_pairs:
                if not isinstance(word, str):
                    raise TypeError("word must be a string")
                self.cog.use_table(self.config.EMBEDDING_SET_TABLE_NAME).put(Record(
                    word, embedding))
        finally:
            self.cog.end_batch()
        return self

    def scan_embeddings(self, limit=100):
        """
        Scan and return a list of words that have embeddings stored.
        
        :param limit: Maximum number of embeddings to return
        :return: Dictionary with 'result' containing list of words with embeddings
        
        Note: This scans the graph vertices and checks which have embeddings.
        """
        result = []
        self.cog.use_namespace(self.graph_name).use_table(self.config.GRAPH_NODE_SET_TABLE_NAME)
        count = 0
        for r in self.cog.scanner():
            if count >= limit:
                break
            word = r.key
            if self.get_embedding(word) is not None:
                result.append({"id": word})
                count += 1
        return {"result": result}

    def embedding_stats(self):
        """
        Return statistics about stored embeddings.
        
        :return: Dictionary with count and dimensions (if available)
        """
        count = 0
        dimensions = None
        # Scan the embedding table directly
        self.cog.use_namespace(self.graph_name).use_table(self.config.EMBEDDING_SET_TABLE_NAME)
        for r in self.cog.scanner():
            count += 1
            if dimensions is None and r.value is not None:
                dimensions = len(r.value)
        return {"count": count, "dimensions": dimensions}

    def k_nearest(self, word, k=10):
        """
        Find the k vertices most similar to the given word based on embeddings.
        
        :param word: The word to find similar vertices for
        :param k: Number of nearest neighbors to return (default 10)
        :return: self for method chaining
        
        Example:
            g.v().k_nearest("machine_learning", k=5).all()
        """
        target_embedding = self.get_embedding(word)
        if target_embedding is None:
            self.last_visited_vertices = []
            return self
        
        # simsimd/fallback requires buffer protocol (e.g. numpy array or python array)
        target_vec = array.array('f', target_embedding)
        similarities = []
        
        # None = no prior traversal, scan entire embedding table
        # [] = prior traversal returned empty, preserve empty semantics
        # [...] = search within visited vertices
        if self.last_visited_vertices is None:
            # Scan embedding table directly for all embeddings
            self.cog.use_namespace(self.graph_name).use_table(self.config.EMBEDDING_SET_TABLE_NAME)
            for r in self.cog.scanner():
                if r.value is not None:
                    v_vec = array.array('f', r.value)
                    distance = self.__cosine_distance(target_vec, v_vec)
                    similarity = 1.0 - float(distance)
                    similarities.append((similarity, Vertex(r.key)))
        elif self.last_visited_vertices:
            # Search within visited vertices
            for v in self.last_visited_vertices:
                v_embedding = self.get_embedding(v.id)
                if v_embedding is not None:
                    v_vec = array.array('f', v_embedding)
                    distance = self.__cosine_distance(target_vec, v_vec)
                    similarity = 1.0 - float(distance)
                    similarities.append((similarity, v))
        # else: empty list, similarities stays empty
        
        # Get top k using heap for efficiency
        top_k = heapq.nlargest(k, similarities, key=lambda x: x[0])
        self.last_visited_vertices = [v for _, v in top_k]
        return self

    def load_glove(self, filepath, limit=None, batch_size=1000):
        """
        Load GloVe embeddings from a text file.
        
        :param filepath: Path to GloVe file (e.g., 'glove.6B.100d.txt')
        :param limit: Maximum number of embeddings to load (None for all)
        :param batch_size: Number of embeddings to batch before writing (default 1000)
        :return: Number of embeddings loaded
        
        Example:
            count = g.load_glove("glove.6B.100d.txt", limit=50000)
        """
        count = 0
        batch = []
        
        with open(filepath, 'r', encoding='utf-8') as f:
            for line in f:
                if limit is not None and count >= limit:
                    break
                parts = line.strip().split()
                if len(parts) < 2:
                    continue
                word = parts[0]
                embedding = [float(x) for x in parts[1:]]
                batch.append((word, embedding))
                count += 1
                
                if len(batch) >= batch_size:
                    self.put_embeddings_batch(batch)
                    batch = []
        
        # Load remaining batch
        if batch:
            self.put_embeddings_batch(batch)
        
        return count

    def load_gensim(self, model, limit=None, batch_size=1000):
        """
        Load embeddings from a Gensim Word2Vec or FastText model.
        
        :param model: A Gensim model with a 'wv' attribute (Word2Vec, FastText)
        :param limit: Maximum number of embeddings to load (None for all)
        :param batch_size: Number of embeddings to batch before writing (default 1000)
        :return: Number of embeddings loaded
        
        Example:
            from gensim.models import Word2Vec
            model = Word2Vec(sentences)
            count = g.load_gensim(model)
        """
        count = 0
        batch = []
        
        # Get word vectors from model
        if hasattr(model, 'wv'):
            wv = model.wv
        else:
            wv = model  # Already a KeyedVectors object
        
        for word in wv.index_to_key:
            if limit is not None and count >= limit:
                break
            embedding = wv[word].tolist()
            batch.append((word, embedding))
            count += 1
            
            if len(batch) >= batch_size:
                self.put_embeddings_batch(batch)
                batch = []
        
        if batch:
            self.put_embeddings_batch(batch)
        
        return count


class View(object):

    def __init__(self, url, html):
        self.url = url
        self.html = html

    def render(self, height=700, width=700):
        """
        :param self:
        :param height:
        :param width:
        :return:
        """
        iframe_html = r"""  <iframe srcdoc='{0}' width="{1}" height="{2}"> </iframe> """.format(self.html, width,
                                                                                                height)
        from IPython.core.display import display, HTML
        display(HTML(iframe_html))

    def persist(self):
        f = open(self.url, "w")
        f.write(self.html)
        f.close()

    def __str__(self):
        return self.url
