#!/usr/bin/env python
# coding=utf-8

import os
import sys
import configobj
if os.name != 'nt':
    import pwd
    import grp

try:
    from setproctitle import setproctitle
    setproctitle  # workaround for pyflakes issue #13
except ImportError:
    setproctitle = None

for path in [
    os.path.join('opt', 'diamond', 'lib'),
    os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src'))
]:
    if os.path.exists(os.path.join(path, 'diamond', '__init__.py')):
        sys.path.append(path)
        break

from diamond.server import Server
from diamond.util import get_diamond_version

import optparse
import signal
import logging.config


def main():
    try:
        # Initialize Options

        defaults = {
            'skip_pidfile': False,
        }

        if os.name == 'nt':
            defaults['skip_pidfile'] = True

        parser = optparse.OptionParser()

        parser.add_option("-c", "--configfile",
                          dest="configfile",
                          default="/etc/diamond/diamond.conf",
                          help="config file")

        parser.add_option("-f", "--foreground",
                          dest="foreground",
                          default=False,
                          action="store_true",
                          help="run in foreground")

        parser.add_option("-l", "--log-stdout",
                          dest="log_stdout",
                          default=False,
                          action="store_true",
                          help="log to stdout")

        parser.add_option("-p", "--pidfile",
                          dest="pidfile",
                          default=None,
                          help="pid file")

        parser.add_option("-r", "--run",
                          dest="collector",
                          default=None,
                          help="run a given collector once and exit")

        parser.add_option("-v", "--version",
                          dest="version",
                          default=False,
                          action="store_true",
                          help="display the version and exit")

        parser.add_option("--skip-pidfile",
                          dest="skip_pidfile",
                          default=defaults['skip_pidfile'],
                          action="store_true",
                          help="Skip creating PID file")

        parser.add_option("-u", "--user",
                          dest="user",
                          default=None,
                          help="Change to specified unprivilegd user")
        parser.add_option("-g", "--group",
                          dest="group",
                          default=None,
                          help="Change to specified unprivilegd group")
        parser.add_option("--skip-change-user",
                          dest="skip_change_user",
                          default=False,
                          action="store_true",
                          help="Skip changing to an unprivilegd user")

        parser.add_option("--skip-fork",
                          dest="skip_fork",
                          default=False,
                          action="store_true",
                          help="Skip forking (damonizing) process")

        # Parse Command Line Args
        (options, args) = parser.parse_args()

        # Initial variables
        uid = -1
        gid = -1

        if options.version:
            print "Diamond version %s" % (get_diamond_version())
            sys.exit(0)

        # Initialize Config
        if os.path.exists(options.configfile):
            config = configobj.ConfigObj(os.path.abspath(options.configfile))
            config['configfile'] = options.configfile
        else:
            print >> sys.stderr, "ERROR: Config file: %s does not exist." % (
                options.configfile)
            parser.print_help(sys.stderr)
            sys.exit(1)

        # Initialize Logging
        log = logging.getLogger('diamond')

        if options.log_stdout:
            log.setLevel(logging.DEBUG)
            # Configure Logging Format
            formatter = logging.Formatter(
                '[%(asctime)s] [%(threadName)s] %(message)s')
            # handler
            streamHandler = logging.StreamHandler(sys.stdout)
            streamHandler.setFormatter(formatter)
            streamHandler.setLevel(logging.DEBUG)
            log.addHandler(streamHandler)
        else:
            try:
                if sys.version_info >= (2, 6):
                    logging.config.fileConfig(options.configfile,
                                              disable_existing_loggers=False)
                else:
                    # python <= 2.5 does not have disable_existing_loggers
                    # default was to always disable them, in our case we want to
                    # keep any logger created by handlers
                    logging.config.fileConfig(options.configfile)
                    for logger in logging.root.manager.loggerDict.values():
                        logger.disabled = 0
            except Exception, e:
                sys.stderr.write("Error occurs when initialize logging: ")
                sys.stderr.write(str(e))
                sys.stderr.write(os.linesep)

    # Pass the exit up stream rather then handle it as an general exception
    except SystemExit, e:
        raise SystemExit

    except Exception, e:
        import traceback
        sys.stderr.write("Unhandled exception: %s" % str(e))
        sys.stderr.write("traceback: %s" % traceback.format_exc())
        sys.exit(1)

    # Switch to using the logging system
    try:
        # PID MANAGEMENT
        if not options.skip_pidfile:
            # Initialize Pid file
            if not options.pidfile:
                options.pidfile = str(config['server']['pid_file'])

            # Read existing pid file
            try:
                pf = file(options.pidfile, 'r')
                pid = int(pf.read().strip())
                pf.close()
            except IOError:
                pid = None

            # Check existing pid file
            if pid:
                # Check if pid is real
                if not os.path.exists("/".join(["/proc", str(pid), "cmdline"])):
                    # Pid is not real
                    os.unlink(options.pidfile)
                    pid = None
                    print >> sys.stderr, (
                        "WARN: Bogus pid file was found. I deleted it.")
                else:
                    print >> sys.stderr, (
                        "ERROR: Pidfile exists. Server already running?")
                    sys.exit(1)

            # Get final GIDs
            if os.name != 'nt':
                if options.group is not None:
                    gid = grp.getgrnam(options.group).gr_gid
                elif len(config['server']['group']):
                    gid = grp.getgrnam(config['server']['group']).gr_gid

            # Get final UID
            if os.name != 'nt':
                if options.user is not None:
                    uid = pwd.getpwnam(options.user).pw_uid
                elif len(config['server']['user']):
                    uid = pwd.getpwnam(config['server']['user']).pw_uid

            # Fix up pid permissions
            if not options.foreground and not options.collector:
                # Write pid file
                pid = str(os.getpid())
                try:
                    pf = file(options.pidfile, 'w+')
                except IOError, e:
                    print >> sys.stderr, "Failed to write PID file: %s" % (e)
                    sys.exit(1)
                pf.write("%s\n" % pid)
                pf.close()
                os.chown(options.pidfile, uid, gid)
                # Log
                log.debug("Wrote First PID file: %s" % (options.pidfile))

        # USER MANAGEMENT
        if not options.skip_change_user:
            # Switch user to specified user/group if required
            try:
                if gid != -1 and os.getgid() != gid:
                    # Set GID
                    os.setgid(gid)

                if uid != -1 and os.getuid() != uid:
                    # Set UID
                    os.setuid(uid)

            except Exception, e:
                print >> sys.stderr, "ERROR: Failed to set UID/GID. %s" % (e)
                sys.exit(1)

            # Log
            log.info('Changed UID: %d (%s) GID: %d (%s).' % (
                os.getuid(),
                config['server']['user'],
                os.getgid(),
                config['server']['group']))

        # DAEMONIZE MANAGEMENT
        if not options.skip_fork:
            # Detatch Process
            if not options.foreground and not options.collector:

                # Double fork to serverize process
                log.info('Detaching Process.')

                # Fork 1
                try:
                    pid = os.fork()
                    if pid > 0:
                        # Exit first paren
                        sys.exit(0)
                except OSError, e:
                    print >> sys.stderr, "Failed to fork process." % (e)
                    sys.exit(1)
                # Decouple from parent environmen
                os.setsid()
                os.umask(0)
                # Fork 2
                try:
                    pid = os.fork()
                    if pid > 0:
                        # Exit second paren
                        sys.exit(0)
                except OSError, e:
                    print >> sys.stderr, "Failed to fork process." % (e)
                    sys.exit(1)
                # Close file descriptors so that we can detach
                sys.stdout.close()
                sys.stderr.close()
                sys.stdin.close()
                os.close(0)
                os.close(1)
                os.close(2)

        # PID MANAGEMENT
        if not options.skip_pidfile:
            # Finish Initialize PID file
            if not options.foreground and not options.collector:
                # Write pid file
                pid = str(os.getpid())
                try:
                    pf = file(options.pidfile, 'w+')
                except IOError, e:
                    log.error("Failed to write child PID file: %s" % (e))
                    sys.exit(1)
                pf.write("%s\n" % pid)
                pf.close()
                # Log
                log.debug("Wrote child PID file: %s" % (options.pidfile))

        # Initialize Server
        server = Server(config)

        def sigint_handler(signum, frame):
                # Log
                log.debug("Signal Received: %d" % (signum))
                # Stop Server
                server.stop()
                # Delete Pidfile
                if not options.skip_pidfile and os.path.exists(options.pidfile):
                    os.remove(options.pidfile)
                    # Log
                    log.debug("Removed PID file: %s" % (options.pidfile))

        # Set the signal handlers
        signal.signal(signal.SIGINT, sigint_handler)
        signal.signal(signal.SIGTERM, sigint_handler)

        if options.collector:
            # Run Server with one collector
            server.run_one(options.collector)
        else:
            # Run Server
            server.run()

    # Pass the exit up stream rather then handle it as an general exception
    except SystemExit, e:
        raise SystemExit

    except Exception, e:
        import traceback
        log.error("Unhandled exception: %s" % str(e))
        log.error("traceback: %s" % traceback.format_exc())
        sys.exit(1)

if __name__ == "__main__":
    if setproctitle:
        setproctitle('diamond')
    main()
