import asyncio
import importlib.util
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 threading import Thread

from .env import OP_DEV_METADATA
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 .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('')

    # TODO: since these three APIs should only ever be called from the context of the op-dev dashboard, you should add
    # some corresponding authentication

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

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

    if self._is_test_instance:
      # This service was brought up in dev mode (using 'op-dev run')

      # An entry point to run the __main__.py script with the provided arguments
      @self.route('/__private/main', { 'args': [ non_empty_str ] }, private=True)
      async def main(data):
        try:
          sys.argv = [ '__main__', *data['args'] ]
          spec = importlib.util.spec_from_file_location('__main__', '/var/www/app/__main__.py')
          module = importlib.util.module_from_spec(spec)
          sys.modules['__main__'] = module
          # Run the module in its own thread so it has its own async context
          t = Thread(target=spec.loader.exec_module, args=(module,))
          t.start()
          t.join()
        except:
          pass
        # try:
        #   if asyncio.iscoroutinefunction(module.main):
        #     await module.main()
        #   else:
        #     module.main()
        # except:
        #   pass

      if jwt_config:
        # If , 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)

  def expose(self, *keys):
    def decorator(fn):
      if not getattr(fn, '_is_route', False):
        raise Exception("Can only use the 'expose' decorator over a 'route' decorator")
      allowed_keys = [ 'ip_address', 'user_agent' ]
      for key in (keys or [ None ]):
        if key not in allowed_keys:
          raise Exception('Unrecognized property: %s' % key)
      fn._expose = keys
      return fn
    return decorator

  # 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
      fn._expose = ()

      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_address = request.headers.get('X-Forwarded-For') or request.client.host
        user_agent = request.headers.get('User-Agent')
        request_state = { 'route': route_name, 'ip': ip_address, 'ua': user_agent }
        self._request_monitor.on_route_start(request_state)

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

          kwargs = {}
          for key in fn._expose:
            if key == 'ip_address':
              kwargs[key] = ip_address
            elif key == 'user_agent':
              kwargs[key] = user_agent
          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_address):
                await self._throttle_monitor.block_ip(ip_address)
            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_DEV_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')
