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 import Middleware
from starlette.middleware.base import BaseHTTPMiddleware
from grpc.experimental.aio import insecure_channel
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 Exception as ex:
      if hasattr(ex, '_state'):
        detail = ex._state.details
        grpc_status = ex._state.code.name
        http_status_code = status_map[StatusCode[grpc_status]]
      else:
        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, shutdown_trigger, port: int = 7070):
  logging.info('starting gateway...')
  debug = 'DELPHAI_ENVIRONMENT' in os.environ and os.environ['DELPHAI_ENVIRONMENT'] == 'development'
  app = Starlette(debug=debug, exception_handlers={HTTPException: http_exception})
  app.add_middleware(AccessLogMiddleware)
  app.add_middleware(PrometheusMiddleware)
  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'/{key}', route=method_get_handler, methods=['get'])
        app.add_route(f'/{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, shutdown_trigger=shutdown_trigger)
