#!/usr/bin/env python
'''
This file handles all the database transactions.
'''

import os
from contextlib import contextmanager
from multiprocessing.util import register_after_fork

from sqlalchemy.orm import scoped_session, sessionmaker
from sqlalchemy import create_engine, exc
from sqlalchemy.pool import NullPool
from sqlalchemy.orm import Session as BaseSession

from framework.utils import FileOperations
from framework.dependency_management.dependency_resolver import BaseComponent
from framework.dependency_management.interfaces import DBInterface
from framework.db import models


class Session(BaseSession):
    def __init__(self, *a, **kw):
        super(Session, self).__init__(*a, **kw)
        self._in_atomic = False

    @contextmanager
    def atomic(self):
        """
        Transaction context manager.

        Will commit the transaction on successful completion
        of the block, or roll it back on error.

        Supports nested usage (via savepoints).
        """
        nested = self._in_atomic
        self.begin(nested=nested)
        self._in_atomic = True
        try:
            yield
        except:
            self.rollback()
            raise
        else:
            self.commit()
        finally:
            if not nested:
                self._in_atomic = False


class DB(BaseComponent, DBInterface):

    COMPONENT_NAME = "db"

    def __init__(self):
        self.register_in_service_locator()
        self.config = self.get_component("config")
        self.error_handler = self.get_component("error_handler")
        self.ErrorDBSession = None
        self.Transaction = None
        self.URL = None
        self.Plugin = None
        self.POutput = None
        self.Target = None
        self.Resource = None
        self.Error = None
        self.CommandRegister = None
        self.Mapping = None
        self._db_settings = self._get_db_settings()
        self.create_session()

    def init(self):
        self.config = self.get_component("config")
        self.Transaction = self.get_component("transaction")
        self.URL = self.get_component("url_manager")
        self.Plugin = self.get_component("db_plugin")
        self.POutput = self.get_component("plugin_output")
        self.Target = self.get_component("target")
        self.Resource = self.get_component("resource")
        self.Config = self.get_component("db_config")
        self.Error = self.get_component("db_error")
        self.CommandRegister = self.get_component("command_register")
        self.Mapping = self.get_component("mapping_db")
        self.error_handler = self.get_component("error_handler")
        self.OWTFSession = self.get_component("session_db")
        self.Worklist = self.get_component("worklist_manager")
        self.DBHealthCheck()

    def get_category(self, plugin_code):
        return self.Mapping.GetCategory(plugin_code)

    def DBHealthCheck(self):
        # TODO: Fix this for postgresql: self.Target.DBHealthCheck()
        return

    def create_session(self):
        self.Session = self.CreateScopedSession()
        self.session = self.Session()

    def clean_up(self):
        """Close the sqlalchemy session opened by DB."""
        self.session.close()

    def _get_db_settings(self):
        """Create DB settings according to the configuration file."""
        config_path = os.path.expanduser(self.config.FrameworkConfigGet('DATABASE_SETTINGS_FILE'))
        settings = {}
        with FileOperations.open(config_path, 'r') as f:
            for line in f:
                line = line.rstrip()
                # Ignore empty/comment lines.
                if not line or line.startswith('#'):
                    continue
                try:
                    key, value = line.split(':')
                    settings[key.strip()] = value.strip()
                except ValueError:
                    self.error_handler.FrameworkAbort("Problem in config file: '%s' -> Cannot parse line: %s" %
                                                      (config_path, line))
        return settings

    def CreateEngine(self, BaseClass):
        try:
            engine = create_engine(
                "postgresql+psycopg2://%s:%s@%s:%s/%s" % (
                    self._db_settings['DATABASE_USER'],
                    self._db_settings['DATABASE_PASS'],
                    self._db_settings['DATABASE_IP'],
                    self._db_settings['DATABASE_PORT'],
                    self._db_settings['DATABASE_NAME']),
                poolclass=NullPool)
            BaseClass.metadata.create_all(engine)
            # Fix for forking
            register_after_fork(engine, engine.dispose)
            return engine
        except ValueError as e:  # Potentially corrupted DB config.
            self.error_handler.FrameworkAbort(
                "Database configuration file is potentially corrupted. Please check %s\n[DB] %s" %
                (self.config.FrameworkConfigGet('DATABASE_SETTINGS_FILE'), str(e)))
        except KeyError:  # Indicates incomplete db config file
            self.error_handler.FrameworkAbort("Incomplete database configuration settings in %s" %
                                              self.config.FrameworkConfigGet('DATABASE_SETTINGS_FILE'))
        except exc.OperationalError as e:
            self.error_handler.FrameworkAbort("[DB] %s\nRun scripts/db_run.sh to start/setup db" % str(e))

    def CreateScopedSession(self):
        # Not to be used apart from main process, use CreateSession instead
        self.engine = self.CreateEngine(models.Base)
        session_factory = sessionmaker(bind=self.engine, autocommit=False, class_=Session)
        return scoped_session(session_factory)

@@ -301,20 +301,21 @@
  import 
-Null
+Queue
 Pool%0Afro
@@ -4526,16 +4526,79 @@
 ass=
-NullPool
+QueuePool,%0A                pool_size=5,%0A                max_overflow=10
 )%0A  

Revert "[db] Prefer NullPool for forked SQLAlchemy connections"