import base64
import hashlib
import io
import os
import uuid
from datetime import datetime
from typing import Callable, List

import rdflib
from naas_abi_core import logger
from naas_abi_core.services.ServiceBase import ServiceBase
from naas_abi_core.services.triple_store.TripleStorePorts import (
    ITripleStorePort,
    ITripleStoreService,
    OntologyEvent,
)
from rdflib import DCTERMS, OWL, RDF, RDFS, XSD, Graph, Literal, URIRef

SCHEMA_TTL = """
@prefix internal: <http://triple-store.internal#> .
@prefix owl: <http://www.w3.org/2002/07/owl#> .
@prefix rdf: <http://www.w3.org/1999/02/22-rdf-syntax-ns#> .
@prefix rdfs: <http://www.w3.org/2000/01/rdf-schema#> .
@prefix xsd: <http://www.w3.org/2001/XMLSchema#> .

internal:Schema a owl:Class ;
    rdfs:label "Schema" ;
    rdfs:comment "Represents a schema file that has been loaded into the triple store" .

internal:hash a owl:DatatypeProperty ;
    rdfs:domain internal:Schema ;
    rdfs:range xsd:string ;
    rdfs:label "hash" ;
    rdfs:comment "SHA-256 hash of the schema content" .

internal:fileLastUpdateTime a owl:DatatypeProperty ;
    rdfs:domain internal:Schema ;
    rdfs:range xsd:dateTime ;
    rdfs:label "file last update time" ;
    rdfs:comment "Last modification timestamp of the schema file" .

internal:filePath a owl:DatatypeProperty ;
    rdfs:domain internal:Schema ;
    rdfs:range xsd:string ;
    rdfs:label "file path" ;
    rdfs:comment "Path to the schema file" .

internal:content a owl:DatatypeProperty ;
    rdfs:domain internal:Schema ;
    rdfs:range xsd:base64Binary ;
    rdfs:label "content" ;
    rdfs:comment "Base64 encoded content of the schema file" .
"""
GRAPH_CLASS = URIRef("http://ontology.naas.ai/abi/Graph")


