# SPDX-License-Identifier: MIT
# Copyright (c) 2023 Anggit Arfanto

import asyncio

from tremolo.utils import parse_int
from .contexts import ConnectionContext
from .http_exceptions import (
    HTTPException,
    BadRequest,
    ExpectationFailed,
    RequestTimeout
)
from .http_header import HTTPHeader
from .http_request import HTTPRequest
from .queue import Queue


# pylint: disable=E0237,E1101
class HTTPProtocol(asyncio.Protocol):
    __slots__ = ('server', 'queue', 'events', 'handlers', '_recv_buf')

    def __init__(self, app, **kwargs):
        self.app = app
        self.extras = kwargs
        self.loop = app.loop
        self.logger = app.logger
        self.globals = app.context  # a worker-level context
        self.context = ConnectionContext()
        self.lock = kwargs['lock']
        self.fileno = -1
        self.request = None  # current request object

        self.server = self.__dict__.copy()  # all the unslotted objects above
        self.queue = [Queue(), Queue()]  # IN, OUT
        self.events = {}
        self.handlers = set()
        self._recv_buf = bytearray()

    @property
    def options(self):
        return self.extras['options']

    @property
    def transport(self):
        return self.context.transport

    def add_close_callback(self, callback):
        self.context.tasks.add(callback)

    def create_task(self, coro):
        task = self.loop.create_task(coro)

        self.context.tasks.add(task)
        task.add_done_callback(self.task_done)

        return task

    def task_done(self, task):
        self.context.tasks.discard(task)

        if not task.cancelled():
            exc = task.exception()

            if exc:
                self.close(exc)

    def connection_made(self, transport):
        self.context.update(transport=transport)
        self.fileno = transport.get_extra_info('socket').fileno()

        self.events['request'] = self.loop.create_future()

        self.app.add_task(self.create_task(self._send_data()))
        self.app.add_task(self.create_task(
            self.set_timeout(self.events['request'],
                             timeout=self.options['request_timeout'],
                             timeout_cb=self.request_timeout)
        ))

    def is_closing(self):
        return self.transport is None or self.transport.is_closing()

    def close(self, exc=None):
        if self.is_closing():
            return

        if exc:
            if isinstance(exc, HTTPException) and exc.code < 600:
                self.transport.write(
                    b'HTTP/1.0 %d %s\r\nContent-Type: %s\r\n\r\n' %
                    (exc.code,
                     exc.message.encode(exc.encoding),
                     exc.content_type.encode(exc.encoding))
                )
                self.transport.write(str(exc).encode(exc.encoding))

            self.print_exception(exc, 'close')

        try:
            self.transport.write_eof()
        except (OSError, RuntimeError) as exc:
            self.logger.info(exc)
        finally:
            self.add_close_callback(
                self.loop.call_at(self.loop.time() +
                                  self.options['app_close_timeout'],
                                  self.transport.abort).cancel
            )
            self.transport.close()

    def request_timeout(self, timeout):
        raise RequestTimeout('request timeout after %gs' % timeout)

    def keepalive_timeout(self, timeout):
        self.logger.debug('keepalive timeout after %gs', timeout)
        self.close()

    def send_timeout(self, timeout):
        self.logger.debug('send timeout after %gs', timeout)

    async def set_timeout(self, waiter, timeout=30, timeout_cb=None):
        timer = self.loop.call_at(self.loop.time() + timeout, waiter.cancel)

        try:
            return await waiter
        except asyncio.CancelledError:
            if not self.is_closing():
                if callable(timeout_cb):
                    timeout_cb(timeout)

                raise
        finally:
            timer.cancel()

    async def put_to_queue(self, data, name=0, rate=-1):
        if self.queue:
            self.queue[name].put_nowait(data)
            queue_size = self.queue[name].qsize()

            if queue_size <= self.options['max_queue_size']:
                if data and rate > 0 and queue_size > rate / len(data):
                    await asyncio.sleep(queue_size * len(data) / rate)

                return True

            self.logger.error('%d exceeds the value of max_queue_size',
                              queue_size)

        self.close()

    async def request_received(self, request, response):
        raise NotImplementedError

    async def error_received(self, exc, response):
        exc = HTTPException(cause=exc)
        i = exc.code - 400

        if 0 <= i < len(self.app.routes[0]) and not response.headers_sent():
            if self.app.routes[0][i][1] is None:
                i = 100

            _, func, kwargs, _ = self.app.routes[0][i]

            response.headers.clear()
            response.set_status(exc.code, exc.message)
            response.set_content_type(exc.content_type)

            return await func(func=func,
                              kwargs=kwargs,
                              request=response.request,
                              response=response,
                              loop=self.loop,
                              exc=exc)

    def handlers_timeout(self):
        if self.request is None or not self.request.upgraded:
            while self.handlers:
                self.handlers.pop().cancel()
                self.logger.error(
                    'handler timeout '
                    '(app_handler_timeout=%g, app_close_timeout=%g)',
                    self.options['app_handler_timeout'],
                    self.options['app_close_timeout']
                )

    def set_handler_timeout(self, timeout):
        if self.handlers:
            return self.loop.call_at(
                self.loop.time() + timeout, self.handlers_timeout
            )

    def print_exception(self, exc, *args):
        cause = exc
        args += (exc.__class__.__name__,)
        errno = -1

        while isinstance(cause, HTTPException):
            if cause.code < 500:
                errno = 0

            if cause.__cause__ is None:
                break

            cause = cause.__cause__
            args += (cause.__class__.__name__,)

        if cause.args and isinstance(cause.args[0], int):
            errno = cause.args[0]
        elif errno == -1:
            errno = getattr(cause, 'errno', -1)

        msg = ': '.join((*args, str(exc)))

        if errno == 0:
            self.logger.info(msg, exc_info=self.options['debug'] and exc)
        elif not isinstance(exc, asyncio.CancelledError):
            self.logger.error(msg, exc_info=self.options['debug'] and exc)

    def send_continue(self):
        if self.request is None or not self.request.http_continue:
            return

        if self.request.content_length > self.options['client_max_body_size']:
            raise ExpectationFailed

        self.transport.write(
            b'HTTP/%s 100 Continue\r\n\r\n' % self.request.version
        )
        self.queue[1].put_nowait(None)

    async def _handle_request(self, request):
        response = request.create_response()
        timer = self.set_handler_timeout(self.options['app_handler_timeout'])

        try:
            if self.options['keepalive_timeout']:  # can be 0 on shutdown
                if b'connection' in request.headers:
                    if b'close' not in request.headers.getlist(b'connection'):
                        self.globals.connections.add(self)
                        request.http_keepalive = True
                elif request.version == b'1.1':
                    self.globals.connections.add(self)
                    request.http_keepalive = True

            if b'transfer-encoding' in request.headers:
                if (request.version != b'1.1' or
                        b'chunked' not in request.transfer_encoding):
                    raise BadRequest('unexpected Transfer-Encoding')

                request.has_body = True

            if b'content-length' in request.headers:
                if (request.has_body or
                        len(request.headers[b'content-length']) != 1):
                    raise BadRequest('ambiguous Content-Length')

                try:
                    request.content_length = parse_int(
                        request.headers[b'content-length'][0]
                    )
                except ValueError as exc:
                    raise BadRequest('bad Content-Length') from exc

                request.has_body = request.content_length != 0

            if (request.has_body and b'expect' in request.headers and
                    request.headers[b'expect'][0].lower() == b'100-continue'):
                # we can handle continue later after the route is found
                # by checking this state
                request.http_continue = True

            # successfully got header,
            # clear either the request or keepalive timeout
            if not self.events['request'].done():
                self.events['request'].set_result(request)

            await self.request_received(request, response)
        except (SystemExit, KeyboardInterrupt):
            raise
        except BaseException as exc:
            if request.protocol is not None:
                data = None

                try:
                    data = await self.error_received(exc, response)
                except Exception as e:
                    exc = e
                finally:
                    await response.handle_exception(exc, data=data)
        finally:
            timer.cancel()
            await request.handler_exit()

    async def _receive_data(self):
        if 'request' not in self.events:
            return

        request = await self.events['request']
        excess = 0

        if not request.upgraded:
            request.body_size += len(self._recv_buf)

            if request.body_size > request.content_length > -1:
                excess = request.body_size - request.content_length

        while len(self._recv_buf) > excess:
            data = self._recv_buf[:min(self.options['buffer_size'],
                                       len(self._recv_buf) - excess)]

            if await self.put_to_queue(data, rate=self.options['upload_rate']):
                del self._recv_buf[:len(data)]
            else:
                del self._recv_buf[:]
                return

        # maybe resume reading, or close
        if 'receive' in self.events:
            del self.events['receive']

            if request.upgraded:
                if self in self.globals.connections:
                    self.transport.resume_reading()
            elif request.body_size > self.options['client_max_body_size']:
                self.close(BadRequest('payload too large'))
            elif request.body_size >= request.content_length > -1:
                self.queue[0].put_nowait(None)
            elif request.has_body and not request.eof():
                self.transport.resume_reading()

    def data_received(self, data):
        if not data:
            return

        self._recv_buf.extend(data)

        if self.request is None:
            header_size = self._recv_buf.find(b'\r\n\r\n') + 2

            if header_size > self.options['client_max_header_size']:
                self.close(BadRequest('request header too large'))
            elif header_size != 1:
                header = HTTPHeader().parse(self._recv_buf, header_size)
                del self._recv_buf[:header_size + 2]

                if header.is_request:
                    self.request = HTTPRequest(self, header)
                    task = self.app.create_task(
                        self._handle_request(self.request)
                    )

                    self.handlers.add(task)
                    task.add_done_callback(self.handlers.discard)
                else:
                    self.close(BadRequest('bad request: not a request'))
            elif len(self._recv_buf) > self.options['client_max_header_size']:
                self.close(BadRequest('bad request'))

            if self.request is None or not self._recv_buf:
                return

        # resumed in _receive_data or _send_data
        self.transport.pause_reading()

        if 'receive' not in self.events:
            self.events['receive'] = self.create_task(self._receive_data())

    def resume_writing(self):
        if 'send' in self.events and not self.events['send'].done():
            self.events['send'].set_result(None)

    def set_watermarks(self, high=65536, low=8192):
        if self.transport is not None:
            self.transport.set_write_buffer_limits(high=high, low=low)

    async def _send_data(self):
        try:
            while self.queue:
                data = await self.queue[1].get()

                if data is None:
                    # close the transport, unless keepalive is enabled
                    if self.request is not None:
                        if self.request.http_continue:
                            self.request.http_continue = False
                            continue

                        if self.request.http_keepalive:
                            self.request.http_keepalive = False

                            if self in self.globals.connections:
                                if not self.request.upgraded:
                                    await self._handle_keepalive()

                                self.transport.resume_reading()
                                continue

                            self.logger.info(
                                'keepalive connection dropped: %d', self.fileno
                            )

                        del self._recv_buf[:]
                        self.request.clear()

                    self.close()
                    return

                # send data
                write_buffer_size = self.transport.get_write_buffer_size()
                low, high = self.transport.get_write_buffer_limits()

                if write_buffer_size > high:
                    self.logger.debug(
                        '%d exceeds the current watermark limits '
                        '(high=%d, low=%d)',
                        write_buffer_size, high, low
                    )
                    self.events['send'] = self.loop.create_future()

                    await self.set_timeout(
                        self.events['send'],
                        timeout=self.options['request_timeout'],
                        timeout_cb=self.send_timeout
                    )

                    if self.transport is None:
                        return

                self.transport.write(data)
        except asyncio.CancelledError:
            self.close()
        except Exception as exc:
            self.close(exc)

    async def _handle_keepalive(self):
        if self.request.has_body and not self.request.eof():
            self.logger.info('request body was not fully consumed')
            self.request.clear()
            self.close()
            return

        if 'receive' in self.events:
            # waits for all incoming data to enter the queue
            await self.events['receive']

        self.events['request'] = self.loop.create_future()

        self.app.add_task(self.create_task(
            self.set_timeout(self.events['request'],
                             timeout=self.options['keepalive_timeout'],
                             timeout_cb=self.keepalive_timeout)
        ))

        # reset. so the next data in data_received will be considered as
        # a fresh http request (not a continuation data)
        self.request.clear()
        self.request = None
        self.server.clear()
        self.server.update(self.__dict__)
        self.queue[1].queue.clear()

        if self._recv_buf:
            self.queue[0].put_nowait(self._recv_buf[:])
            del self._recv_buf[:]

        while self.queue[0].qsize():
            # this data is supposed to be the next header
            self.data_received(self.queue[0].get_nowait())

    def connection_lost(self, _):
        while self.context.tasks:
            task = self.context.tasks.pop()

            try:
                getattr(task, 'cancel', task)()
            except Exception as exc:
                self.print_exception(exc, 'connection_lost')

        while self.queue:
            self.queue.pop().queue.clear()

        self.request = None

        self.context.clear()
        self.server.clear()
        self.events.clear()

        self.set_handler_timeout(self.options['app_close_timeout'])
        self.globals.connections.discard(self)
