import asyncio
import functools
import importlib.util
import inspect
import json
import logging
import op_alarm.factory
import op_storage.factory
import os
import re
import signal
import socket
import sys
import traceback

from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor
from contextlib import contextmanager, suppress
from datetime import datetime, timezone
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 Condition, Thread

from .env import OP_DEV_CONFIG, OP_DEV_METADATA
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 .op_json_encoder import OPJSONEncoderFallback
from .task_manager import OPTaskManager
from .type_validation import describe, validate


_REQUEST_LOGGER = logging.getLogger('op-request')

# This structure describes the order / priority for printing out request state information - the lists are read from top
# to bottom, and keys that come first in a list will be printed instead of any following keys.
_STATE_GROUPS = [
  [ 'profile_validated', 'profile' ],
  [ 'input_data_validated', 'input_data' ]
]

# key -> (header, transform)
_STATE_KEY_OUTPUTS = {
  'profile': 'p', # Input profile
  'profile_validated': 'pv', # Input profile (validated)
  'input_data': 'i', # Input data
  'input_data_validated': 'iv' # Input data (validated)
}

def _on_route_start(request_state):
  _REQUEST_LOGGER.info(json.dumps({ '_': 0, 'r': request_state['route'], 'ip': request_state['ip'], 'ua': request_state['ua'] }, separators=(',', ':')))

def _on_route_end(request_state):
  msg = { '_': 1, 'r': request_state['route'], 'ip': request_state['ip'], 'ua': request_state['ua'], 'c': request_state['response_code'] }
  for state_group in _STATE_GROUPS:
    for key in state_group:
      if key in request_state:
        msg[ _STATE_KEY_OUTPUTS[key] ] = request_state[key]
        break
  if 'o' in msg:
    # This is a special case: this output is already serialized
    del msg['o']
  packed = json.dumps(msg, cls=OPJSONEncoderFallback, separators=(',', ':'))
  if 'response_data' in request_state:
    # The 'response_data' is a special case because it is already serialized (and potentially large), so we can avoid
    # reserializing it if we just inject it directly
    packed = packed[:-1] + ',"o":' + request_state['response_data'] + '}'
  _REQUEST_LOGGER.info(packed)


_IP_ADDR_V4_REGEX = re.compile(r'^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}')

def _get_ip_address(request):
  # In non-test environments, all requests should have the X-Forwarded-For header set correctly by our proxy
  x_forwarded_for = request.headers.get('X-Forwarded-For')
  if not x_forwarded_for:
    return request.client.host
  # https://docs.aws.amazon.com/elasticloadbalancing/latest/application/x-forwarded-headers.html#x-forwarded-for
  with suppress(Exception):
    return _IP_ADDR_V4_REGEX.match(x_forwarded_for)[0]
  # This must be an IP V6 address - AWS may wrap the V6 address in brackets if it includes a port
  if x_forwarded_for.startswith('['):
    ip_v6 = x_forwarded_for.split(']', 1)[0][1:]
  else:
    ip_v6 = x_forwarded_for.split(',', 1)[0]
  with suppress(Exception):
    socket.inet_pton(socket.AF_INET6, ip_v6)
    return ip_v6
  # We couldn't determine the IP address from the provided string. This shouldn't happen as all requests must have an
  # IP address, but I'm not confident enough in my logic / configuration yet to throw an exception in this case.
  return ''


# The allowed argument names for functions decorated with @route
_ALLOWED_ARGS = [ 'request', 'ip_address', 'user_agent', 'profile', 'data' ]

