import logging
from contextlib import contextmanager
from functools import lru_cache, wraps
from typing import Iterator, Optional, Callable
from datetime import datetime

from sqlalchemy.ext.declarative import (
    declarative_base,
)
from sqlalchemy import (
    engine as sa_engine,
    create_engine as sa_create_engine,
    Column,
    types,
)
from sqlalchemy.orm import (
    Session,
    sessionmaker as sa_sessionmaker,
)

from stackbot.config import settings


BaseModel = declarative_base()
logger = logging.getLogger(__name__)


def now():
    return datetime.now().astimezone(settings.DEFAULT_TIMEZONE)


class TimestampedModelMixin:
    created_at = Column(
        types.DateTime(timezone=True),
        default=now,
    )
    updated_at = Column(
        types.DateTime(timezone=True),
        default=now,
        onupdate=now,
    )


@lru_cache()
def get_cached_engine(url: Optional[str] = None) -> sa_engine.Engine:
    if url is None:
        url = settings.database_url
    return sa_create_engine(url=url)


@lru_cache()
def get_cached_sessionmaker(engine: Optional[sa_engine.Engine] = None) -> sa_sessionmaker:
    if engine is None:
        engine = get_cached_engine()
    return sa_sessionmaker(bind=engine)


@contextmanager
def get_db() -> Iterator[Session]:
    sessionmaker = get_cached_sessionmaker()
    session = sessionmaker()
    try:
        yield session
        session.commit()
    except Exception as exc:
        session.rollback()
        raise exc
    finally:
        session.close()


def dbconnect(func: Callable):
    @wraps(func)
    def wrapper(*args, **kwargs):
        session = kwargs.get('session')
        if not session:
            engine = sa_create_engine(settings.database_url)
            session = Session(bind=engine, )
            kwargs['session'] = session
        try:
            result = func(*args, **kwargs)
            session.commit()
            return result
        except Exception as exc:
            logger.exception(f'Unexpected error in {func.__name__}: {repr(exc)}')
            session.rollback()
            raise
        finally:
            session.close()
    return wrapper
