
from __future__ import annotations

import dataclasses
import functools
import logging
import os
import typing as t
from importlib import import_module

from nr.util.process import Pidfile, replace_stdio, spawn_daemon

from . import AWSGILauncher, BaseAWSGIConfig

if t.TYPE_CHECKING:
  import flask  # type: ignore

logger = logging.getLogger(__name__)


def load_entrypoint(ep: str) -> t.Any:
  """
  Loads an entrypoint definition of the form `<module>:<member>`.
  """

  module_name, member_name = ep.split(':', 1)
  return getattr(import_module(module_name), member_name)



@dataclasses.dataclass
class FlaskRunner(AWSGILauncher, BaseAWSGIConfig):
  """
  Runs a #flask.Flask application using the Flask development server. This is not suitable
  for production, though very useful during development due to the built-in debugging
  capabilities.

  Note: Daemonizing the Flask runner is currently only available on systems that support
  #os.fork().
  """

  #: Enable the Flask server's debug capabilities.
  debug: bool = False

  #: Automatically defaults to `True` if #debug is enabled.
  use_reloader: bool | None = None

  #: Enable stdio redirects.
  redirect_stdio: bool = False

  def get_status(self) -> Pidfile.Status:
    """
    Returns the status of the application, i.e. whether it is currently running or not. This
    method is only relevant when the application was started in the background.
    """

    filename = self.files.pidfile
    if not filename:
      return Pidfile.Status.UNKNOWN
    return Pidfile(filename).get_status()

  def is_main(self) -> bool:
    """
    Checks whether this is the main process, which can be the reloader thread in case
    the #FlaskRunner is used.
    """

    if os.getenv('_FLASK_USE_RELOAER') == 'true':
      return not os.getenv('WERKZEUG_RUN_MAIN')

    return True

  def launch(self, entrypoint: str, factory: bool) -> None:
    import flask  # type: ignore[module]

    # Set up the environment.
    use_reloader = self.debug if self.use_reloader is None else self.use_reloader
    if use_reloader:
      os.environ['_FLASK_USE_RELOAER'] = 'true'
    if self.debug:
      os.environ['FLASK_DEBUG'] = 'true'

    # Load the Flask application.
    app = load_entrypoint(entrypoint)
    if factory:
      app = app()
    if not isinstance(app, flask.Flask):
      raise RuntimeError('entrypoint ({}) must be a Flask application.'.format(entrypoint))

    # Setup stdio redirects.
    if self.redirect_stdio:
      stdout, stderr = self._init_redirects(self.files)
    else:
      stdout, stderr = None, None

    if self.files.pidfile:
      os.makedirs(os.path.dirname(self.files.pidfile), exist_ok=True)

    if self.ssl:
      ssl_context = (self.ssl.cert, self.ssl.key)
    else:
      ssl_context = None

    run = functools.partial(
      self._run,
      app=app,
      stdout=stdout,
      stderr=stderr,
      pidfile=self.files.pidfile,
      daemon=self.daemonize,
      host=self.host,
      port=self.port,
      debug=self.debug,
      use_reloader=use_reloader,
      ssl_context=ssl_context)

    if self.is_main() and self.get_status() == Pidfile.Status.RUNNING:
      raise RuntimeError('Application is already running.')

    if self.daemonize:
      spawn_daemon(run)
    else:
      run()

  def _init_redirects(self, files: BaseAWSGIConfig.Files) -> tuple[t.TextIO | None, t.TextIO | None]:
    if files.stdout:
      os.makedirs(os.path.dirname(files.stdout), exist_ok=True)
      stdout = open(files.stdout, 'a+')
    else:
      stdout = None

    stderr: t.TextIO | None
    if files.stderr and files.stderr != files.stdout:
      os.makedirs(os.path.dirname(files.stderr), exist_ok=True)
      stderr = open(files.stderr, 'a+')
    elif files.stderr == files.stdout:
      stderr = stdout
    else:
      stderr = None

    return stdout, stderr

  def _run(
    self,
    app: 'flask.Flask',
    stdout: t.TextIO | None,
    stderr: t.TextIO | None,
    pidfile: str | None,
    daemon: bool,
    host: str,
    port: int,
    debug: bool,
    use_reloader: bool,
    ssl_context: tuple[str, str]
  ) -> None:
    """
    Internal function to actually invoke the Flask application after everything was prepared.
    """

    if stdout or stderr:
      replace_stdio(None, stdout, stderr)

    if pidfile:
      with open(pidfile, 'w') as fp:
        fp.write(str(os.getpid()))

    if daemon:
      logger.info('Process %s started.', os.getpid())

    try:
      app.run(
        host=host,
        port=port,
        debug=debug,
        use_reloader=use_reloader,
        ssl_context=ssl_context,
      )
    finally:
      if (not use_reloader or os.getenv('WERKZEUG_RUN_MAIN') == 'true') and pidfile:
        try:
          logger.info('Removing pidfile "%s" from PID %s.', pidfile, os.getpid())
          os.remove(pidfile)
        except OSError as exc:
          logger.exception('Unable to remove "%s".', pidfile)


from __future__ import annotations

import dataclasses
import functools
import logging
import os
import typing as t
from importlib import import_module

from nr.util.process import Pidfile, replace_stdio, spawn_daemon

from . import AWSGILauncher, BaseAWSGIConfig

if t.TYPE_CHECKING:
  import flask  # type: ignore

logger = logging.getLogger(__name__)


