#  coding=utf-8
#  Copyright 2021-present, the Recognai S.L. team.
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.
import datetime
import json
import logging
import re
from typing import Any, Callable, Dict, List, Optional, Tuple

from argilla.client.models import Record, TextClassificationRecord, TokenClassificationRecord
from argilla.monitoring.base import BaseMonitor
from argilla.utils.dependency import require_version

require_version("starlette>=0.13.0")
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import JSONResponse, Response, StreamingResponse
from starlette.types import Message, Receive

_logger = logging.getLogger(__name__)
_default_tokenization_pattern = re.compile(r"\W+")


def token_classification_mapper(inputs, outputs):
    if isinstance(inputs, str):
        text = inputs
    elif isinstance(inputs, dict):
        text = inputs.get("text", "")
    else:
        text = ""
    tokens = outputs.get("tokens") if isinstance(outputs, dict) else None
    return TokenClassificationRecord(
        text=text,
        tokens=tokens or _default_tokenization_pattern.split(text),
        prediction=[
            (entity["label"], entity["start"], entity["end"])
            for entity in (outputs.get("entities") if isinstance(outputs, dict) else outputs)
        ],
        event_timestamp=datetime.datetime.now(),
    )


def text_classification_mapper(inputs, outputs):
    return TextClassificationRecord(
        inputs=inputs,
        prediction=[(label, score) for label, score in zip(outputs.get("labels", []), outputs.get("scores", []))],
        event_timestamp=datetime.datetime.now(),
    )


class CachedJsonRequest(Request):
    """
    We must a cached version of incoming requests since request body cannot be read from middleware directly.
    See <https://github.com/encode/starlette/issues/847> for more information

    TODO Remove usage of CachedRequest when https://github.com/encode/starlette/pull/848 is released
    """

    @property
    def receive(self) -> Receive:
        body = None
        if hasattr(self, "_body"):
            body = self._body
        if body is not None:

            async def cached_receive() -> Message:
                return dict(type="http.request", body=body)

            return cached_receive
        return self._receive


class ArgillaLogHTTPMiddleware(BaseHTTPMiddleware):
    """An standard starlette middleware that enables argilla logs for http prediction requests"""

    def __init__(
        self,
        api_endpoint: str,
        dataset: str,
        records_mapper: Optional[Callable[[dict, dict], Record]],
        sample_rate: float = 1.0,
        log_interval: float = 1.0,
        agent: Optional[str] = None,
        tags: Dict[str, str] = None,
        *args,
        **kwargs,
    ):
        BaseHTTPMiddleware.__init__(self, *args, **kwargs)

        self._endpoint = api_endpoint
        self._dataset = dataset
        self._records_mapper = records_mapper
        self._monitor_cfg = dict(
            dataset=dataset,
            sample_rate=sample_rate,
            log_interval=log_interval,
            agent=agent,
            tags=tags,
        )
        self._monitor: Optional[BaseMonitor] = None

    def init(self):
        if self._monitor:
            return
        from argilla.client.api import active_api

        self._monitor = BaseMonitor(
            self,
            api=active_api(),
            **self._monitor_cfg,
        )
        self._monitor._prepare_log_data = self._prepare_argilla_data

    async def dispatch(
        self,
        request: Request,
        call_next: RequestResponseEndpoint,
    ) -> Response:
        self.init()

        if self._endpoint != request.url.path:  # Filtering endpoint path
            return await call_next(request)

        cached_request = CachedJsonRequest(
            scope=request.scope,
            receive=request.receive,
            send=request._send,
        )

        try:
            # Must obtain input parameters from request
            if cached_request.method in ["POST", "PUT"]:
                content_type = request.headers.get("Content-type", None)
                if content_type is None:
                    if "application/json" not in content_type:
                        return await call_next(request)
                inputs = await cached_request.json()
            elif cached_request.method == "GET":
                inputs = cached_request.query_params._dict
            else:
                raise NotImplementedError("Only request methods POST, PUT and GET are implemented.")

            # Must obtain response from request
            response: Response = await call_next(cached_request)
            if not isinstance(response, (JSONResponse, StreamingResponse)) or response.status_code >= 400:
                return response

            new_response, outputs = await self._extract_response_content(response)
            self._monitor.send_records(inputs=inputs, outputs=outputs)
            return new_response
        except Exception as ex:
            _logger.error("Cannot log to argilla", exc_info=ex)
            return await call_next(request)

    async def _extract_response_content(self, response: Response) -> Tuple[Response, List[Dict[str, Any]]]:
        """Extracts response body content from response and returns a new processable response"""
        body = b""
        new_response = response
        if isinstance(response, StreamingResponse):
            async for chunk in response.body_iterator:
                body += chunk
            new_response = StreamingResponse(
                content=(chunk for chunk in [body]),
                status_code=response.status_code,
                headers={k: v for k, v in response.headers.items()},
                media_type=response.media_type,
                background=response.background,
            )
        else:
            body = response.body
        return new_response, json.loads(body)

    def _prepare_argilla_data(self, inputs: List[Dict[str, Any]], outputs: List[Dict[str, Any]], **tags):
        # using the base monitor, we only need to provide the input data to the rg.log function
        # and the monitor will handle the sample rate, queue and argilla interaction
        try:
            records = self._records_mapper(inputs, outputs)
            assert records, ValueError("The records_mapper returns and empty record list.")
            if not isinstance(records, list):
                records = [records]
        except Exception as ex:
            records = []
            _logger.error("Cannot log to argilla. Error in records mapper.", exc_info=ex)

        for record in records:
            if self._monitor.agent is not None and not record.prediction_agent:
                record.prediction_agent = self._monitor.agent

        return dict(
            records=records or [],
            name=self._dataset,
            tags=tags,
        )
