import rdflib
from rdflib import Graph, Literal, RDF, URIRef
from rdflib.graph import ConjunctiveGraph
from rdflib.plugins.sparql import prepareQuery
from rdflib.plugins.sparql.evaluate import evalPart, evalBGP
from rdflib.plugins.sparql.sparql import SPARQLError
from rdflib.plugins.sparql.evalutils import _eval
from rdflib.plugins.sparql.processor import translateQuery as translateQuery
# from rdflib.plugins.sparql.algebra import pprintAlgebra

from fastapi import FastAPI, Request, Response, Query
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from starlette.responses import FileResponse 
from typing import Optional

import pkg_resources
import logging
import re
import os
from urllib import parse

class SparqlEndpoint(FastAPI):
    """
    Class to deploy a SPARQL endpoint using a RDFLib Graph
    """

    def __init__(self,
            *args,
            title: str = "SPARQL endpoint for RDFLib graph", 
            description="A SPARQL endpoint to serve machine learning models, or any other logic implemented in Python. \n[Source code](https://github.com/vemonet/rdflib-endpoint)",
            version="0.1.0",
            graph=ConjunctiveGraph(), 
            functions={},
            enable_update=False,
            cors_enabled=True,
            public_url='https://sparql.openpredict.semanticscience.org/sparql',
            example_query="""Example query:\n
```
PREFIX myfunctions: <https://w3id.org/um/sparql-functions/>
SELECT ?concat ?concatLength WHERE {
    BIND("First" AS ?first)
    BIND(myfunctions:custom_concat(?first, "last") AS ?concat)
}
```""", 
            **kwargs
        ) -> None:
        """
        Constructor of the SPARQL endpoint, everything happens here.
        FastAPI calls are defined in this constructor
        """
        self.graph=graph
        self.functions=functions
        self.title=title
        self.description=description
        self.version=version
        self.public_url=public_url
        self.example_query=example_query
        self.enable_update=enable_update

        # Instantiate FastAPI
        super().__init__(
            title=title, description=description, version=version, 
        )

        # Save custom function in custom evaluation dictionary
        # Handle multiple functions directly in the evalCustomFunctions function
        rdflib.plugins.sparql.CUSTOM_EVALS['evalCustomFunctions'] = self.evalCustomFunctions
        
        if cors_enabled:
            self.add_middleware(
                CORSMiddleware,
                allow_origins=["*"],
                allow_credentials=True,
                allow_methods=["*"],
                allow_headers=["*"],
            )

        api_responses={
            200: {
                "description": "SPARQL query results",
                "content": {
                    "application/sparql-results+json": {
                        "results": {"bindings": []}, 'head': {'vars': []}
                    },
                    "application/json": {
                        "results": {"bindings": []}, 'head': {'vars': []}
                    },
                    "text/csv": {
                        "example": "s,p,o"
                    },
                    "application/sparql-results+csv": {
                        "example": "s,p,o"
                    },
                    "text/turtle": {
                        "example": "service description"
                    },
                    "application/sparql-results+xml": {
                        "example": "<root></root>"
                    },
                    "application/xml": {
                        "example": "<root></root>"
                    },
                    # "application/rdf+xml": {
                    #     "example": '<?xml version="1.0" encoding="UTF-8"?> <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"></rdf:RDF>'
                    # },
                },
            },
            400:{
                "description": "Bad Request",
            },
            403:{
                "description": "Forbidden",
            }, 
            422:{
                "description": "Unprocessable Entity",
            },
        }

        mimetype={
            'turtle': 'text/turtle',
            'xml_results': 'application/sparql-results+xml'
        }

        @self.get(
            "/sparql",
            name="SPARQL endpoint",
            description=self.example_query,
            responses=api_responses
        )
        async def sparql_endpoint(
            request: Request,
            query: Optional[str] = Query(None)):
            """
            Send a SPARQL query to be executed through HTTP GET operation.
            \f
            :param request: The HTTP GET request
            :param query: SPARQL query input.
            """
            if not query:
                # Return the SPARQL endpoint service description
                service_graph = Graph()
                # service_graph.parse('app/service-description.ttl', format="ttl")
                service_graph.parse(data=service_description_ttl, format="ttl")

                # Add custom functions URI to the service description
                for custom_function_uri in self.functions.keys():
                    service_graph.add((URIRef(custom_function_uri), RDF.type, URIRef('http://www.w3.org/ns/sparql-service-description#Function')))
                    service_graph.add((URIRef(self.public_url), URIRef('http://www.w3.org/ns/sparql-service-description#extensionFunction'), URIRef(custom_function_uri)))

                # Return the service description RDF as turtle or XML
                if request.headers['accept'] == mimetype['turtle']:
                    return Response(service_graph.serialize(format = 'turtle'), media_type=mimetype['turtle'])
                else:
                    return Response(service_graph.serialize(format = 'xml'), media_type='application/xml')

            # Pretty print the query object 
            # parsed_query = parser.parseQuery(query)
            # tq = algebraTranslateQuery(parsed_query)
            # pprintAlgebra(tq)

            try:
                # Query the graph with the custom functions loaded
                parsed_query = prepareQuery(query)
                query_operation = re.sub(r"(\w)([A-Z])", r"\1 \2", parsed_query.algebra.name)
            except Exception as e:
                logging.error("Error parsing the SPARQL query: " + str(e))
                return JSONResponse(status_code=400, content={"message": "Error parsing the SPARQL query: " + str(e)})

            # Useless: RDFLib dont support SPARQL insert (Expected {SelectQuery | ConstructQuery | DescribeQuery | AskQuery}, found 'INSERT')
            # if not self.enable_update:
            #     if query_operation == "Insert Query" or query_operation == "Delete Query":
            #         return JSONResponse(status_code=403, content={"message": "INSERT and DELETE queries are not allowed."})
            # if os.getenv('RDFLIB_APIKEY') and (query_operation == "Insert Query" or query_operation == "Delete Query"):
            #     if apikey != os.getenv('RDFLIB_APIKEY'):
            #         return JSONResponse(status_code=403, content={"message": "Wrong API KEY."})

            try:
                query_results = self.graph.query(parsed_query)
            except Exception as e:
                logging.error("Error executing the SPARQL query on the RDFLib Graph: " + str(e))
                return JSONResponse(status_code=400, content={"message": "Error executing the SPARQL query on the RDFLib Graph: " + str(e)})

            # Format and return results depending on Accept mime type in request header
            output_mime_type = request.headers['accept']
            if not output_mime_type:
                output_mime_type = 'application/xml'
            
            # Handle mime type for construct queries
            if query_operation == "Construct Query" and (output_mime_type == 'application/json' or output_mime_type == 'text/csv'):
                output_mime_type = mimetype['turtle']
                # TODO: support JSON-LD for construct query?
                # g.serialize(format='json-ld', indent=4)
            if query_operation == "Construct Query" and output_mime_type == 'application/xml':
                output_mime_type = 'application/rdf+xml'

            try:
                if output_mime_type == 'text/csv' or output_mime_type == 'application/sparql-results+csv':
                    return Response(query_results.serialize(format = 'csv'), media_type=output_mime_type)
                elif output_mime_type == 'application/json' or output_mime_type == 'application/sparql-results+json':
                    return Response(query_results.serialize(format = 'json'), media_type=output_mime_type)
                elif output_mime_type == 'application/xml' or output_mime_type == mimetype['xml_results']:
                    return Response(query_results.serialize(format = 'xml'), media_type=output_mime_type)
                elif output_mime_type == mimetype['turtle']:
                    # .serialize(format='turtle').decode("utf-8")
                    return Response(query_results.serialize(format = 'turtle'), media_type=output_mime_type)
                else:
                    ## XML by default for federated queries
                    return Response(query_results.serialize(format = 'xml'), media_type=mimetype['xml_results'])
            except Exception as e:
                logging.error("Error serializing the SPARQL query results with RDFLib: " + str(e))
                return JSONResponse(status_code=422, content={"message": "Error serializing the SPARQL query results: " + str(e)})

        @self.post(
            "/sparql",
            name="SPARQL endpoint",
            description=self.example_query,
            responses=api_responses
        )
        async def post_sparql_endpoint(
            request: Request,
            query: Optional[str] = Query(None)):
            """Send a SPARQL query to be executed through HTTP POST operation.
            \f
            :param request: The HTTP POST request with a .body()
            :param query: SPARQL query input.
            """
            if not query:
                # Handle federated query services which provide the query in the body
                query_body = await request.body()
                body = parse.unquote(query_body.decode('utf-8'))
                parsed_query = parse.parse_qsl(body)
                for params in parsed_query:
                    if params[0] == 'query':
                        query = params[1]
            return await sparql_endpoint(request, query)


        @self.get("/", include_in_schema=False)
        async def serve_yasgui():
            """Serve YASGUI interface"""
            return FileResponse(pkg_resources.resource_filename('rdflib_endpoint', 'yasgui.html'))


        # Service description returned when no query provided
        service_description_ttl = """@prefix sd: <http://www.w3.org/ns/sparql-service-description#> .
        @prefix ent: <http://www.w3.org/ns/entailment/> .
        @prefix prof: <http://www.w3.org/ns/owl-profile/> .
        @prefix void: <http://rdfs.org/ns/void#> .
        @prefix dc: <http://purl.org/dc/elements/1.1/> .
        @prefix rdfs: <http://www.w3.org/2000/01/rdf-schema#> .

        <{public_url}> a sd:Service ;
            rdfs:label "{title}" ;
            dc:description "{description}" ;
            sd:endpoint <{public_url}> ;
            sd:supportedLanguage sd:SPARQL11Query ;
            sd:resultFormat <http://www.w3.org/ns/formats/SPARQL_Results_JSON>, <http://www.w3.org/ns/formats/SPARQL_Results_CSV> ;
            sd:feature sd:DereferencesURIs ;
            sd:defaultEntailmentRegime ent:RDFS ;
            sd:defaultDataset [
                a sd:Dataset ;
                sd:defaultGraph [
                    a sd:Graph ;
                ] 
            ] .""".format(public_url=self.public_url, title=self.title, description=self.description.replace("\n", ""))

    def evalCustomFunctions(self, ctx:object, part:object) -> object:
        """Retrieve variables from a SPARQL-query, then execute registered SPARQL functions
        The results are then stored in Literal objects and added to the query results.

        :param ctx:     <class 'rdflib.plugins.sparql.sparql.QueryContext'>
        :param part:    <class 'rdflib.plugins.sparql.parserutils.CompValue'>
        :return:        <class 'rdflib.plugins.sparql.processor.SPARQLResult'>
        """

        # This part holds basic implementation for adding new functions
        if part.name == 'Extend':
            query_results = []

            # Information is retrieved and stored and passed through a generator
            for eval_part in evalPart(ctx, part.p):
                # Checks if the function is a URI (custom function)
                if hasattr(part.expr, 'iri'):

                    # Iterate through the custom functions passed in the constructor
                    for function_uri, custom_function in self.functions.items():
                        # Check if URI correspond to a registered custom function
                        if part.expr.iri == URIRef(function_uri):
                            # Execute each function
                            query_results, ctx, part, eval_part = custom_function(query_results, ctx, part, eval_part)

                else:
                    # For built-in SPARQL functions (that are not URIs)
                    evaluation = [_eval(part.expr, eval_part.forget(ctx, _except=part._vars))]
                    if isinstance(evaluation[0], SPARQLError):
                        raise evaluation[0]
                    # Append results for built-in SPARQL functions
                    for result in evaluation:
                        query_results.append(eval_part.merge({part.var: Literal(result)}))
                        
            return query_results
        raise NotImplementedError()
