# Cherrypy-foundation
# Copyright (C) 2026 IKUS Software
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.

import logging
import threading
from contextlib import contextmanager
from datetime import datetime, timezone

import cherrypy
from cherrypy.lib.sessions import Session
from sqlalchemy import Column, DateTime, PickleType, String, TypeDecorator, delete, func, select

from cherrypy_foundation.plugins import db  # noqa

CONTEXT = 'DBSESSION'


class UtcDateTime(TypeDecorator):
    cache_ok = True
    impl = DateTime

    def process_bind_param(self, value: datetime, dialect):
        """Convert from aware to naive UTC."""
        if value is None:
            return None
        if value.tzinfo is None:
            raise ValueError('datetime without tzinfo: %s' % value)
        return value.astimezone(timezone.utc)

    def process_result_value(self, value, dialect):
        """Convert from naive to UTC timezone."""
        if value is None:
            return None
        if value.tzinfo is None:
            return value.replace(tzinfo=timezone.utc)
        return value.astimezone(timezone.utc)


class SessionModel:
    """
    Abstract session model. Subclass this and define __tablename__ to create the table.

    Example::

        class MySession(SessionModel, Base):
            __tablename__ = 'sessions'
    """

    session_id = Column('SessionID', String(255), unique=True, nullable=False, primary_key=True)
    data = Column('Data', PickleType)
    # Store naive UTC dates
    expiration_time = Column('ExpirationTime', UtcDateTime, nullable=False)


class DbSession(Session):
    _locks = {}
    _lock_counts = {}
    _lock_meta = threading.Lock()

    model_class = None  # Must be set by the user

    def __init__(self, id=None, **kwargs):
        Session.__init__(self, id=id, **kwargs)

    @contextmanager
    def _transaction(self):
        """
        Return a transaction context — existing one if one exists, new if not.
        With this helper, this module support autobegin=True or False.
        """
        if self.model_class is None:
            raise RuntimeError('tools.sessions.model_class is required')
        session = self.model_class.session()  # underlying Session
        if session.in_transaction():
            yield session  # return current session
        else:
            # Return new transaction
            with session.begin():
                yield session

    def _exists(self):
        return self._load() is not None

    def _load(self):
        with self._transaction() as session:
            stmt = select(self.model_class).where(
                self.model_class.session_id == self.id,
                self.model_class.expiration_time > self.now(),
            )
            sess = session.execute(stmt).scalars().first()
            if not sess:
                return None
            try:
                return (sess.data, sess.expiration_time)
            except TypeError:
                cherrypy.log('fail to read session data', context=CONTEXT, severity=logging.ERROR, traceback=False)
            return None

    def _save(self, expiration_time):
        with self._transaction() as session:
            stmt = select(self.model_class).where(self.model_class.session_id == self.id)
            sess = session.execute(stmt).scalars().first()
            if not sess:
                sess = self.model_class(session_id=self.id)
                session.add(sess)
            sess.data = dict(self._data)
            sess.expiration_time = expiration_time
            session.flush()

    def _delete(self):
        with self._transaction() as session:
            stmt = delete(self.model_class).where(self.model_class.session_id == self.id)
            session.execute(stmt)

    def clean_up(self):
        """Clean up expired sessions."""
        with self._transaction() as session:
            try:
                stmt = delete(self.model_class).where(self.model_class.expiration_time < self.now())
                session.execute(stmt)
            except Exception:
                cherrypy.log('fail to clean-up sessions', context=CONTEXT, severity=logging.ERROR, traceback=True)

    def __len__(self):
        """Return the number of active sessions."""
        with self._transaction() as session:
            stmt = (
                select(func.count()).select_from(self.model_class).where(self.model_class.expiration_time > self.now())
            )
            return session.execute(stmt).scalar()

    def acquire_lock(self):
        self.locked = True
        with self._lock_meta:
            self._locks.setdefault(self.id, threading.RLock())
            self._lock_counts[self.id] = self._lock_counts.get(self.id, 0) + 1
        self._locks[self.id].acquire()

    def release_lock(self):
        self._locks[self.id].release()
        self.locked = False
        with self._lock_meta:
            self._lock_counts[self.id] -= 1
            if self._lock_counts[self.id] == 0:
                self._locks.pop(self.id, None)
                self._lock_counts.pop(self.id, None)

    def now(self):
        """Generate a timezone aware, versions of 'now'."""
        return datetime.now(tz=timezone.utc)
