import asyncio
import inspect
import json
import logging
import os
import sys
import traceback

from collections import defaultdict
from contextlib import contextmanager, suppress
from functools import wraps
from op_service_client import OPJSONEncoder
from starlette.applications import Starlette
from starlette.middleware import Middleware
from starlette.middleware.cors import CORSMiddleware
from starlette.responses import Response

from .config import load_config
from .error_monitor import OPErrorMonitor
from .exceptions import AuthenticationException, ClientException, InvalidAuthorizationHeader, InvalidInput, ServiceException, ignore_exceptions
from .instance_id import SERVICE_INSTANCE_ID
from .jwt_manager import JWTManager
from .logging_config import get_relative_sys_filename
from .request_monitor import RequestMonitor
from .op_version_metadata import op_version_metadata
from .throttle_monitor import ThrottleMonitor
from .type_validation import describe, validate
from .type_validation_functions import non_empty_str, regex


class OPService:

  def __init__(self, jwt_config=None, alarm=None, throttle=None):

    # Note: I generally don't like the concept of "middleware" as it's less clear where / how this code runs.
    # Generally, I prefer to explicitly wrap "middleware" around the "run" function inside the "route" decorator,
    # except where Starlette's built-ins are super convenient...
    middleware = [ Middleware(CORSMiddleware, allow_origins=[ '*' ], allow_headers=["*"], allow_methods=["*"]) ]

    self.starlette = Starlette(middleware=middleware)

    self._routes = []

    self._is_test_instance = os.environ.get('OP_DEV_SERVICE_TEST', False)

    self.jwt = JWTManager(jwt_config) if jwt_config else None
    self.alarm = alarm
    self._error_monitor = OPErrorMonitor(alarm) if alarm else None
    self._request_monitor = RequestMonitor()
    self._throttle_monitor = ThrottleMonitor(throttle) if throttle else None

    # Every production service has a '/' health-check endpoint
    @self.starlette.route('/', methods=[ 'GET' ])
    async def health(request):
      return Response('')

    # Returns the commit / config IDs
    @self.route('/__private/metadata', private=True)
    def metadata():
      return op_version_metadata

    # Returns documentation about the routes
    @self.route('/__private/explore', private=True)
    def explore():
      return self.build_explore_data()

    if self._is_test_instance and jwt_config:
      # If the service was brought up in debug mode (using 'op-dev run'), then add a private API that allows us to
      # generate custom JWTs
      @self.route('/__private/generate-jwt', dict, private=True)
      def generate_jwt(data):
        return self.jwt.encode(data)


  def build_explore_data(self):
    grouped_routes = defaultdict(list)
    for route in self._routes:
      grouped_routes[route['grouping']].append(route)
    output = []
    for key in sorted(grouped_routes):
      output.append(
        {
          'group': key,
          'routes': [
            {
              'name': i['name'],
              'profile': describe(i['fn']._profile_format),
              'data': describe(i['fn']._data_format),
              'documentation': i['fn'].__doc__
            } for i in grouped_routes[key]
          ]
        }
      )
    return output


  # This will only ever get called if we have the 'jwt' property initialized (someone used it as a decorator)
  def _get_input_profile(self, request, profile_format, request_state):
    headers = request.headers
    auth_header_value = headers.get('Authorization')
    if auth_header_value:
      if not auth_header_value.startswith('Bearer '):
        raise InvalidAuthorizationHeader("Authorization header is malformatted - does not start with 'Bearer '")
      token = auth_header_value[7:]
      profile = self.jwt.decode(token)
    else:
      profile = None
    request_state['profile'] = profile
    try:
      request_state['profile_validated'] = validated = validate(profile, profile_format)
      return validated
    except Exception as e:
      raise AuthenticationException() from e


  async def _get_input_data(self, request, data_format, request_state):
    input_data = None
    if data_format is not None:
      content_type = request.headers.get('content-type', '')
      if content_type == 'application/json':
        body = await request.body()
        if body:
          try:
            input_data = json.loads(body)
          except Exception as e:
            raise ClientException('Error parsing JSON') from e
      elif content_type.startswith('multipart/form-data'):
        try:
          form = await request.form()
          def inject_files(d):
            file_id = len(d) == 1 and d.get('__file__')
            return form[file_id] if file_id else d
          input_data = json.loads(form['_'], object_hook=inject_files)
        except Exception as e:
          raise ClientException('Error parsing form data') from e
      else:
        raise ClientException('Unrecognized content-type')
    request_state['input_data'] = input_data
    try:
      request_state['input_data_validated'] = validated = validate(input_data, data_format)
      return validated
    except Exception as e:
      raise InvalidInput(e)


  # Functions with this decorator can still be called directly, though they won't get any input / profile validation
  def route(self, route_name, data_format=None, private=False):
    route_name = validate(route_name, regex('(/[a-zA-Z0-9-_]+)+'))
    for route in self._routes:
      if route['name'] == route_name:
        raise Exception('Duplicate routes: ' + route_name)

    def decorator(fn):
      fn._is_route = True
      fn._data_format = data_format
      fn._profile_format = None

      if not private:
        calling_module = inspect.getmodule(inspect.currentframe().f_back)
        calling_module_name = calling_module and calling_module.__name__ or None
        self._routes.append({ 'name': route_name, 'fn': fn, 'grouping': calling_module_name })

      @self.starlette.route(route_name, methods=[ 'POST' ])
      async def run(request):
        ip = request.headers.get('X-Forwarded-For')
        request_state = { 'route': route_name, 'ip': ip, 'ua': request.headers.get('User-Agent') }
        self._request_monitor.on_route_start(request_state)

        try:
          if self._throttle_monitor:
            self._throttle_monitor.check(ip)

          kwargs = {}
          if fn._profile_format is not None:
            kwargs['profile'] = self._get_input_profile(request, fn._profile_format, request_state)
          if fn._data_format is not None:
            kwargs['data'] = await self._get_input_data(request, fn._data_format, request_state)
          output = fn(**kwargs)
          if asyncio.iscoroutinefunction(fn):
            output = await output
          request_state['response_data'] = content = json.dumps(output, cls=OPJSONEncoder, separators=(',', ':'))
          request_state['response_code'] = 200
          return Response(content=content, media_type='application/json')

        except Exception as e:
          # Make it harder to run a DOS attack by just repeatedly calling APIs with bad input
          await asyncio.sleep(3)

          request_state['exc_info'] = sys.exc_info()
          request_state['response_code'] = status_code = e.status_code if isinstance(e, ServiceException) else 500
          if status_code == 429:
            if self._throttle_monitor:
              with ignore_exceptions(log_msg='Failed to block IP: %s' % ip):
                await self._throttle_monitor.block_ip(ip)
            if self._error_monitor:
              with ignore_exceptions(log_msg='Alarm failed'):
                self._error_monitor.on_route_end(request_state)
          elif status_code == 500:
            logging.exception('Uncaught exception on route: ' + route_name)
            if self._error_monitor:
              with ignore_exceptions(log_msg='Alarm failed'):
                self._error_monitor.on_route_end(request_state)
            return Response(content='', status_code=500)
          request_state['response_data'] = content = json.dumps({ 'err': e.__class__.__name__, 'msg': str(e) }, separators=(',', ':'))
          return Response(content=content, status_code=status_code, media_type='application/json')

        finally:
          self._request_monitor.on_route_end(request_state)

      return fn
    return decorator

  @contextmanager
  def report_exceptions(self, msg):
    kwargs = { 'metadata': op_version_metadata }
    try:
      yield
    except Exception:
      logging.exception(msg)
      if self.alarm:
        try:
          _, exc_value, exc_traceback = sys.exc_info()
          stack_summary = traceback.extract_tb(exc_traceback)
          kwargs['exception'] = { 'msg': repr(exc_value), 'traceback': stack_summary.format() }
          self.alarm.report_issue(msg, **kwargs)
        except:
          logging.exception('Failed to report exception')