def load_entrypoint(ep: str) -> t.Any:
  """
  Loads an entrypoint definition of the form `<module>:<member>`.
  """

  module_name, member_name = ep.split(':', 1)
  return getattr(import_module(module_name), member_name)



@dataclasses.dataclass
class FlaskRunner(AWSGILauncher, BaseAWSGIConfig):
  """
  Runs a #flask.Flask application using the Flask development server. This is not suitable
  for production, though very useful during development due to the built-in debugging
  capabilities.

  Note: Daemonizing the Flask runner is currently only available on systems that support
  #os.fork().
  """

  #: Enable the Flask server's debug capabilities.
  debug: bool = False

  #: Automatically defaults to `True` if #debug is enabled.
  use_reloader: bool | None = None

  #: Enable stdio redirects.
  redirect_stdio: bool = False

  def get_status(self) -> Pidfile.Status:
    """
    Returns the status of the application, i.e. whether it is currently running or not. This
    method is only relevant when the application was started in the background.
    """

    filename = self.files.pidfile
    if not filename:
      return Pidfile.Status.UNKNOWN
    return Pidfile(filename).get_status()

  def is_main(self) -> bool:
    """
    Checks whether this is the main process, which can be the reloader thread in case
    the #FlaskRunner is used.
    """

    if os.getenv('_FLASK_USE_RELOAER') == 'true':
      return not os.getenv('WERKZEUG_RUN_MAIN')

    return True

  def launch(self, entrypoint: str, factory: bool) -> None:
    import flask  # type: ignore[module]

    # Set up the environment.
    use_reloader = self.debug if self.use_reloader is None else self.use_reloader
    if use_reloader:
      os.environ['_FLASK_USE_RELOAER'] = 'true'
    if self.debug:
      os.environ['FLASK_DEBUG'] = 'true'

    # Load the Flask application.
    app = load_entrypoint(entrypoint)
    if factory:
      app = app()
    if not isinstance(app, flask.Flask):
      raise RuntimeError('entrypoint ({}) must be a Flask application.'.format(entrypoint))

    # Setup stdio redirects.
    if self.redirect_stdio:
      stdout, stderr = self._init_redirects(self.files)
    else:
      stdout, stderr = None, None

    if self.files.pidfile:
      os.makedirs(os.path.dirname(self.files.pidfile), exist_ok=True)

    if self.ssl:
      ssl_context = (self.ssl.cert, self.ssl.key)
    else:
      ssl_context = None

    run = functools.partial(
      self._run,
      app=app,
      stdout=stdout,
      stderr=stderr,
      pidfile=self.files.pidfile,
      daemon=self.daemonize,
      host=self.host,
      port=self.port,
      debug=self.debug,
      use_reloader=use_reloader,
      ssl_context=ssl_context)

    if self.is_main() and self.get_status() == Pidfile.Status.RUNNING:
      raise RuntimeError('Application is already running.')

    if self.daemonize:
      spawn_daemon(run)
    else:
      run()

  def _init_redirects(self, files: BaseAWSGIConfig.Files) -> tuple[t.TextIO | None, t.TextIO | None]:
    if files.stdout:
      os.makedirs(os.path.dirname(files.stdout), exist_ok=True)
      stdout = open(files.stdout, 'a+')
    else:
      stdout = None

    stderr: t.TextIO | None
    if files.stderr and files.stderr != files.stdout:
      os.makedirs(os.path.dirname(files.stderr), exist_ok=True)
      stderr = open(files.stderr, 'a+')
    elif files.stderr == files.stdout:
      stderr = stdout
    else:
      stderr = None

    return stdout, stderr

  def _run(
    self,
    app: 'flask.Flask',
    stdout: t.TextIO | None,
    stderr: t.TextIO | None,
    pidfile: str | None,
    daemon: bool,
    host: str,
    port: int,
    debug: bool,
    use_reloader: bool,
    ssl_context: tuple[str, str]
  ) -> None:
    """
    Internal function to actually invoke the Flask application after everything was prepared.
    """

    if stdout or stderr:
      replace_stdio(None, stdout, stderr)

    if pidfile:
      with open(pidfile, 'w') as fp:
        fp.write(str(os.getpid()))

    if daemon:
      logger.info('Process %s started.', os.getpid())

    try:
      app.run(
        host=host,
        port=port,
        debug=debug,
        use_reloader=use_reloader,
        ssl_context=ssl_context,
      )
    finally:
      if (not use_reloader or os.getenv('WERKZEUG_RUN_MAIN') == 'true') and pidfile:
        try:
          logger.info('Removing pidfile "%s" from PID %s.', pidfile, os.getpid())
          os.remove(pidfile)
        except OSError as exc:
          logger.exception('Unable to remove "%s".', pidfile)

# type: ignore Loads an entrypoint definition of the form `<module>:<member>`. Runs a #flask.Flask application using the Flask development server. This is not suitable   for production, though very useful during development due to the built-in debugging   capabilities.    Note: Daemonizing the Flask runner is currently only available on systems that support   #os.fork(). #: Enable the Flask server's debug capabilities. #: Automatically defaults to `True` if #debug is enabled. #: Enable stdio redirects. Returns the status of the application, i.e. whether it is currently running or not. This     method is only relevant when the application was started in the background. Checks whether this is the main process, which can be the reloader thread in case     the #FlaskRunner is used. # type: ignore[module] # Set up the environment. # Load the Flask application. # Setup stdio redirects. Internal function to actually invoke the Flask application after everything was prepared.