import re

from sqlalchemy import create_engine, Column, Integer, String, schema
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.exc import IntegrityError


BASE = declarative_base()


_RE_DB = {
    "sqlite": re.compile(r"^.*columns?([^)]+)(is|are)\s+not\s+unique$"),
    "postgresql": re.compile(r"^.*duplicate\s+key.*\"([^\"]+)\"\s*\n.*$"),
    "mysql": re.compile(r"^.*\(1062,.*'([^\']+)'\"\)$")
}


class NovaException(Exception):
    message = "An unknown exception occurred."
    code = 500
    headers = {}
    safe = False
    def __init__(self, message=None, **kwargs):
        self.kwargs = kwargs
        if 'code' not in self.kwargs:
            try:
                self.kwargs['code'] = self.code
            except AttributeError:
                pass
        if not message:
            try:
                message = self.message % kwargs
            except Exception as e:
                # kwargs doesn't match a variable in the message
                # log the issue and the kwargs
                LOG.exception(_('Exception in string format operation'))
                for name, value in kwargs.iteritems():
                    LOG.error("%s: %s" % (name, value))
                if CONF.fatal_exception_format_errors:
                    raise e
                else:
                    # at least get the core message out if something happened
                    message = self.message
        super(NovaException, self).__init__(message)


class DBError(NovaException):
    """Wraps an implementation specific exception."""
    def __init__(self, inner_exception=None):
        self.inner_exception = inner_exception
        super(DBError, self).__init__(str(inner_exception))


class DBDuplicateEntry(NovaException):
    """Wraps an implementation specific exception."""
    def __init__(self, columns=[], inner_exception=None):
        self.columns = columns
        super(DBDuplicateEntry, self).__init__(inner_exception)


class Album(BASE):
    __tablename__ = 'albums'
    __table_args__ = (schema.UniqueConstraint("artist", "album"), )
    id = Column(Integer, primary_key=True)
    artist = Column(String(255))
    album = Column(String(255))
    def __init__(self, artist, album):
        self.album = album
        self.artist = artist


def raise_if_duplicate_entry_error(integrity_error, engine_name):
    """ In this function will be raised DBDuplicateEntry exception if integrity
        error wrap unique constraint violation. """
    def get_columns_from_uniq_cons_or_name(columns):
        uniqbase = "uniq_"
        if not columns.startswith(uniqbase):
            if engine_name == "postgresql":
                return [columns[columns.index("_") + 1:columns.rindex("_")]]
            return [columns]
        return columns[len(uniqbase):].split("_x_")
    if engine_name not in ["mysql", "sqlite", "postgresql"]:
        return
    m = _RE_DB[engine_name].match(integrity_error.message)
    if not m:
        return
    columns = m.group(1)
    if engine_name == "sqlite":
        columns = columns.strip().split(", ")
    else:
        columns = get_columns_from_uniq_cons_or_name(columns)
    raise DBDuplicateEntry(columns, integrity_error)


def wrap_db_error(f):
    def _wrap(*args, **kwargs):
        try:
            return f(*args, **kwargs)
        except IntegrityError, e:
            raise_if_duplicate_entry_error(e, engine.name)
            raise DBError(e)
        except Exception, e:
            raise DBError(e)
    _wrap.func_name = f.func_name
    return _wrap


engine = create_engine('mysql://root:test@localhost/test?charset=utf8', echo=False)
Session = sessionmaker(bind=engine, autocommit=True)
session = Session()
session.flush = wrap_db_error(session.flush)

#create the test table
BASE.metadata.create_all(engine)

# add a unique record
ed = Album('Ed Jones', 'ed')
session.add(ed)
session.flush()

# add duplicate .... this should explode
session.add(Album('Ed Jones', 'ed'))
session.flush()
