import asyncio
from asyncio.base_events import BaseEventLoop
import threading
from typing import List
from google.protobuf.descriptor import Descriptor, FieldDescriptor, FileDescriptor, MethodDescriptor, ServiceDescriptor
from google.protobuf.json_format import MessageToDict
from google.protobuf.message import Message
from grpc import Channel, Server, StatusCode
from delphai_utils.logging import logging
from google.protobuf.descriptor_pb2 import MethodOptions
from google.api.http_pb2 import HttpRule
from starlette.applications import Starlette
from starlette.responses import JSONResponse
from starlette.routing import Route
from starlette.exceptions import HTTPException
from starlette.middleware.base import BaseHTTPMiddleware
from grpc.experimental.aio import insecure_channel, AioRpcError
from starlette.middleware.cors import CORSMiddleware
import os
from google.protobuf import symbol_database
from hypercorn.config import Config
from hypercorn.asyncio import serve
from grpc import StatusCode
from urllib.parse import urlparse
from time import perf_counter
from starlette_prometheus import metrics, PrometheusMiddleware
from delphai_utils.utils import generate_default_message
supported_methods = ['get', 'put', 'post', 'delete', 'patch']

status_map = {
  StatusCode.OK: 200,
  StatusCode.CANCELLED: 499,
  StatusCode.UNKNOWN: 500,
  StatusCode.INVALID_ARGUMENT: 400,
  StatusCode.DEADLINE_EXCEEDED: 504,
  StatusCode.NOT_FOUND: 404,
  StatusCode.ALREADY_EXISTS: 409,
  StatusCode.PERMISSION_DENIED: 403,
  StatusCode.UNAUTHENTICATED: 401,
  StatusCode.RESOURCE_EXHAUSTED: 429,
  StatusCode.FAILED_PRECONDITION: 412,
  StatusCode.ABORTED: 499,
  StatusCode.OUT_OF_RANGE: 416,
  StatusCode.UNIMPLEMENTED: 404,
  StatusCode.INTERNAL: 418,
  StatusCode.UNAVAILABLE: 503,
  StatusCode.DATA_LOSS: 420
}


class AccessLogMiddleware(BaseHTTPMiddleware):
  async def dispatch(self, request, call_next):
    start = perf_counter()
    response = await call_next(request)
    if response.status_code < 400:
      path = urlparse(str(request.url)).path
      end = perf_counter()
      elapsed = round((end - start) * 1000, 2)
      logging.info(f'[{response.status_code}] {request.method} {path} [{elapsed}ms]')
    return response


async def process_grpc_request(descriptor: FileDescriptor, service_name: str, method: str, input: Message,
                               channel: Channel):
  service = descriptor.services_by_name[service_name]
  method_descriptor = service.methods_by_name[method]
  input_prototype = symbol_database.Default().GetPrototype(method_descriptor.input_type)
  output_prototype = symbol_database.Default().GetPrototype(method_descriptor.output_type)
  method_callable = channel.unary_unary(f'/{service.full_name}/{method}',
                                        request_serializer=input_prototype.SerializeToString,
                                        response_deserializer=output_prototype.FromString)
  response = await method_callable(input_prototype(**input))
  return response


async def http_exception(request, exc):
  if 'favicon.ico' in str(request.url):
    detail = 'not found'
    status_code = 404
  else:
    path = urlparse(str(request.url)).path
    logging.error(f'[{exc.status_code}] {request.method} {path} - {exc.detail}')
    detail = exc.detail
    status_code = exc.status_code
  return JSONResponse({'detail': detail, 'status': exc.status_code}, status_code=status_code)


def get_http_handlers(descriptor: FileDescriptor, service_name: str, method: str, channel: Channel):
  async def method_get_handler(request):
    service = descriptor.services_by_name[service_name]
    method_descriptor = service.methods_by_name[method]
    input = generate_default_message(method_descriptor.input_type)
    output = generate_default_message(method_descriptor.output_type)

    function_name = f'{service.full_name}.{method}'
    return JSONResponse({'function_name': function_name, 'input': input, 'output': output})

  async def request_handler(request):
    body = {}
    if len(await request.body()) > 0:
      body = await request.json()
    input = {**request.path_params, **request.query_params, **body}
    try:
      raw_output = await process_grpc_request(descriptor, service_name, method, input, channel)
      output = MessageToDict(raw_output, preserving_proto_field_name=True)
      return JSONResponse(output)
    except AioRpcError as ex:
      detail = ex.details()
      grpc_status = ex.code()
      http_status_code = status_map[grpc_status]
      raise HTTPException(http_status_code, detail=detail)
    except Exception as ex:
      detail = str(ex).replace('\n', ' ')
      http_status_code = 500
      raise HTTPException(http_status_code, detail=detail)

  return method_get_handler, request_handler


async def start_gateway(descriptor: FileDescriptor, port: int = 7070):
  try:
    logging.info('starting gateway...')
    debug = 'DELPHAI_ENVIRONMENT' in os.environ and os.environ['DELPHAI_ENVIRONMENT'] == 'development'
    app = Starlette(debug=debug)
    app.add_exception_handler(HTTPException, http_exception)
    app.add_middleware(AccessLogMiddleware)
    app.add_middleware(PrometheusMiddleware)
    app.add_middleware(CORSMiddleware,
                       allow_origins=['*'],
                       allow_methods=['*'],
                       allow_headers=['*'],
                       allow_credentials=True)
    app.add_exception_handler(HTTPException, http_exception)
    app.add_route("/metrics/", metrics)
    channel = insecure_channel('localhost:8080')
    config = Config()
    config.bind = [f'0.0.0.0:{port}']
    for service_name in descriptor.services_by_name.keys():
      service_handler: ServiceDescriptor = descriptor.services_by_name[service_name]
      if service_handler.full_name.startswith('grpc.'):
        logging.info(f'skipping service {service_handler.name}')
      else:
        logging.info(f'processing service {service_handler.name}')
        for key in service_handler.methods_by_name.keys():
          method_get_handler, request_handler = get_http_handlers(descriptor, service_handler.name, key, channel)
          service_name = key[1:].split('/')[0].split('.')[-1]
          app.add_route(f'/{service_handler.full_name}.{key}', route=method_get_handler, methods=['get'])
          app.add_route(f'/{service_handler.full_name}.{key}', route=request_handler, methods=['post'])
          logging.info(f'  processing {key}')
          method_descriptor: MethodDescriptor = service_handler.methods_by_name[key]
          method_options: MethodOptions = method_descriptor.GetOptions()
          fields: List(FieldDescriptor) = method_options.ListFields()
          for field in fields:
            if field[0].full_name == 'google.api.http':
              http_rule: HttpRule = field[1]
              for supported_method in supported_methods:
                http_path = getattr(http_rule, supported_method)
                if http_path != '':
                  app.add_route(http_path, route=request_handler, methods=[supported_method])
    logging.info(f'started gateway on port {port}')
    return await serve(app, config)
  except Exception as ex:
    logging.error(str(ex))
