import json
from urllib.parse import urlparse
import logging

from fastapi.responses import PlainTextResponse
from starlette.types import Message, Scope, Receive, Send

from metlo.utils import map_log_level_to_string, DEFAULT_BODY_SIZE

import metlo_python_agent_bindings_common


def get_src_ip(scope: Scope):
    src_ip = ""
    try:
        src_ip = (
            scope["client"][0]
            if "1.0.0.127.in-addr.arpa" not in scope["client"][0]
            else "localhost"
        )
    except:
        pass
    return src_ip


def get_dest_ip(scope: Scope):
    dest_ip = ""
    try:
        dest_ip = scope["server"][0]
    except:
        pass
    return dest_ip


def get_src_port(scope: Scope):
    source_port = "0"
    try:
        source_port = scope["client"][1] if scope["client"][1] is not None else "0"
    except:
        source_port = "0"
    return int(source_port)


def get_dest_port(scope: Scope):
    dest_port = "0"
    try:
        dest_port = scope["server"][1] if scope["server"][1] is not None else "0"
    except:
        source_port = "0"
    return int(dest_port)


class ASGIMiddleware:
    def compile_request(self, scope: Scope):
        user = None
        try:
            user = self.get_user(scope) if self.get_user else None
        except Exception as e:
            self.logger.error("Error obtaining user info from function")
            self.logger.error(e)
            user = None
        return {
            "url": {
                "host": self.api_host
                if self.api_host is not None
                else f"{get_dest_ip(scope)}:{get_dest_port(scope)}",
                "path": scope["route"].path,
                "parameters": list(
                    map(
                        lambda x: {"name": x[0], "value": x[1]},
                        [
                            x.split("=")
                            for x in scope["query_string"].decode().split("&")
                            if x
                        ],
                    )
                ),
            }
            if ("query_string" in scope and isinstance(scope["query_string"], bytes))
            else [],
            "headers": list(
                map(
                    lambda x: {"name": x[0].decode(), "value": x[1].decode()},
                    [
                        x
                        for x in scope["headers"]
                        if isinstance(x[0], bytes) and isinstance(x[1], bytes)
                    ],
                )
            )
            if ("headers" in scope and isinstance(scope["headers"], list))
            else [],
            "body": (
                scope["state"]["req_body"]
                if (
                    "req_body" in scope["state"]
                    and isinstance(scope["state"]["req_body"], bytes)
                )
                else b""
            ).decode(),
            "method": scope["method"],
            "user": user,
        }

    def compile_meta(self, scope: Scope):
        return {
            "environment": "production",
            "incoming": True,
            "source": get_src_ip(scope),
            "sourcePort": get_src_port(scope),
            "destination": get_dest_ip(scope),
            "destinationPort": get_dest_port(scope),
            "metloSource": "python/fastapi",
        }

    def compile_response(self, scope: Scope):
        res_body: str = (
            scope["state"]["res_body"] if "res_body" in scope["state"] else ""
        )
        return {
            "status": scope["state"]["status"] if "status" in scope["state"] else [],
            "headers": list(
                map(
                    lambda x: {"name": x[0].decode(), "value": x[1].decode()},
                    [
                        x
                        for x in scope["state"]["res_headers"]
                        if isinstance(x[0], bytes) and isinstance(x[1], bytes)
                    ],
                )
            )
            if (
                "res_headers" in scope["state"]
                and isinstance(scope["state"]["res_headers"], list)
            )
            else [],
            "body": scope["state"]["res_body"]
            if "res_body" in scope["state"]
            and isinstance(scope["state"]["res_body"], str)
            else "",
        }

    def check_blocking(self, scope: Scope):
        try:
            compiledRequest = self.compile_request(scope)
            compiledMeta = self.compile_meta(scope)
            will_reject = metlo_python_agent_bindings_common.block_trace(
                json.dumps(compiledRequest), json.dumps(compiledMeta)
            )
            self.logger.debug("Will be rejecting response from Metlo")
            return will_reject
        except Exception as e:
            self.logger.error(e)
            return False

    def process_trace(self, scope):
        try:
            compiledRequest = self.compile_request(scope)
            compiledMeta = self.compile_meta(scope)
            compiledResponse = self.compile_response(scope)
            will_reject = metlo_python_agent_bindings_common.process_trace(
                json.dumps(
                    {
                        "request": compiledRequest,
                        "response": compiledResponse,
                        "meta": compiledMeta,
                    }
                )
            )
        except Exception as e:
            self.logger.error(e)

    async def process_rejection(self, scope, receive, send):
        try:
            if self.reject_response:
                return await self.reject_response(scope, receive, send)
            else:
                return await PlainTextResponse("Forbidden", 403)(scope, receive, send)
        except Exception as e:
            self.logger.error(
                "Encountered error while sending rejection response. Attempting to send default rejection response 403"
            )
            self.logger.error(e)
            return await PlainTextResponse("Forbidden", 403)(scope, receive, send)

    def __init__(self, app, host, api_key, **kwargs):
        self.app = app
        try:
            self.disabled = kwargs.get("disabled", False)
            if not (self.disabled):
                self.logger = logging.getLogger("metlo")
                ch = logging.StreamHandler()

                level = kwargs.get("log_level", logging.INFO)

                ch.setLevel(level)
                formatter = logging.Formatter(
                    "[%(asctime)s] [%(process)d] [%(levelname)s] [%(name)s]  %(message)s",
                    datefmt="%Y-%m-%d %H:%M:%S %z",
                )
                ch.setFormatter(formatter)
                self.logger.addHandler(ch)

                assert (
                    host is not None
                ), "Metlo for FastAPI __init__ is missing host parameter"

                assert (
                    api_key is not None
                ), "Metlo for FASTAPI __init__ is missing api_key parameter"

                assert urlparse(host).scheme in [
                    "http",
                    "https",
                ], f"Metlo for FLASK has invalid host scheme. Host must be in format http[s]://example.com"

                self.host = host
                self.key = api_key
                params = {}

                settings = kwargs
                shouldStart = True

                params["log_level"] = map_log_level_to_string(level)

                # Collector Port
                if "collector_port" in settings:
                    if (settings["collector_port"] is not None) and isinstance(
                        settings["collector_port"], int
                    ):
                        params["collector_port"] = settings["collector_port"]
                    else:
                        self.logger.error(
                            "Metlo collector_port arg is incorrectly specified. Metlo will not start"
                        )
                        shouldStart = False

                # Backend Port
                if "backend_port" in settings:
                    if (settings["backend_port"] is not None) and isinstance(
                        settings["backend_port"], int
                    ):
                        params["backend_port"] = settings["backend_port"]
                    else:
                        self.logger.error(
                            "Metlo backend_port arg is incorrectly specified. Metlo will not start"
                        )
                        shouldStart = False

                # Encryption Key
                if "encryption_key" in settings:
                    if (settings["encryption_key"] is not None) and isinstance(
                        settings["encryption_key"], str
                    ):
                        params["encryption_key"] = settings["encryption_key"]
                    else:
                        self.logger.error(
                            "Metlo encryption_key arg is incorrectly specified. Metlo will not start"
                        )
                        shouldStart = False

                # Encryption Key
                if "api_host" in settings:
                    if (settings["api_host"] is not None) and isinstance(
                        settings["api_host"], str
                    ):
                        self.api_host = settings["api_host"]
                    else:
                        self.api_host = None
                        self.logger.error(
                            "Metlo api_host arg is incorrectly specified. Defaulting to value from connection"
                        )
                else:
                    self.api_host = None

                # Get User
                if "get_user" in settings and callable(settings["get_user"]):
                    self.get_user = settings["get_user"]
                else:
                    self.logger.debug(
                        "Metlo get_user arg is incorrectly specified or unspecified"
                    )
                    self.get_user = None

                # Reject Response
                if "reject_response" in settings and callable(
                    settings["reject_response"]
                ):
                    self.reject_response = settings["reject_response"]
                else:
                    self.logger.debug(
                        "Metlo reject_response arg is incorrectly specified or unspecified"
                    )
                    self.reject_response = None

                if "max_body" in settings and isinstance(settings["max_body"], int):
                    self.max_body = settings["max_body"]
                else:
                    self.max_body = DEFAULT_BODY_SIZE

                can_start = False
                try:
                    if shouldStart:
                        metlo_python_agent_bindings_common.setup(
                            self.host, self.key, **params
                        )
                        can_start = True
                except Exception as e:
                    self.logger.error(e)

                self.disabled = not (can_start)
        except Exception as e:
            print("Metlo encountered an error during initialization")
            print(e)
            self.disabled = True

    async def __call__(self, scope, receive, send):
        if not (self.disabled):

            async def rec_wrapper():
                return await self.receivefn(scope, receive, send)

            async def send_wrapper(message):
                await self.sendfn(scope, send, message)

            await self.app(scope, rec_wrapper, send_wrapper)
        else:
            await self.app(scope, receive, send)

    async def receivefn(self, scope: Scope, receive: Receive, send: Send) -> None:
        # Receive this part of the request body
        req = await receive()
        if "http" in req["type"]:
            try:
                if "body" in req:
                    if "req_body" in scope["state"]:
                        curr_length = len(req["body"])
                        remaining_buffer = max(self.max_body - curr_length, 0)
                        if remaining_length > 0:
                            copy_length = min(remaining_buffer, len(req["body"]))
                            scope["state"]["req_len_stored"] += copy_length
                            scope["state"]["req_body"] = (
                                scope["state"]["req_body"] + req["body"][:copy_length]
                            )
                    else:
                        scope["state"]["req_len_stored"] = len(req["body"])
                        scope["state"]["req_body"] = req["body"]
                else:
                    scope["state"]["req_len_stored"] = 0
                    scope["state"]["req_body"] = ""
                if not ("blocking_checked" in scope["state"]):
                    # blocking_checked has not been set, so we haven't checked for blocking yet.
                    if "more_body" in req and req["more_body"]:
                        # There is still more body to read
                        if scope["state"]["req_len_stored"] == self.max_body:
                            scope["state"]["blocking_checked"] = True
                            # Do blocking
                            if self.check_blocking(scope):
                                scope["state"]["rejected"] = True
                                return await self.process_rejection(
                                    scope, receive, send
                                )
                    else:
                        # No more body to read
                        scope["state"]["blocking_checked"] = True
                        # Do blocking
                        if self.check_blocking(scope):
                            scope["state"]["rejected"] = True
                            return await self.process_rejection(scope, receive, send)
            except Exception as e:
                self.logger.error(
                    "Encountered an error while processing http request. Processing as normal"
                )
                self.logger.error(e)
            return req
        elif "websocket" in req["type"]:
            return req
        else:
            return req

    async def sendfn(self, scope: Scope, send: Send, message: Message) -> None:
        if "rejected" in scope["state"]:
            return
        else:
            if "http" in message["type"]:
                try:
                    # We received http response start. Can record status and headers from this
                    if message["type"] == "http.response.start":
                        scope["state"]["status"] = message["status"]
                        scope["state"]["res_headers"] = message["headers"]
                    elif message["type"] == "http.response.body":
                        # If message has a body, then see if we can add body
                        if "body" in message:
                            # Have we already recorded some part of the response body ?
                            if "res_body" in scope["state"]:
                                # Get length of body currently stored
                                curr_length = scope["state"]["res_len_stored"]
                                # Remaining space in the buffer
                                remaining_buffer = max(self.max_body - curr_length, 0)
                                # If remaining space in the buffer is more than zero, then try to add part of body to buffer
                                if remaining_buffer > 0:
                                    # Figure out how much we can add to the buffer
                                    copy_length = min(
                                        remaining_buffer, len(message["body"])
                                    )
                                    scope["state"]["res_len_stored"] += copy_length
                                    # Add decoded body to buffer
                                    scope["state"]["res_body"] = (
                                        scope["state"]["res_body"]
                                        + message["body"].decode()[:copy_length]
                                    )
                                else:
                                    # Buffer is full, do nothing
                                    pass
                            # We haven't recorded any part of the body yet.
                            else:
                                # We can either fill the entire body or upto self.max_body
                                scope["state"]["res_len_stored"] = min(
                                    len(message["body"]), self.max_body
                                )
                                # We want decoded bodies only
                                scope["state"]["res_body"] = message["body"].decode()[
                                    : self.max_body
                                ]
                        else:
                            # Somehow body is missing from message of type 'http.response.body'
                            scope["state"]["res_len_stored"] = 0
                            scope["state"]["res_body"] = ""
                    # We haven't capture response yet
                    if not ("response_captured" in scope["state"]):
                        # Check if there's still more body
                        if "more_body" in message and message["more_body"]:
                            # Check if we have actually filled the body
                            if scope["state"]["res_len_stored"] == self.max_body:
                                # Filled the body completely, Set the flag and process trace
                                scope["state"]["response_capture"] = True
                                self.process_trace(scope)
                            else:
                                # We haven't filled the body yet and there's still more data to consume
                                pass
                        # There's no more data in the buffer and we haven't already processed the data.
                        else:
                            # Set the flag and process the trace
                            scope["state"]["response_capture"] = True
                            self.process_trace(scope)
                except Exception as e:
                    self.logger.error(
                        "Encountered an error while sending response. Processing as normal"
                    )
                    self.logger.error(e)
            await send(message)