class TripleStoreService(ServiceBase, ITripleStoreService):
    """TripleStoreService provides CRUD operations and SPARQL querying capabilities for ontologies.

    This service acts as a facade for ontology storage and retrieval operations. It handles storing,
    retrieving, merging and querying of RDF ontologies while providing optional filtering of
    non-named individuals.

    Attributes:
        __triple_store_adapter (ITripleStorePort): The storage adapter implementation used for
            persisting and retrieving ontologies.

    Example:
        >>> store = TripleStoreService(FileSystemTripleStore("ontologies/"))
        >>> ontology = Graph()
        >>> # ... populate ontology ...
        >>> store.store("my_ontology", ontology)
        >>> results = store.query("SELECT ?s WHERE { ?s a owl:Class }")
    """

    def __init__(
        self,
        triple_store_adapter: ITripleStorePort,
    ):
        super().__init__()
        self.__triple_store_adapter = triple_store_adapter
        self.__schema_graph = URIRef("http://ontology.naas.ai/graph/schema")

        # Load SCHEMA_TTL in IOBuffer
        schema_ttl_buffer = io.StringIO(SCHEMA_TTL)
        self.insert(
            Graph().parse(schema_ttl_buffer, format="turtle"),
            graph_name=self.__schema_graph,
        )

    def _hash_value(self, value: object) -> str:
        return hashlib.sha256(str(value).encode("utf-8")).hexdigest()[:32]

    def _subscription_graph_topic_token(self, graph_name: URIRef | str) -> str:
        if graph_name == "*":
            return "*"
        return self._hash_value(graph_name)

    def insert(
        self,
        triples: Graph,
        graph_name: URIRef,
    ) -> None:

        assert graph_name is not None

        # Insert the triples into the store
        self.__triple_store_adapter.insert(triples, graph_name=graph_name)

        if self.services_wired is False:
            return

        # Notify listeners of the insert
        for s, p, o in triples.triples((None, None, None)):
            triple_bytes = f"{s.n3()} {p.n3()} {o.n3()} .\n".encode("utf-8")

            try:
                topic = f"ts.insert.g.{self._hash_value(str(graph_name))}.s.{self._hash_value(s)}.p.{self._hash_value(p)}.o.{self._hash_value(o)}"
                self.services.bus.topic_publish(
                    "triple_store",
                    topic,
                    triple_bytes,
                )
            except Exception as e:
                logger.error(f"Error publishing triple: {e}")
                raise e

    def remove(
        self,
        triples: Graph,
        graph_name: URIRef,
    ) -> None:
        assert graph_name is not None

        # Remove the triples from the store
        self.__triple_store_adapter.remove(triples, graph_name=graph_name)

        if self.services_wired is False:
            return

        # Notify listeners of the delete
        for s, p, o in triples.triples((None, None, None)):
            triple_bytes = f"{s.n3()} {p.n3()} {o.n3()} .\n".encode("utf-8")

            try:
                topic = f"ts.delete.g.{self._hash_value(str(graph_name))}.s.{self._hash_value(s)}.p.{self._hash_value(p)}.o.{self._hash_value(o)}"
                self.services.bus.topic_publish(
                    "triple_store",
                    topic,
                    triple_bytes,
                )
            except Exception as e:
                logger.error(f"Error publishing triple: {e}")
                raise e

    def get(self) -> Graph:
        return self.__triple_store_adapter.get()

    def query(self, query: str) -> rdflib.query.Result:
        return self.__triple_store_adapter.query(query)

    def query_view(self, view: str, query: str) -> rdflib.query.Result:
        return self.__triple_store_adapter.query_view(view, query)

    def _insert_graph_metadata(self, graph_name: URIRef) -> None:
        assert graph_name is not None
        assert isinstance(graph_name, URIRef)
        label = graph_name.split("/")[-1].split("#")[-1]
        g = Graph()
        g.add((graph_name, RDF.type, OWL.NamedIndividual))
        g.add((graph_name, RDF.type, GRAPH_CLASS))
        g.add((graph_name, RDFS.label, Literal(label)))
        g.add(
            (
                graph_name,
                DCTERMS.created,
                Literal(
                    datetime.now().strftime("%Y-%m-%dT%H:%M:%S%z"),
                    datatype=XSD.dateTime,
                ),
            )
        )
        self.insert(g, graph_name=graph_name)

    def create_graph(self, graph_name: URIRef) -> None:
        assert graph_name is not None
        assert isinstance(graph_name, URIRef)

        self.__triple_store_adapter.create_graph(graph_name)

        # Insert schema graph into the new graph
        self._insert_graph_metadata(graph_name)

    def clear_graph(self, graph_name: URIRef) -> None:
        assert graph_name is not None
        assert isinstance(graph_name, URIRef)

        self.__triple_store_adapter.clear_graph(graph_name)

        # Insert schema graph into the new graph
        self._insert_graph_metadata(graph_name)

    def drop_graph(self, graph_name: URIRef) -> None:
        self.__triple_store_adapter.drop_graph(graph_name)

    def list_graphs(self) -> list[URIRef]:
        return self.__triple_store_adapter.list_graphs()

    def subscribe(
        self,
        topic: tuple[URIRef | None, URIRef | None, URIRef | None],
        callback: Callable[[bytes], None],
        event_type: OntologyEvent | None = None,
        graph_name: URIRef | str = "*",
    ) -> None:
        _event_type: str = ""
        if event_type == OntologyEvent.INSERT:
            _event_type = "insert"
        elif event_type == OntologyEvent.DELETE:
            _event_type = "delete"
        elif event_type is None:
            _event_type = "*"

        graph_topic = self._subscription_graph_topic_token(graph_name)
        s = "*" if topic[0] is None else self._hash_value(topic[0])
        p = "*" if topic[1] is None else self._hash_value(topic[1])
        o = "*" if topic[2] is None else self._hash_value(topic[2])

        topic_str = f"ts.{_event_type}.g.{graph_topic}.s.{s}.p.{p}.o.{o}"

        self.services.bus.topic_consume(
            "triple_store",
            topic_str,
            callback,
        )

    def get_subject_graph(self, subject: str, graph_name: str = "*") -> Graph:
        return self.__triple_store_adapter.get_subject_graph(
            URIRef(subject), graph_name
        )

    ############################################################
    # Schema Management
    ############################################################

    def load_schemas(self, filepaths: List[str]):
        # First build a cache of all schemas to speed up the process.
        schema_cache = Graph()

        results = self.query(f"""
            PREFIX internal: <http://triple-store.internal#>
            SELECT ?schema ?filePath ?hash ?fileLastUpdateTime ?content
            WHERE {{
                GRAPH <{str(self.__schema_graph)}> {{
                ?schema a internal:Schema ;
                    internal:filePath ?filePath ;
                    internal:hash ?hash ;
                    internal:fileLastUpdateTime ?fileLastUpdateTime ;
                    internal:content ?content .
                }}
            }}
        """)

        for row in results:
            assert isinstance(row, rdflib.query.ResultRow)
            schema, filePath, hash, fileLastUpdateTime, content = row
            schema_cache.add(
                (schema, RDF.type, URIRef("http://triple-store.internal#Schema"))
            )
            schema_cache.add(
                (schema, URIRef("http://triple-store.internal#filePath"), filePath)
            )
            schema_cache.add(
                (schema, URIRef("http://triple-store.internal#hash"), hash)
            )
            schema_cache.add(
                (
                    schema,
                    URIRef("http://triple-store.internal#fileLastUpdateTime"),
                    fileLastUpdateTime,
                )
            )
            schema_cache.add(
                (schema, URIRef("http://triple-store.internal#content"), content)
            )

        for filepath in filepaths:
            self.load_schema(filepath, schema_cache)

    def load_schema(self, filepath: str, schema_cache: Graph | None = None):
        logger.debug(f"Loading schema: {filepath}")
        if schema_cache is not None:

            def _read_query_func(query: str):
                return schema_cache.query(query)

            read_query_func = _read_query_func
        else:
            read_query_func = self.query

        def _load_schema_metadata(subject: URIRef) -> dict[str, str]:
            triples: rdflib.query.Result = read_query_func(
                f"""
                PREFIX internal: <http://triple-store.internal#>
                SELECT ?p ?o 
                WHERE {{
                    GRAPH <{str(self.__schema_graph)}> {{
                        <{subject}> ?p ?o .
                    }}
                }}
                """
            )

            schema_dict: dict[str, str] = {}
            for row in triples:
                assert isinstance(row, rdflib.query.ResultRow)
                p, o = row
                schema_dict[str(p).replace("http://triple-store.internal#", "")] = str(
                    o
                )

            return schema_dict

        def _remove_schema_subject(subject: URIRef) -> None:
            triples: rdflib.query.Result = read_query_func(
                f"""
                SELECT ?p ?o 
                WHERE {{ 
                    GRAPH <{str(self.__schema_graph)}> {{ 
                    <{str(subject)}> ?p ?o . 
                    }} 
                }}
                """
            )

            cleanup_graph = Graph()
            for row in triples:
                assert isinstance(row, rdflib.query.ResultRow)
                p, o = row
                cleanup_graph.add((subject, p, o))

            if len(cleanup_graph) > 0:
                self.remove(cleanup_graph, graph_name=self.__schema_graph)

        try:
            query = f"""
            PREFIX internal: <http://triple-store.internal#>

            SELECT *
            WHERE {{
                GRAPH <{str(self.__schema_graph)}> {{
                    ?s internal:filePath "{filepath}" .
                }}
            }}
            """
            # logger.debug(f"Query: {query}")
            # Check if schema with filePath == filepath already exists and grab all triples.
            schema_triples: rdflib.query.Result = read_query_func(query)

            schema_rows = list(schema_triples)
            # logger.debug(f"len(schema_rows): {len(schema_rows)}")
            # If schema with filePath == filepath already exists, we check if the file has been modified.
            schema_exists_in_store = len(schema_rows) > 0
            # logger.debug(f"Schema exists in store: {schema_exists_in_store}")
            if schema_exists_in_store:
                # Open file and get content.
                with open(filepath, "r") as file:
                    new_content = file.read()

                new_content_hash = hashlib.sha256(
                    new_content.encode("utf-8")
                ).hexdigest()

                schema_entries: list[tuple[URIRef, dict[str, str]]] = []
                for row in schema_rows:
                    assert isinstance(row, rdflib.query.ResultRow)
                    _SUBJECT_TUPLE_INDEX = 0
                    subject = row[_SUBJECT_TUPLE_INDEX]
                    assert isinstance(subject, URIRef)
                    schema_entries.append((subject, _load_schema_metadata(subject)))

                matching_entry = next(
                    (
                        (subject, schema_dict)
                        for subject, schema_dict in schema_entries
                        if schema_dict.get("hash") == new_content_hash
                    ),
                    None,
                )

                if matching_entry is not None:
                    subject, schema_dict = matching_entry
                else:
                    subject, schema_dict = schema_entries[0]

                # Get file last update time
                file_last_update_time = os.path.getmtime(filepath)

                duplicate_subjects = [
                    candidate_subject
                    for candidate_subject, _ in schema_entries
                    if candidate_subject != subject
                ]

                # If fileLastUpdateTime is the same, return. Otherwise we continue as we need to update the schema.
                if schema_dict["hash"] == new_content_hash:
                    if len(duplicate_subjects) > 0:
                        logger.debug(
                            f"Cleaning up {len(duplicate_subjects)} duplicate schema metadata entries for {filepath}."
                        )
                        for duplicate_subject in duplicate_subjects:
                            _remove_schema_subject(duplicate_subject)

                    logger.debug("Schema is up to date, no need to update.")
                    return

                logger.debug("Schema is not up to date, updating.")

                # Decode old content
                old_content = base64.b64decode(schema_dict["content"]).decode("utf-8")

                # Parse old and new schema
                old_schema = Graph().parse(io.StringIO(old_content), format="turtle")

                new_schema = Graph().parse(io.StringIO(new_content), format="turtle")

                # Compute addition and deletion triples
                addition_triples = new_schema - old_schema
                deletion_triples = old_schema - new_schema

                # Insert addition and remove deletion triples
                self.insert(addition_triples, graph_name=self.__schema_graph)
                self.remove(deletion_triples, graph_name=self.__schema_graph)

                # Update schema information in the triple store.

                self.remove(
                    Graph().parse(
                        io.StringIO(f'''
                    @prefix internal: <http://triple-store.internal#> .
                    
                    <{subject}> internal:hash "{schema_dict["hash"]}" ;
                        internal:fileLastUpdateTime "{schema_dict["fileLastUpdateTime"]}" ;
                        internal:content "{schema_dict["content"]}" .
                '''),
                        format="turtle",
                    ),
                    graph_name=self.__schema_graph,
                )

                self.insert(
                    Graph().parse(
                        io.StringIO(f'''
                    @prefix internal: <http://triple-store.internal#> .
                    
                    <{subject}> internal:hash "{new_content_hash}" ;
                        internal:fileLastUpdateTime "{file_last_update_time}" ;
                        internal:content "{base64.b64encode(new_content.encode("utf-8")).decode("utf-8")}" .
                '''),
                        format="turtle",
                    ),
                    graph_name=self.__schema_graph,
                )

                if len(duplicate_subjects) > 0:
                    logger.debug(
                        f"Cleaning up {len(duplicate_subjects)} duplicate schema metadata entries for {filepath}."
                    )
                    for duplicate_subject in duplicate_subjects:
                        _remove_schema_subject(duplicate_subject)

                # Return as we don't need to continue as we have already updated the schema.
                return
            elif not schema_exists_in_store:
                logger.debug("Loading schema in graph as it doesn't exist in store.")

                # Open file and get content.
                with open(filepath, "r") as file:
                    content = file.read()

                # Compute base64 content
                base64_content = base64.b64encode(content.encode("utf-8")).decode(
                    "utf-8"
                )

                # Compute hash of content
                content_hash = hashlib.sha256(content.encode("utf-8")).hexdigest()

                # Parse schema
                g = Graph().parse(filepath)

                # Insert schema into the triple store
                self.insert(g, graph_name=self.__schema_graph)

                # Get file last update time
                file_last_update_time = os.path.getmtime(filepath)

                # Insert Schema with hash, filePath, fileLastUpdateTime and content to be able to track changes.
                self.insert(
                    Graph().parse(
                        io.StringIO(f'''
                    @prefix internal: <http://triple-store.internal#> .
                                                    
                    <http://triple-store.internal/{uuid.uuid4()}> a internal:Schema ;
                        internal:hash "{content_hash}" ;
                        internal:filePath "{filepath}" ;
                        internal:fileLastUpdateTime "{file_last_update_time}" ;
                        internal:content "{base64_content}" .
                '''),
                        format="turtle",
                    ),
                    graph_name=self.__schema_graph,
                )
        except Exception as e:
            import traceback

            logger.error(f"Error loading schema ({filepath}): {str(e)}")
            traceback.print_exc()

    def get_schema_graph(self) -> Graph:
        contents: rdflib.query.Result = self.query(
            f"""
            PREFIX internal: <http://triple-store.internal#>
            SELECT ?s ?o WHERE {{
                GRAPH <{str(self.__schema_graph)}> {{
                    ?s internal:content ?o .
                }}
            }}
            """
        )

        graph = Graph()

        for row in contents:
            assert isinstance(row, rdflib.query.ResultRow)
            _, o = row

            g = Graph().parse(
                io.StringIO(base64.b64decode(o).decode("utf-8")), format="turtle"
            )

            graph += g

            for prefix, namespace in g.namespaces():
                graph.bind(prefix, namespace)

        return graph