class OPService:

  def __init__(self):
    # 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, on_startup=[ self._on_startup ])

    self._init_thread_pool = ThreadPoolExecutor(max_workers=64)
    self._request_thread_pool = ThreadPoolExecutor(max_workers=64)
    self._routes = []
    self._init_fns = []
    self._init_condition = Condition()
    self._init_complete = False

    self.alarm = op_alarm.factory.build(OP_DEV_CONFIG['op_alarm']['type'], OP_DEV_CONFIG['op_alarm']['config'])
    self.db = op_storage.factory.build(OP_DEV_CONFIG['op_storage']['type'], OP_DEV_CONFIG['op_storage']['config'])
    self.jwt = JWTManager(OP_DEV_CONFIG['jwt'])
    self.task = OPTaskManager(self)

    # Every production service has a '/' health-check endpoint
    @self.route('/', methods=[ 'GET' ], private=True)
    async def health(request):
      with self._init_condition:
        self._init_condition.wait_for(lambda: self._init_complete)
      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(path='/__private/metadata', private=True)
    async def metadata():
      return OP_DEV_METADATA

    # Returns documentation about the routes
    @self.route(path='/__private/explore', private=True)
    async def explore():
      grouped_routes = defaultdict(list)
      for route in self._routes:
        if not route['private']:
          grouped_routes[route['grouping']].append(route)
      output = []
      for key in sorted(grouped_routes):
        output.append(
          {
            'group': key,
            'routes': [
              {
                'name': x['name'],
                'profile': describe(x['profile_fmt']),
                'data': describe(x['data_fmt']),
                'methods': x['methods'],
                'documentation': x['fn'].__doc__
              } for x in grouped_routes[key]
            ]
          }
        )
      return output

    if os.environ.get('OP_DEV_SERVICE_TEST', False):
      # This service was brought up in dev mode (using 'op-dev run')
      # Add a private API that allows us to generate custom JWTs
      @self.route(path='/__private/generate-jwt', data=dict, private=True)
      async def generate_jwt(data):
        return self.jwt.encode(data)


  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, path=None, data=None, profile=None, methods=[ 'POST' ], private=False):
    data_fmt = data
    profile_fmt = profile

    def decorator(fn):
      nonlocal path, data_fmt, profile_fmt, methods, private

      route_name = path or ('/' + fn.__name__)
      for route in self._routes:
        if route['name'] == route_name:
          raise Exception(f'Duplicate route: {route_name}')

      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,
        'profile_fmt': profile_fmt,
        'data_fmt': data_fmt,
        'methods': methods,
        'private': private,
        'grouping': calling_module_name
      })

      args = set(inspect.signature(fn).parameters)
      for arg in args:
        if arg not in _ALLOWED_ARGS:
          raise Exception('Unexpected parameter name for route function (allowed values: %s)' % _ALLOWED_ARGS)

      kwargs = { x: None for x in args }

      @self.starlette.route(route_name, methods=methods)
      async def run(request):
        ip_address = _get_ip_address(request)
        user_agent = request.headers.get('User-Agent')
        request_state = { 'route': route_name, 'ip': ip_address, 'ua': user_agent }

        # TODO: if the called URL is different than the route name, include it in the request_state (i.e. query string
        # parameters on a GET request)
        # if not request.url.endswith(route_name):
        #   request_state['route_actual'] = request.url

        _on_route_start(request_state)

        try:
          if profile_fmt:
            profile = self._get_input_profile(request, profile_fmt, request_state)
          if data_fmt:
            data = await self._get_input_data(request, data_fmt, request_state)
          for arg in args:
            # We've already validated that 'args' does not contain anything not in '_ALLOWED_ARGS', so this is safe
            kwargs[arg] = eval(arg)

          if asyncio.iscoroutinefunction(fn):
            output = await fn(**kwargs)
          else:
            # Make sure we don't block the main loop from accepting more requests. Note that bascially eliminates any
            # advantage we get using Starlette since everything is ultimately run in a thread anyway...
            output = await asyncio.get_event_loop().run_in_executor(self._request_thread_pool, functools.partial(fn, **kwargs))

          if issubclass(type(output), Response):
            request_state['response_type'] = type(output)
            request_state['response_code'] = output.status_code
            return 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(1)

          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 == 500:
            logging.exception('Uncaught exception on route: ' + route_name)
            with ignore_exceptions(log_msg='Alarm failed'):
              self._report(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:
          _on_route_end(request_state)

      return fn
    return decorator



  def init(self, block=True):
    def decorator(fn):
      if self._init_complete:
        raise Exception('Service is already initialized')
      self._init_fns.append((fn, block))
      return fn
    return decorator


  def _on_startup(self):
    logging.info('Starting worker...\nDomain: %(domain)s\nPackage ID: %(pkg_id)s\nCommit ID: %(commit_id)s\nConfig ID: %(config_id)s' % OP_DEV_METADATA)
    tasks = []
    for fn, block in self._init_fns:
      if asyncio.iscoroutinefunction(fn):
        task = self._init_thread_pool.submit(functools.partial(asyncio.run, fn()))
      else:
        task = self._init_thread_pool.submit(fn)
      tasks.append((task, block))
    for task, block in tasks:
      if block:
        task.result()
    self._init_complete = True
    with self._init_condition:
      self._init_condition.notify_all()
    logging.info('Finished starting worker')

    if 'OP_DEV_SERVICE_MAIN' in os.environ:
      sys.argv = [ '__main__', *json.loads(os.environ['OP_DEV_SERVICE_MAIN']) ]
      spec = importlib.util.spec_from_file_location('__main__', '/var/www/app/__main__.py')
      module = importlib.util.module_from_spec(spec)
      sys.modules['__main__'] = module
      try:
        spec.loader.exec_module(module)
      except Exception as e:
        print(traceback.format_exc(), end='', flush=True)
      finally:
        signal.raise_signal(signal.SIGTERM)

    for task, block in tasks:
      if not block:
        task.result()
    self._init_thread_pool.shutdown()


  def _report(self, request_state):
    # Any 500-level error surfaced here should have 'exc_info' attached
    exc_type, exc_value, exc_traceback = request_state['exc_info']
    stack_summary = traceback.extract_tb(exc_traceback)
    cause = stack_summary[-1]
    rel_file = get_relative_sys_filename(cause.filename)

    msg = '%s -> %s - %s %s#%s' % (request_state['route'], request_state['response_code'], exc_type.__name__, rel_file, cause.lineno)
    now = datetime.now(timezone.utc)
    issue = {
      'exception': {
        'msg': repr(exc_value),
        'traceback': stack_summary.format()
      },
      'request': {
        'route': request_state['route'],
        'profile': request_state.get('profile_validated'),
        'input_data': request_state.get('input_data'),
        'response_code': request_state['response_code'],
        'response_data': request_state.get('response_data')
      },
      'metadata': OP_DEV_METADATA,
      'count': {
        'n': 1,
        't_start': now.isoformat(),
        't_end': now.isoformat()
      }
    }

    for key in [ 'profile', 'input_data', 'response_data' ]:
      if issue['request'][key]:
        issue['request'][key] = json.dumps(issue['request'][key], cls=OPJSONEncoderFallback)

    self.alarm.report_issue(msg, **issue)


  @contextmanager
  def report_exceptions(self, msg):
    kwargs = { 'metadata': OP_DEV_METADATA }
    try:
      yield
    except Exception:
      logging.exception(msg)
      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')


service = OPService()
