import asyncio
import io
import logging
import os
import re
import struct
import urllib.parse
import traceback
from typing import Any, Dict, Iterator, List, Optional, Tuple

from aioquic.asyncio import QuicConnectionProtocol, serve
from aioquic.quic.configuration import QuicConfiguration
from aioquic.quic.connection import END_STATES
from aioquic.quic.events import StreamDataReceived, QuicEvent
from aioquic.tls import SessionTicket

SERVER_NAME = 'aioquic-transport'
logger = logging.getLogger(__name__)

handlers_path = ''


class EventHandler:
    def __init__(self, connection: QuicConnectionProtocol, global_dict: Dict[str, Any]):
        self.connection = connection
        self.global_dict = global_dict

    def handle_client_indication(
            self,
            origin: str,
            query: Dict[str, str]) -> None:
        name = 'handle_client_indication'
        if name in self.global_dict:
            self.global_dict[name](self.connection, origin, query)

    def handle_event(self, event: QuicEvent) -> None:
        name = 'handle_event'
        if name in self.global_dict:
            self.global_dict[name](self.connection, event)


class QuicTransportProtocol(QuicConnectionProtocol):  # type: ignore
    def __init__(self, *args: Any, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.pending_events: List[QuicEvent] = []
        self.client_indication_finished = False
        self.client_indication_data = b''
        self.handler: Optional[EventHandler] = None

    def quic_event_received(self, event: QuicEvent) -> None:
        prefix = '!!'
        logger.info('QUIC event: %s', type(event))
        try:
            if (not self.client_indication_finished and
                    isinstance(event, StreamDataReceived) and
                    event.stream_id == 2):
                # client indication process
                prefix = 'Client indication error: '
                self.client_indication_data += event.data
                if len(self.client_indication_data) > 65535:
                    raise Exception('too large data for client indication')
                if event.end_stream:
                    self.process_client_indication()
                    assert self.handler is not None
                    if self.is_closing_or_closed():
                        return
                    prefix = 'Event handling Error: '
                    for e in self.pending_events:
                        self.handler.handle_event(e)
                    self.pending_events.clear()
            elif not self.client_indication_finished:
                self.pending_events.append(event)
            elif self.handler is not None:
                prefix = 'Event handling Error: '
                self.handler.handle_event(event)
        except Exception as e:
            self.handler = None
            logger.warn(prefix + str(e))
            traceback.print_exc()
            self.close()

    def parse_client_indication(self, bs: io.BytesIO) -> Iterator[Tuple[int, bytes]]:
        while True:
            key_b = bs.read(2)
            if len(key_b) == 0:
                return
            length_b = bs.read(2)
            if len(key_b) != 2:
                raise Exception('failed to get "Key" field')
            if len(length_b) != 2:
                raise Exception('failed to get "Length" field')
            key = struct.unpack('!H', key_b)[0]
            length = struct.unpack('!H', length_b)[0]
            value = bs.read(length)
            if len(value) != length:
                raise Exception('truncated "Value" field')
            yield (key, value)

    def process_client_indication(self) -> None:
        origin = None
        origin_string = None
        path = None
        path_string = None
        KEY_ORIGIN = 0
        KEY_PATH = 1
        for (key, value) in self.parse_client_indication(
                io.BytesIO(self.client_indication_data)):
            if key == KEY_ORIGIN:
                origin_string = value.decode()
                origin = urllib.parse.urlparse(origin_string)
            elif key == KEY_PATH:
                path_string = value.decode()
                path = urllib.parse.urlparse(path_string)
            else:
                # We must ignore unrecognized fields.
                pass
        logger.info('origin = %s, path = %s', origin_string, path_string)
        if origin is None:
            raise Exception('No origin is given')
        assert origin_string is not None
        if path is None:
            raise Exception('No path is given')
        assert path_string is not None
        if origin.scheme != 'https' and origin.scheme != 'http':
            raise Exception('Invalid origin: %s' % origin_string)
        if origin.netloc == '':
            raise Exception('Invalid origin: %s' % origin_string)

        # To make the situation simple we accept only simple path strings.
        m = re.compile(r'^/([a-zA-Z0-9\._\-]+)$').match(path.path)
        if m is None:
            raise Exception('Invalid path: %s' % path_string)

        handler_name = m.group(1)
        query = dict(urllib.parse.parse_qsl(path.query))
        self.handler = self.create_event_handler(handler_name)
        self.handler.handle_client_indication(origin_string, query)
        if self.is_closing_or_closed():
            return
        self.client_indication_finished = True
        logger.info('Client indication finished')

    def create_event_handler(self, handler_name: str) -> EventHandler:
        global_dict: Dict[str, Any] = {}
        with open(handlers_path + '/' + handler_name) as f:
            exec(f.read(), global_dict)
        return EventHandler(self, global_dict)

    def is_closing_or_closed(self) -> bool:
        if self._quic._close_pending:
            return True
        if self._quic._state in END_STATES:
            return True
        return False


class SessionTicketStore:
    '''
    Simple in-memory store for session tickets.
    '''

    def __init__(self) -> None:
        self.tickets: Dict[bytes, SessionTicket] = {}

    def add(self, ticket: SessionTicket) -> None:
        self.tickets[ticket.ticket] = ticket

    def pop(self, label: bytes) -> Optional[SessionTicket]:
        return self.tickets.pop(label, None)


def start(**kwargs: Any) -> None:
    configuration = QuicConfiguration(
        alpn_protocols=['wq-vvv-01'] + ['siduck'],
        is_client=False,
        max_datagram_frame_size=65536,
    )

    global handlers_path
    handlers_path = os.path.abspath(os.path.expanduser(
        kwargs['handlers_path']))
    logger.info('port = %s', kwargs['port'])
    logger.info('handlers path = %s', handlers_path)

    # load SSL certificate and key
    configuration.load_cert_chain(kwargs['certificate'], kwargs['private_key'])

    ticket_store = SessionTicketStore()

    loop = asyncio.get_event_loop()
    loop.run_until_complete(
        serve(
            kwargs['host'],
            kwargs['port'],
            configuration=configuration,
            create_protocol=QuicTransportProtocol,
            session_ticket_fetcher=ticket_store.pop,
            session_ticket_handler=ticket_store.add,
        )
    )
    try:
        loop.run_forever()
    except KeyboardInterrupt:
        pass